[PB-5850] Add interceptor to the axios client to handle 429

src/infra/drive-server/client/axios.d.ts

@@ -0,0 +1,7 @@
+import { RETRY_CONFIG_KEY } from './drive-server.constants';
+
+declare module 'axios' {
+  interface InternalAxiosRequestConfig {
+    [RETRY_CONFIG_KEY]?: number;
+  }
+}

src/infra/drive-server/client/drive-server.client.instance.test.ts

@@ -1,8 +1,7 @@
 import { createClient } from '../drive-server.client';
-import Bottleneck from 'bottleneck';
 import eventBus from '../../../apps/main/event-bus';
 import { logout } from '../../../apps/main/auth/service';
-import { Mock } from 'vitest';
+import { call } from 'tests/vitest/utils.helper';
 
 vi.mock('../drive-server.client', () => ({
   createClient: vi.fn(() => ({})),
@@ -36,22 +35,19 @@ describe('driveServerClient instance', () => {
 
   it('should call createClient with expected options', async () => {
     await import('./drive-server.client.instance');
-    expect(createClient).toHaveBeenCalledWith(
-      expect.objectContaining({
-        baseUrl: expect.any(String),
-        limiter: expect.any(Bottleneck),
-        onUnauthorized: expect.any(Function),
-      }),
-    );
+    call(createClient).toMatchObject({
+      baseUrl: expect.any(String),
+      onUnauthorized: expect.any(Function),
+    });
   });
 
   it('should call eventBus.emit and logout when onUnauthorized is triggered', async () => {
     await import('./drive-server.client.instance');
-    const clientOptionsArg = (createClient as Mock).mock.calls[0][0];
+    const clientOptions = vi.mocked(createClient).mock.calls[0]![0]!;
 
-    clientOptionsArg.onUnauthorized();
+    clientOptions.onUnauthorized!();
 
-    expect(eventBus.emit).toHaveBeenCalledWith('USER_WAS_UNAUTHORIZED');
+    call(eventBus.emit).toEqual('USER_WAS_UNAUTHORIZED');
     expect(logout).toHaveBeenCalled();
   });
 
@@ -63,8 +59,6 @@ describe('driveServerClient instance', () => {
 
     await import('./drive-server.client.instance');
 
-    const mostRecentCall = (createClient as Mock).mock.calls[(createClient as Mock).mock.calls.length - 1];
-    const clientOptions = mostRecentCall[0];
-    expect(clientOptions.baseUrl).toBe('https://mock.api');
+    call(createClient).toMatchObject({ baseUrl: 'https://mock.api' });
   });
 });

src/infra/drive-server/client/drive-server.client.instance.ts

@@ -1,5 +1,4 @@
 import { paths } from '../../schemas';
-import Bottleneck from 'bottleneck';
 import { logout } from '../../../apps/main/auth/service';
 import eventBus from '../../../apps/main/event-bus';
 import { ClientOptions, createClient } from '../drive-server.client';
@@ -9,14 +8,8 @@ function handleOnUserUnauthorized(): void {
   logout();
 }
 
-const limiter = new Bottleneck({
-  maxConcurrent: 2,
-  minTime: 500,
-});
-
 const clientOptions: ClientOptions = {
   baseUrl: process.env.NEW_DRIVE_URL || '',
-  limiter,
   onUnauthorized: handleOnUserUnauthorized,
 };
 

src/infra/drive-server/client/drive-server.constants.ts

@@ -0,0 +1,2 @@
+export const RETRY_CONFIG_KEY = '__rateLimiterRetryCount';
+export const MAX_RETRIES = 3;

src/infra/drive-server/client/interceptors/rate-limiter/add-jitter.test.ts

@@ -0,0 +1,40 @@
+import { addJitter } from './add-jitter';
+
+describe('addJitter - spreads retry timing to avoid thundering herd', () => {
+  it('should never return less than the base delay', () => {
+    const result = addJitter(1000);
+    expect(result).toBeGreaterThanOrEqual(1000);
+  });
+
+  it('should add up to maxJitter ms on top of the base delay', () => {
+    const result = addJitter(1000, 200);
+    expect(result).toBeGreaterThanOrEqual(1000);
+    expect(result).toBeLessThan(1200);
+  });
+
+  it(' should default maxJitter to 100ms', () => {
+    const result = addJitter(500);
+    expect(result).toBeGreaterThanOrEqual(500);
+    expect(result).toBeLessThan(600);
+  });
+
+  it('should add no jitter when randomness is 0', () => {
+    vi.spyOn(Math, 'random').mockReturnValue(0);
+    expect(addJitter(1000, 200)).toBe(1000);
+  });
+
+  it('should add the maximum jitter when randomness is near 1', () => {
+    vi.spyOn(Math, 'random').mockReturnValue(0.999);
+    expect(addJitter(1000, 100)).toBe(1099);
+  });
+
+  it('should return the base delay exactly when maxJitter is 0', () => {
+    expect(addJitter(500, 0)).toBe(500);
+  });
+
+  it('should work with a base delay of 0', () => {
+    const result = addJitter(0, 50);
+    expect(result).toBeGreaterThanOrEqual(0);
+    expect(result).toBeLessThan(50);
+  });
+});

src/infra/drive-server/client/interceptors/rate-limiter/add-jitter.ts

@@ -0,0 +1,3 @@
+export function addJitter(baseMs: number, maxJitter = 100): number {
+  return baseMs + Math.floor(Math.random() * maxJitter);
+}

src/infra/drive-server/client/interceptors/rate-limiter/attach-rate-limiter-interceptors.test.ts

@@ -0,0 +1,69 @@
+import { type Mock } from 'vitest';
+import { call } from 'tests/vitest/utils.helper';
+import { attachRateLimiterInterceptors } from './attach-rate-limiter-interceptors';
+import { createRequestInterceptor } from './create-request-interceptor';
+import { createResponseInterceptor } from './create-response-interceptor';
+
+vi.mock('./create-request-interceptor');
+vi.mock('./create-response-interceptor');
+
+describe('attachRateLimiterInterceptors', () => {
+  const mockRequestInterceptor = vi.fn();
+  const mockOnFulfilled = vi.fn();
+  const mockOnRejected = vi.fn();
+
+  const mockRequestUse = vi.fn();
+  const mockResponseUse = vi.fn();
+
+  const instance = {
+    interceptors: {
+      request: { use: mockRequestUse },
+      response: { use: mockResponseUse },
+    },
+  } as any;
+
+  beforeEach(() => {
+    (createRequestInterceptor as Mock).mockReturnValue(mockRequestInterceptor);
+    (createResponseInterceptor as Mock).mockReturnValue({
+      onFulfilled: mockOnFulfilled,
+      onRejected: mockOnRejected,
+    });
+  });
+
+  it('should create a request interceptor with a fresh delay state', () => {
+    attachRateLimiterInterceptors(instance);
+
+    call(createRequestInterceptor).toMatchObject({ pending: null });
+  });
+
+  it('should register the request interceptor on the instance', () => {
+    attachRateLimiterInterceptors(instance);
+
+    call(mockRequestUse).toMatchObject(mockRequestInterceptor);
+  });
+
+  it('should create a response interceptor with the instance, fresh rate limit state, and delay state', () => {
+    attachRateLimiterInterceptors(instance);
+
+    call(createResponseInterceptor).toMatchObject([
+      instance,
+      { limit: null, remaining: null, reset: null },
+      { pending: null },
+    ]);
+  });
+
+  it('should register the response interceptor on the instance', () => {
+    attachRateLimiterInterceptors(instance);
+
+    call(mockResponseUse).toMatchObject([mockOnFulfilled, mockOnRejected]);
+  });
+
+  it('should share the same delay state between request and response interceptors', () => {
+    const delayState = { pending: null };
+
+    attachRateLimiterInterceptors(instance);
+
+    call(createRequestInterceptor).toMatchObject(delayState);
+    call(createResponseInterceptor).toMatchObject([expect.anything(), expect.anything(), delayState]);
+  });
+});

src/infra/drive-server/client/interceptors/rate-limiter/attach-rate-limiter-interceptors.ts

@@ -0,0 +1,14 @@
+import type { AxiosInstance } from 'axios';
+import { DelayState, RateLimitState } from './rate-limiter.types';
+import { createRequestInterceptor } from './create-request-interceptor';
+import { createResponseInterceptor } from './create-response-interceptor';
+
+export function attachRateLimiterInterceptors(instance: AxiosInstance): void {
+  const state: RateLimitState = { limit: null, remaining: null, reset: null };
+  const delayState: DelayState = { pending: null };
+
+  instance.interceptors.request.use(createRequestInterceptor(delayState));
+
+  const { onFulfilled, onRejected } = createResponseInterceptor(instance, state, delayState);
+  instance.interceptors.response.use(onFulfilled, onRejected);
+}

src/infra/drive-server/client/interceptors/rate-limiter/create-request-interceptor.test.ts

@@ -0,0 +1,64 @@
+import type { InternalAxiosRequestConfig } from 'axios';
+import { createRequestInterceptor } from './create-request-interceptor';
+import { DelayState } from './rate-limiter.types';
+
+describe('createRequestInterceptor', () => {
+  const mockConfig = { url: '/test' } as InternalAxiosRequestConfig;
+
+  it('should return the config immediately when there is no pending delay', async () => {
+    const state: DelayState = { pending: null };
+    const interceptor = createRequestInterceptor(state);
+
+    const result = await interceptor(mockConfig);
+
+    expect(result).toBe(mockConfig);
+  });
+
+  it('should wait for the pending delay before returning the config', async () => {
+    let resolveDelay!: () => void;
+    const state: DelayState = {
+      pending: new Promise((resolve) => {
+        resolveDelay = resolve;
+      }),
+    };
+    const interceptor = createRequestInterceptor(state);
+
+    let resolved = false;
+    const resultPromise = interceptor(mockConfig).then((config) => {
+      resolved = true;
+      return config;
+    });
+
+    await Promise.resolve();
+    expect(resolved).toBe(false);
+
+    resolveDelay();
+    const result = await resultPromise;
+
+    expect(resolved).toBe(true);
+    expect(result).toBe(mockConfig);
+  });
+
+  it('should make multiple concurrent requests wait for the same delay', async () => {
+    let resolveDelay!: () => void;
+    const state: DelayState = {
+      pending: new Promise((resolve) => {
+        resolveDelay = resolve;
+      }),
+    };
+    const interceptor = createRequestInterceptor(state);
+
+    const configA = { url: '/a' } as InternalAxiosRequestConfig;
+    const configB = { url: '/b' } as InternalAxiosRequestConfig;
+
+    const promiseA = interceptor(configA);
+    const promiseB = interceptor(configB);
+
+    resolveDelay();
+
+    const [resultA, resultB] = await Promise.all([promiseA, promiseB]);
+
+    expect(resultA).toBe(configA);
+    expect(resultB).toBe(configB);
+  });
+});

src/infra/drive-server/client/interceptors/rate-limiter/create-request-interceptor.ts

@@ -0,0 +1,14 @@
+import type { InternalAxiosRequestConfig } from 'axios';
+import { DelayState } from './rate-limiter.types';
+
+export function createRequestInterceptor(
+  delayState: DelayState,
+): (config: InternalAxiosRequestConfig) => Promise<InternalAxiosRequestConfig> {
+  return async (config: InternalAxiosRequestConfig) => {
+    if (delayState.pending) {
+      await delayState.pending;
+    }
+
+    return config;
+  };
+}