Merge remote-tracking branch 'origin/main' into feat/gateway-v2

backend/src/keystore/key-value-store-dal.ts

@@ -0,0 +1,91 @@
+import { Knex } from "knex";
+
+import { TDbClient } from "@app/db";
+import { TableName } from "@app/db/schemas";
+import { ormify, TOrmify } from "@app/lib/knex";
+import { logger } from "@app/lib/logger";
+import { QueueName } from "@app/queue";
+
+export interface TKeyValueStoreDALFactory extends TOrmify<TableName.KeyValueStore> {
+  incrementBy: (key: string, dto: { incr?: number; tx?: Knex; expiresAt?: Date }) => Promise<number>;
+  findOneInt: (key: string, tx?: Knex) => Promise<number | undefined>;
+  pruneExpiredKeys: () => Promise<void>;
+}
+
+const QUERY_TIMEOUT_MS = 10 * 60 * 1000; // 10 minutes
+const CACHE_KEY_PRUNE_BATCH_SIZE = 10000;
+const MAX_RETRY_ON_FAILURE = 3;
+
+export const keyValueStoreDALFactory = (db: TDbClient): TKeyValueStoreDALFactory => {
+  const keyValueStoreOrm = ormify(db, TableName.KeyValueStore);
+
+  const incrementBy: TKeyValueStoreDALFactory["incrementBy"] = async (key, { incr = 1, tx, expiresAt }) => {
+    return (tx || db)(TableName.KeyValueStore)
+      .insert({ key, integerValue: 1, expiresAt })
+      .onConflict("key")
+      .merge({
+        integerValue: db.raw(`"${TableName.KeyValueStore}"."integerValue" + ?`, [incr]),
+        expiresAt
+      })
+      .returning("integerValue")
+      .then((result) => Number(result[0]?.integerValue || 0));
+  };
+
+  const findOneInt: TKeyValueStoreDALFactory["findOneInt"] = async (key, tx) => {
+    const doc = await (tx || db.replicaNode())(TableName.KeyValueStore)
+      .where({ key })
+      .andWhere(
+        (builder) =>
+          void builder
+            .whereNull("expiresAt") // no expiry
+            .orWhere("expiresAt", ">", db.fn.now()) // or not expired
+      )
+      .first()
+      .select("integerValue");
+    return Number(doc?.integerValue || 0);
+  };
+
+  const pruneExpiredKeys: TKeyValueStoreDALFactory["pruneExpiredKeys"] = async () => {
+    let deletedIds: { key: string }[] = [];
+    let numberOfRetryOnFailure = 0;
+    let isRetrying = false;
+
+    logger.info(`${QueueName.DailyResourceCleanUp}: db key value store clean up started`);
+    do {
+      try {
+        // eslint-disable-next-line no-await-in-loop
+        deletedIds = await db.transaction(async (trx) => {
+          await trx.raw(`SET statement_timeout = ${QUERY_TIMEOUT_MS}`);
+
+          const findExpiredKeysSubQuery = trx(TableName.KeyValueStore)
+            .where("expiresAt", "<", db.fn.now())
+            .select("key")
+            .limit(CACHE_KEY_PRUNE_BATCH_SIZE);
+
+          // eslint-disable-next-line no-await-in-loop
+          const results = await trx(TableName.KeyValueStore)
+            .whereIn("key", findExpiredKeysSubQuery)
+            .del()
+            .returning("key");
+
+          return results;
+        });
+
+        numberOfRetryOnFailure = 0; // reset
+      } catch (error) {
+        numberOfRetryOnFailure += 1;
+        deletedIds = [];
+        logger.error(error, "Failed to clean up db key value");
+      } finally {
+        // eslint-disable-next-line no-await-in-loop
+        await new Promise((resolve) => {
+          setTimeout(resolve, 10); // time to breathe for db
+        });
+      }
+      isRetrying = numberOfRetryOnFailure > 0;
+    } while (deletedIds.length > 0 || (isRetrying && numberOfRetryOnFailure < MAX_RETRY_ON_FAILURE));
+    logger.info(`${QueueName.DailyResourceCleanUp}: db key value store clean up completed`);
+  };
+
+  return { ...keyValueStoreOrm, incrementBy, findOneInt, pruneExpiredKeys };
+};

backend/src/keystore/keystore.ts

@@ -1,9 +1,14 @@
+import { Cluster, Redis } from "ioredis";
+import { Knex } from "knex";
+
 import { buildRedisFromConfig, TRedisConfigKeys } from "@app/lib/config/redis";
 import { pgAdvisoryLockHashText } from "@app/lib/crypto/hashtext";
 import { applyJitter } from "@app/lib/dates";
 import { delay as delayMs } from "@app/lib/delay";
+import { ms } from "@app/lib/ms";
 import { ExecutionResult, Redlock, Settings } from "@app/lib/red-lock";
