fix(FIR-44193): auth concurrency (#139)

Author: ptiurin

package-lock.json

@@ -24,6 +24,7 @@
   "devDependencies": {
     "@types/jest": "^29.5.12",
     "@types/node-fetch": "^2.5.12",
+    "@types/rwlock": "^5.0.6",
     "@typescript-eslint/eslint-plugin": "^5.4.0",
     "@typescript-eslint/parser": "^5.4.0",
     "allure-commandline": "^2.29.0",
@@ -46,6 +47,7 @@
     "abort-controller": "^3.0.0",
     "agentkeepalive": "^4.5.0",
     "json-bigint": "^1.0.0",
-    "node-fetch": "^2.6.6"
+    "node-fetch": "^2.6.6",
+    "rwlock": "^5.0.0"
   }
 }

package.json

@@ -7,6 +7,7 @@ import {
   UsernamePasswordAuth
 } from "../types";
 import { TokenKey, inMemoryCache, noneCache } from "../common/tokenCache";
+import ReadWriteLock from "rwlock";
 
 type Login = {
   access_token: string;
@@ -22,6 +23,7 @@ export class Authenticator {
   options: ConnectionOptions;
 
   accessToken?: string;
+  private readonly rwlock = new ReadWriteLock();
 
   constructor(context: Context, options: ConnectionOptions) {
     context.httpClient.authenticator = this;
@@ -157,21 +159,65 @@ export class Authenticator {
     );
   }
 
-  async authenticate() {
-    const options = this.options.auth || this.options;
-    const cachedToken = this.getCachedToken();
+  async authenticate(): Promise<void> {
+    // Try to get token from cache using read lock
+    const cachedToken = await this.tryGetCachedToken();
     if (cachedToken) {
       this.accessToken = cachedToken;
       return;
     }
 
+    // No cached token, acquire write lock and authenticate
+    await this.acquireWriteLockAndAuthenticate();
+  }
+
+  private async tryGetCachedToken(): Promise<string | undefined> {
+    return new Promise((resolve, reject) => {
+      this.rwlock.readLock(releaseReadLock => {
+        try {
+          const cachedToken = this.getCachedToken();
+          resolve(cachedToken);
+        } catch (error) {
+          reject(error instanceof Error ? error : new Error(String(error)));
+        } finally {
+          releaseReadLock();
+        }
+      });
+    });
+  }
+
+  private async acquireWriteLockAndAuthenticate(): Promise<void> {
+    return new Promise((resolve, reject) => {
+      this.rwlock.writeLock(async releaseWriteLock => {
+        try {
+          // Double-check cache in case another thread authenticated while waiting
+          const cachedToken = this.getCachedToken();
+          if (cachedToken) {
+            this.accessToken = cachedToken;
+            return resolve();
+          }
+
+          await this.performAuthentication();
+
+          resolve();
+        } catch (error) {
+          reject(error instanceof Error ? error : new Error(String(error)));
+        } finally {
+          releaseWriteLock();
+        }
+      });
+    });
+  }
+
+  private async performAuthentication(): Promise<void> {
+    const options = this.options.auth || this.options;
+
     if (this.isUsernamePassword()) {
-      await this.authenticateUsernamePassword(options as UsernamePasswordAuth);
-      return;
+      return this.authenticateUsernamePassword(options as UsernamePasswordAuth);
     }
+
     if (this.isServiceAccount()) {
-      await this.authenticateServiceAccount(options as ServiceAccountAuth);
-      return;
+      return this.authenticateServiceAccount(options as ServiceAccountAuth);
     }
 
     throw new Error("Please provide valid auth credentials");

src/auth/index.ts

@@ -2,6 +2,7 @@ import { makeConnection } from "../connection";
 import { Authenticator } from "../auth";
 import { Context, ConnectionOptions, FireboltClientOptions } from "../types";
 import { ResourceManager } from "../service";
+import { NodeHttpClient } from "../http/node";
 
 export class FireboltCore {
   private options: FireboltClientOptions;
@@ -13,24 +14,39 @@ export class FireboltCore {
     this.options = options;
   }
 
-  async connect(connectionOptions: ConnectionOptions) {
-    const auth = new Authenticator(this.context, connectionOptions);
-    const connection = makeConnection(this.context, connectionOptions);
+  private async prepareConnection(connectionOptions: ConnectionOptions) {
+    // Create a new httpClient instance for each connection
+    const httpClient =
+      this.options.dependencies?.httpClient || new NodeHttpClient();
+
+    // Create a new context with the new httpClient
+    const connectionContext = {
+      ...this.context,
+      httpClient
+    };
+
+    const auth = new Authenticator(connectionContext, connectionOptions);
+    const connection = makeConnection(connectionContext, connectionOptions);
     await auth.authenticate();
     await connection.resolveEngineEndpoint();
-    const context = {
+
+    return { connection, connectionContext };
+  }
+
+  async connect(connectionOptions: ConnectionOptions) {
+    const { connection, connectionContext } = await this.prepareConnection(
+      connectionOptions
+    );
+    const resourceContext = {
       connection,
-      ...this.context
+      ...connectionContext
     };
-    this.resourceManager = new ResourceManager(context);
+    this.resourceManager = new ResourceManager(resourceContext);
     return connection;
   }
 
   async testConnection(connectionOptions: ConnectionOptions) {
-    const auth = new Authenticator(this.context, connectionOptions);
-    const connection = makeConnection(this.context, connectionOptions);
-    await auth.authenticate();
-    await connection.resolveEngineEndpoint();
+    const { connection } = await this.prepareConnection(connectionOptions);
     await connection.testConnection();
   }
 }

src/core/index.ts

@@ -136,14 +136,17 @@ describe.each([
   ]
 ])("token caching %s", (_, authUrl: string, auth: AuthOptions) => {
   const server = setupServer();
-  const httpClient = new NodeHttpClient();
+  let httpClient = new NodeHttpClient();
 
   beforeAll(() => {
     server.listen();
   });
   afterAll(() => {
     server.close();
   });
+  beforeEach(() => {
+    httpClient = new NodeHttpClient();
+  });
 
   it("caches and reuses access token", async () => {
     const authenticator = new Authenticator(
@@ -163,7 +166,7 @@ describe.each([
         return res(
           ctx.json({
             access_token: "fake_access_token",
-            expires_in: 2 ^ 30
+            expires_in: 2 ** 30
           })
         );
       })
@@ -230,7 +233,7 @@ describe.each([
         return res(
           ctx.json({
             access_token: "fake_access_token",
-            expires_in: 2 ^ 30
+            expires_in: 2 ** 30
           })
         );
       }),
@@ -273,7 +276,7 @@ describe.each([
         return res(
           ctx.json({
             access_token: "fake_access_token",
-            expires_in: 2 ^ 30
+            expires_in: 2 ** 30
           })
         );
       })
@@ -295,10 +298,11 @@ describe.each([
         useCache: true
       }
     );
-
+    // Different Authenticator has to have different httpClient
+    const httpClient2 = new NodeHttpClient();
     const authenticator2 = new Authenticator(
       {
-        httpClient,
+        httpClient: httpClient2,
         apiEndpoint: "api.fake2.firebolt.io",
         logger
       },
@@ -318,7 +322,7 @@ describe.each([
         return res(
           ctx.json({
             access_token: "fake_access_token",
-            expires_in: 2 ^ 30
+            expires_in: 2 ** 30
           })
         );
       }),
@@ -327,7 +331,7 @@ describe.each([
         return res(
           ctx.json({
             access_token: "fake_access_token",
-            expires_in: 2 ^ 30
+            expires_in: 2 ** 30
           })
         );
       })
@@ -342,4 +346,37 @@ describe.each([
     expect(authenticator2.accessToken).toEqual("fake_access_token");
     expect(calls).toEqual(2);
   });
+  it("does not overwrite token when run concurrently", async () => {
+    const authenticator = new Authenticator(
+      { httpClient, apiEndpoint, logger },
+      {
+        auth,
+        account: "my_account",
+        useCache: true
+      }
+    );
+
+    let calls = 0;
+
+    server.use(
+      rest.post(authUrl, (req, res, ctx) => {
+        calls++;
+        return res(
+          ctx.json({
+            access_token: "fake_access_token",
+            expires_in: 2 ** 30
+          })
+        );
+      })
+    );
+
+    authenticator.clearCache();
+
+    const promises = Array(10)
+      .fill(null)
+      .map(() => authenticator.authenticate());
+
+    await Promise.all(promises);
+    expect(calls).toEqual(1);
+  });
 });