Merge pull request #2317 from akhilmhdh/feat/debounce-last-used

feat: added identity and service token postgres update debounced

backend/src/services/access-token-queue/access-token-queue.ts

@@ -0,0 +1,125 @@
+import { z } from "zod";
+
+import { KeyStorePrefixes, KeyStoreTtls, TKeyStoreFactory } from "@app/keystore/keystore";
+import { applyJitter, secondsToMillis } from "@app/lib/dates";
+import { logger } from "@app/lib/logger";
+import { QueueJobs, QueueName, TQueueServiceFactory } from "@app/queue";
+
+import { TIdentityAccessTokenDALFactory } from "../identity-access-token/identity-access-token-dal";
+import { TServiceTokenDALFactory } from "../service-token/service-token-dal";
+
+type TAccessTokenQueueServiceFactoryDep = {
+  queueService: TQueueServiceFactory;
+  keyStore: Pick<TKeyStoreFactory, "getItem" | "setItemWithExpiry">;
+  identityAccessTokenDAL: Pick<TIdentityAccessTokenDALFactory, "updateById">;
+  serviceTokenDAL: Pick<TServiceTokenDALFactory, "updateById">;
+};
+
+export type TAccessTokenQueueServiceFactory = ReturnType<typeof accessTokenQueueServiceFactory>;
+
+export const AccessTokenStatusSchema = z.object({
+  lastUpdatedAt: z.string().datetime(),
+  numberOfUses: z.number()
+});
+
+export const accessTokenQueueServiceFactory = ({
+  queueService,
+  keyStore,
+  identityAccessTokenDAL,
+  serviceTokenDAL
+}: TAccessTokenQueueServiceFactoryDep) => {
+  const getIdentityTokenDetailsInCache = async (identityAccessTokenId: string) => {
+    const tokenDetailsInCache = await keyStore.getItem(
+      KeyStorePrefixes.IdentityAccessTokenStatusUpdate(identityAccessTokenId)
+    );
+    if (tokenDetailsInCache) {
+      return AccessTokenStatusSchema.parseAsync(JSON.parse(tokenDetailsInCache));
+    }
+  };
+
+  const updateServiceTokenStatus = async (serviceTokenId: string) => {
+    await keyStore.setItemWithExpiry(
+      KeyStorePrefixes.ServiceTokenStatusUpdate(serviceTokenId),
+      KeyStoreTtls.AccessTokenStatusUpdateInSeconds,
+      JSON.stringify({ lastUpdatedAt: new Date() })
+    );
+    await queueService.queue(
+      QueueName.AccessTokenStatusUpdate,
+      QueueJobs.ServiceTokenStatusUpdate,
+      {
+        serviceTokenId
+      },
+      {
+        delay: applyJitter(secondsToMillis(KeyStoreTtls.AccessTokenStatusUpdateInSeconds / 2), secondsToMillis(10)),
+        // https://docs.bullmq.io/guide/jobs/job-ids
+        jobId: KeyStorePrefixes.ServiceTokenStatusUpdate(serviceTokenId).replaceAll(":", "_"),
+        removeOnFail: true,
+        removeOnComplete: true
+      }
+    );
+  };
+
+  const updateIdentityAccessTokenStatus = async (identityAccessTokenId: string, numberOfUses: number) => {
+    await keyStore.setItemWithExpiry(
+      KeyStorePrefixes.IdentityAccessTokenStatusUpdate(identityAccessTokenId),
+      KeyStoreTtls.AccessTokenStatusUpdateInSeconds,
+      JSON.stringify({ lastUpdatedAt: new Date(), numberOfUses })
+    );
+    await queueService.queue(
+      QueueName.AccessTokenStatusUpdate,
+      QueueJobs.IdentityAccessTokenStatusUpdate,
+      {
+        identityAccessTokenId,
+        numberOfUses
+      },
+      {
+        delay: applyJitter(secondsToMillis(KeyStoreTtls.AccessTokenStatusUpdateInSeconds / 2), secondsToMillis(10)),
+        jobId: KeyStorePrefixes.IdentityAccessTokenStatusUpdate(identityAccessTokenId).replaceAll(":", "_"),
+        removeOnFail: true,
+        removeOnComplete: true
+      }
+    );
+  };
+
+  queueService.start(QueueName.AccessTokenStatusUpdate, async (job) => {
+    // for identity token update
+    if (job.name === QueueJobs.IdentityAccessTokenStatusUpdate && "identityAccessTokenId" in job.data) {
+      const { identityAccessTokenId } = job.data;
+      const tokenDetails = { lastUpdatedAt: new Date(job.timestamp), numberOfUses: job.data.numberOfUses };
+      const tokenDetailsInCache = await getIdentityTokenDetailsInCache(identityAccessTokenId);
+      if (tokenDetailsInCache) {
+        tokenDetails.numberOfUses = tokenDetailsInCache.numberOfUses;
+        tokenDetails.lastUpdatedAt = new Date(tokenDetailsInCache.lastUpdatedAt);
+      }
+
+      await identityAccessTokenDAL.updateById(identityAccessTokenId, {
+        accessTokenLastUsedAt: tokenDetails.lastUpdatedAt,
+        accessTokenNumUses: tokenDetails.numberOfUses
+      });
+      return;
+    }
+
+    // for service token
+    if (job.name === QueueJobs.ServiceTokenStatusUpdate && "serviceTokenId" in job.data) {
+      const { serviceTokenId } = job.data;
+      const tokenDetailsInCache = await keyStore.getItem(KeyStorePrefixes.ServiceTokenStatusUpdate(serviceTokenId));
+      let lastUsed = new Date(job.timestamp);
+      if (tokenDetailsInCache) {
+        const tokenDetails = await AccessTokenStatusSchema.pick({ lastUpdatedAt: true }).parseAsync(
+          JSON.parse(tokenDetailsInCache)
+        );
+        lastUsed = new Date(tokenDetails.lastUpdatedAt);
+      }
+
+      await serviceTokenDAL.updateById(serviceTokenId, {
+        lastUsed
+      });
+    }
+  });
+
+  queueService.listen(QueueName.AccessTokenStatusUpdate, "failed", (_, err) => {
+    logger.error(err, `${QueueName.AccessTokenStatusUpdate}: Failed to updated access token status`);
+  });
+
+  return { updateIdentityAccessTokenStatus, updateServiceTokenStatus, getIdentityTokenDetailsInCache };
+};

