optimized UA login for higher throughput

backend/src/services/identity-ua/identity-ua-service.ts

@@ -84,23 +84,6 @@ export const identityUaServiceFactory = ({
 
     const LOCKOUT_KEY = `lockout:identity:${identityUa.identityId}:${IdentityAuthMethod.UNIVERSAL_AUTH}:${clientId}`;
 
-    let lock: Awaited<ReturnType<typeof keyStore.acquireLock>> | undefined;
-    if (identityUa.lockoutEnabled) {
-      try {
-        lock = await keyStore.acquireLock([KeyStorePrefixes.IdentityLockoutLock(LOCKOUT_KEY)], 500, {
-          retryCount: 3,
-          retryDelay: 300,
-          retryJitter: 100
-        });
-      } catch (e) {
-        logger.info(
-          `identity login failed to acquire lock [identityId=${identityUa.identityId}] [authMethod=${IdentityAuthMethod.UNIVERSAL_AUTH}]`
-        );
-        throw new RateLimitError({ message: "Failed to acquire lock: rate limit exceeded" });
-      }
-    }
-
-    try {
     const lockoutRaw = await keyStore.getItem(LOCKOUT_KEY);
 
     let lockout: LockoutObject | undefined;
@@ -140,13 +123,31 @@ export const identityUaServiceFactory = ({
 
     if (!validClientSecretInfo) {
       if (identityUa.lockoutEnabled) {
-          if (!lockout) {
+        let lock: Awaited<ReturnType<typeof keyStore.acquireLock>> | undefined;
+        try {
+          lock = await keyStore.acquireLock([KeyStorePrefixes.IdentityLockoutLock(LOCKOUT_KEY)], 300, {
+            retryCount: 3,
+            retryDelay: 300,
+            retryJitter: 100
+          });
+
+          // Re-fetch the latest lockout data while holding the lock
+          const lockoutRawNew = await keyStore.getItem(LOCKOUT_KEY);
+          if (lockoutRawNew) {
+            lockout = JSON.parse(lockoutRawNew) as LockoutObject;
+          } else {
             lockout = {
               lockedOut: false,
               failedAttempts: 0
             };
           }
 
+          if (lockout.lockedOut) {
+            throw new UnauthorizedError({
+              message: "This identity auth method is temporarily locked, please try again later"
+            });
+          }
+
           lockout.failedAttempts += 1;
           if (lockout.failedAttempts >= identityUa.lockoutThreshold) {
             lockout.lockedOut = true;
@@ -157,10 +158,24 @@ export const identityUaServiceFactory = ({
             lockout.lockedOut ? identityUa.lockoutDurationSeconds : identityUa.lockoutCounterResetSeconds,
             JSON.stringify(lockout)
           );
+        } catch (e) {
+          if (lock === undefined) {
+            logger.info(
+              `identity login failed to acquire lock [identityId=${identityUa.identityId}] [authMethod=${IdentityAuthMethod.UNIVERSAL_AUTH}]`
+            );
+            throw new RateLimitError({ message: "Failed to acquire lock: rate limit exceeded" });
+          }
+          throw e;
+        } finally {
+          if (lock) {
+            await lock.release();
+          }
+        }
       }
 
       throw new UnauthorizedError({ message: "Invalid credentials" });
     } else if (lockout) {
+      // If credentials are valid, clear any existing lockout record
       await keyStore.deleteItem(LOCKOUT_KEY);
     }
 
@@ -258,9 +273,6 @@ export const identityUaServiceFactory = ({
       identityMembershipOrg,
       ...accessTokenTTLParams
     };
-    } finally {
-      if (lock) await lock.release();
-    }
   };
 
   const attachUniversalAuth = async ({