-import { Redis, Cluster } from "ioredis";
+
+import { TKeyValueStoreDALFactory } from "./key-value-store-dal";
 
 export const PgSqlLock = {
   BootUpMigration: 2023,
@@ -97,13 +102,17 @@ export type TKeyStoreFactory = {
   deleteItemsByKeyIn: (keys: string[]) => Promise<number>;
   deleteItems: (arg: TDeleteItems) => Promise<number>;
   incrementBy: (key: string, value: number) => Promise<number>;
+  getKeysByPattern: (pattern: string, limit?: number) => Promise<string[]>;
+  // pg
+  pgIncrementBy: (key: string, dto: { incr?: number; expiry?: string; tx?: Knex }) => Promise<number>;
+  pgGetIntItem: (key: string, prefix?: string) => Promise<number | undefined>;
+  // locks
   acquireLock(
     resources: string[],
     duration: number,
     settings?: Partial<Settings>
   ): Promise<{ release: () => Promise<ExecutionResult> }>;
   waitTillReady: ({ key, waitingCb, keyCheckCb, waitIteration, delay, jitter }: TWaitTillReady) => Promise<void>;
-  getKeysByPattern: (pattern: string, limit?: number) => Promise<string[]>;
 };
 
 const pickPrimaryOrSecondaryRedis = (primary: Redis | Cluster, secondaries?: Array<Redis | Cluster>) => {
@@ -116,7 +125,10 @@ interface TKeyStoreFactoryDTO extends TRedisConfigKeys {
   REDIS_READ_REPLICAS?: { host: string; port: number }[];
 }
 
-export const keyStoreFactory = (redisConfigKeys: TKeyStoreFactoryDTO): TKeyStoreFactory => {
+export const keyStoreFactory = (
+  redisConfigKeys: TKeyStoreFactoryDTO,
+  keyValueStoreDAL: TKeyValueStoreDALFactory
+): TKeyStoreFactory => {
   const primaryRedis = buildRedisFromConfig(redisConfigKeys);
   const redisReadReplicas = redisConfigKeys.REDIS_READ_REPLICAS?.map((el) => {
     if (redisConfigKeys.REDIS_URL) {
@@ -191,29 +203,6 @@ export const keyStoreFactory = (redisConfigKeys: TKeyStoreFactoryDTO): TKeyStore
 
   const setExpiry = async (key: string, expiryInSeconds: number) => primaryRedis.expire(key, expiryInSeconds);
 
-  const waitTillReady = async ({
-    key,
-    waitingCb,
-    keyCheckCb,
-    waitIteration = 10,
-    delay = 1000,
-    jitter = 200
-  }: TWaitTillReady) => {
-    let attempts = 0;
-    let isReady = keyCheckCb(await getItem(key));
-    while (!isReady) {
-      if (attempts > waitIteration) return;
-      // eslint-disable-next-line
-      await new Promise((resolve) => {
-        waitingCb?.();
-        setTimeout(resolve, Math.max(0, applyJitter(delay, jitter)));
-      });
-      attempts += 1;
-      // eslint-disable-next-line
-      isReady = keyCheckCb(await getItem(key));
-    }
-  };
-
   const getKeysByPattern = async (pattern: string, limit?: number) => {
     let cursor = "0";
     const allKeys: string[] = [];
@@ -238,6 +227,37 @@ export const keyStoreFactory = (redisConfigKeys: TKeyStoreFactoryDTO): TKeyStore
     return allKeys;
   };
 
+  const pgIncrementBy: TKeyStoreFactory["pgIncrementBy"] = async (key, { incr = 1, tx, expiry }) => {
+    const expiresAt = expiry ? new Date(Date.now() + ms(expiry)) : undefined;
+    return keyValueStoreDAL.incrementBy(key, { incr, expiresAt, tx });
+  };
+
+  const pgGetIntItem = async (key: string, prefix?: string) =>
+    keyValueStoreDAL.findOneInt(prefix ? `${prefix}:${key}` : key);
+
+  const waitTillReady = async ({
+    key,
+    waitingCb,
+    keyCheckCb,
+    waitIteration = 10,
+    delay = 1000,
+    jitter = 200
+  }: TWaitTillReady) => {
+    let attempts = 0;
+    let isReady = keyCheckCb(await getItem(key));
+    while (!isReady) {
+      if (attempts > waitIteration) return;
+      // eslint-disable-next-line
+      await new Promise((resolve) => {
+        waitingCb?.();
+        setTimeout(resolve, Math.max(0, applyJitter(delay, jitter)));
+      });
+      attempts += 1;
+      // eslint-disable-next-line
+      isReady = keyCheckCb(await getItem(key));
+    }
+  };
+
   return {
     setItem,
     getItem,
@@ -252,6 +272,8 @@ export const keyStoreFactory = (redisConfigKeys: TKeyStoreFactoryDTO): TKeyStore
     waitTillReady,
     getKeysByPattern,
     deleteItemsByKeyIn,
-    getItems
+    getItems,
+    pgGetIntItem,
+    pgIncrementBy
   };
 };

backend/src/keystore/memory.ts

@@ -53,6 +53,15 @@ export const inMemoryKeyStore = (): TKeyStoreFactory => {
       }
       return null;
     },
+    pgGetIntItem: async (key) => {
+      const value = store[key];
+      if (typeof value === "number") {
+        return Number(value);
+      }
+    },
+    pgIncrementBy: async () => {
+      return 1;
+    },
     incrementBy: async () => {
       return 1;
     },