backend/src/services/identity-access-token/identity-access-token-service.ts

@@ -5,6 +5,7 @@ import { getConfig } from "@app/lib/config/env";
 import { BadRequestError, UnauthorizedError } from "@app/lib/errors";
 import { checkIPAgainstBlocklist, TIp } from "@app/lib/ip";
 
+import { TAccessTokenQueueServiceFactory } from "../access-token-queue/access-token-queue";
 import { AuthTokenType } from "../auth/auth-type";
 import { TIdentityOrgDALFactory } from "../identity/identity-org-dal";
 import { TIdentityAccessTokenDALFactory } from "./identity-access-token-dal";
@@ -13,19 +14,24 @@ import { TIdentityAccessTokenJwtPayload, TRenewAccessTokenDTO } from "./identity
 type TIdentityAccessTokenServiceFactoryDep = {
   identityAccessTokenDAL: TIdentityAccessTokenDALFactory;
   identityOrgMembershipDAL: TIdentityOrgDALFactory;
+  accessTokenQueue: Pick<
+    TAccessTokenQueueServiceFactory,
+    "updateIdentityAccessTokenStatus" | "getIdentityTokenDetailsInCache"
+  >;
 };
 
 export type TIdentityAccessTokenServiceFactory = ReturnType<typeof identityAccessTokenServiceFactory>;
 
 export const identityAccessTokenServiceFactory = ({
   identityAccessTokenDAL,
-  identityOrgMembershipDAL
+  identityOrgMembershipDAL,
+  accessTokenQueue
 }: TIdentityAccessTokenServiceFactoryDep) => {
   const validateAccessTokenExp = async (identityAccessToken: TIdentityAccessTokens) => {
     const {
       id: tokenId,
-      accessTokenTTL,
       accessTokenNumUses,
+      accessTokenTTL,
       accessTokenNumUsesLimit,
       accessTokenLastRenewedAt,
       createdAt: accessTokenCreatedAt
@@ -83,7 +89,12 @@ export const identityAccessTokenServiceFactory = ({
     });
     if (!identityAccessToken) throw new UnauthorizedError();
 
-    await validateAccessTokenExp(identityAccessToken);
+    let { accessTokenNumUses } = identityAccessToken;
+    const tokenStatusInCache = await accessTokenQueue.getIdentityTokenDetailsInCache(identityAccessToken.id);
+    if (tokenStatusInCache) {
+      accessTokenNumUses = tokenStatusInCache.numberOfUses;
+    }
+    await validateAccessTokenExp({ ...identityAccessToken, accessTokenNumUses });
 
     const { accessTokenMaxTTL, createdAt: accessTokenCreatedAt, accessTokenTTL } = identityAccessToken;
 
@@ -164,14 +175,14 @@ export const identityAccessTokenServiceFactory = ({
       throw new UnauthorizedError({ message: "Identity does not belong to any organization" });
     }
 
-    await validateAccessTokenExp(identityAccessToken);
+    let { accessTokenNumUses } = identityAccessToken;
+    const tokenStatusInCache = await accessTokenQueue.getIdentityTokenDetailsInCache(identityAccessToken.id);
+    if (tokenStatusInCache) {
+      accessTokenNumUses = tokenStatusInCache.numberOfUses;
+    }
+    await validateAccessTokenExp({ ...identityAccessToken, accessTokenNumUses });
 
-    await identityAccessTokenDAL.updateById(identityAccessToken.id, {
-      accessTokenLastUsedAt: new Date(),
-      $incr: {
-        accessTokenNumUses: 1
-      }
-    });
+    await accessTokenQueue.updateIdentityAccessTokenStatus(identityAccessToken.id, Number(accessTokenNumUses) + 1);
     return { ...identityAccessToken, orgId: identityOrgMembership.orgId };
   };
 

backend/src/services/service-token/service-token-service.ts

@@ -8,6 +8,7 @@ import { ProjectPermissionActions, ProjectPermissionSub } from "@app/ee/services
 import { getConfig } from "@app/lib/config/env";
 import { BadRequestError, UnauthorizedError } from "@app/lib/errors";
 
+import { TAccessTokenQueueServiceFactory } from "../access-token-queue/access-token-queue";
 import { ActorType } from "../auth/auth-type";
 import { TProjectDALFactory } from "../project/project-dal";
 import { TProjectEnvDALFactory } from "../project-env/project-env-dal";
@@ -26,6 +27,7 @@ type TServiceTokenServiceFactoryDep = {
   permissionService: Pick<TPermissionServiceFactory, "getProjectPermission">;
   projectEnvDAL: Pick<TProjectEnvDALFactory, "findBySlugs">;
   projectDAL: Pick<TProjectDALFactory, "findById">;
+  accessTokenQueue: Pick<TAccessTokenQueueServiceFactory, "updateServiceTokenStatus">;
 };
 
 export type TServiceTokenServiceFactory = ReturnType<typeof serviceTokenServiceFactory>;
@@ -35,7 +37,8 @@ export const serviceTokenServiceFactory = ({
   userDAL,
   permissionService,
   projectEnvDAL,
-  projectDAL
+  projectDAL,
+  accessTokenQueue
 }: TServiceTokenServiceFactoryDep) => {
   const createServiceToken = async ({
     iv,
@@ -166,11 +169,9 @@ export const serviceTokenServiceFactory = ({
 
     const isMatch = await bcrypt.compare(TOKEN_SECRET, serviceToken.secretHash);
     if (!isMatch) throw new UnauthorizedError();
-    const updatedToken = await serviceTokenDAL.updateById(serviceToken.id, {
-      lastUsed: new Date()
-    });
+    await accessTokenQueue.updateServiceTokenStatus(serviceToken.id);
 
-    return { ...serviceToken, lastUsed: updatedToken.lastUsed, orgId: project.orgId };
+    return { ...serviceToken, lastUsed: new Date(), orgId: project.orgId };
   };
 
   return {