Throttling in common api (#15338)

2025-10-24 17:32:51 +02:00
parent d7bda9576f
commit 93ab1c4118
11 changed files with 295 additions and 14 deletions
@@ -0,0 +1,148 @@
+import { Test, type TestingModule } from '@nestjs/testing';
+
+import { type CacheStorageService } from 'src/engine/core-modules/cache-storage/services/cache-storage.service';
+import { CacheStorageNamespace } from 'src/engine/core-modules/cache-storage/types/cache-storage-namespace.enum';
+import { ThrottlerException } from 'src/engine/core-modules/throttler/throttler.exception';
+import { ThrottlerService } from 'src/engine/core-modules/throttler/throttler.service';
+
+describe('ThrottlerService', () => {
+  let service: ThrottlerService;
+  let cacheStorageService: jest.Mocked<CacheStorageService>;
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      providers: [
+        ThrottlerService,
+        {
+          provide: CacheStorageNamespace.EngineWorkspace,
+          useValue: {
+            get: jest.fn(),
+            set: jest.fn(),
+          },
+        },
+      ],
+    }).compile();
+
+    service = module.get<ThrottlerService>(ThrottlerService);
+    cacheStorageService = module.get(CacheStorageNamespace.EngineWorkspace);
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  describe('tokenBucketThrottle', () => {
+    const key = 'test-throttle-key';
+    const maxTokens = 100;
+    const timeWindow = 1000; // 1 second
+
+    it('should allow request when tokens are available (first request)', async () => {
+      cacheStorageService.get.mockResolvedValue(null);
+
+      await service.tokenBucketThrottle(key, 10, maxTokens, timeWindow);
+
+      expect(cacheStorageService.get).toHaveBeenCalledWith(key);
+      expect(cacheStorageService.set).toHaveBeenCalledWith(
+        key,
+        {
+          tokens: 90, // maxTokens - tokensToConsume
+          lastRefillAt: expect.any(Number),
+        },
+        timeWindow * 2,
+      );
+    });
+
+    it('should allow request when sufficient tokens are available', async () => {
+      const now = Date.now();
+
+      cacheStorageService.get.mockResolvedValue({
+        tokens: 50,
+        lastRefillAt: now - 100,
+      });
+
+      await service.tokenBucketThrottle(key, 10, maxTokens, timeWindow);
+
+      expect(cacheStorageService.get).toHaveBeenCalledWith(key);
+      expect(cacheStorageService.set).toHaveBeenCalledWith(
+        key,
+        {
+          tokens: expect.any(Number),
+          lastRefillAt: expect.any(Number),
+        },
+        timeWindow * 2,
+      );
+    });
+
+    it('should throw ThrottlerException when tokens are insufficient', async () => {
+      const now = Date.now();
+
+      cacheStorageService.get.mockResolvedValue({
+        tokens: 5,
+        lastRefillAt: now,
+      });
+
+      await expect(
+        service.tokenBucketThrottle(key, 10, maxTokens, timeWindow),
+      ).rejects.toThrow(ThrottlerException);
+
+      await expect(
+        service.tokenBucketThrottle(key, 10, maxTokens, timeWindow),
+      ).rejects.toThrow('Limit reached');
+
+      expect(cacheStorageService.set).not.toHaveBeenCalled();
+    });
+
+    it('should refill tokens over time', async () => {
+      const now = Date.now();
+      const refillRate = maxTokens / timeWindow; // 100 tokens per 1000ms = 0.1 tokens/ms
+      const timePassed = 500; // 500ms
+      const expectedRefill = Math.floor(timePassed * refillRate); // 50 tokens
+
+      cacheStorageService.get.mockResolvedValue({
+        tokens: 20,
+        lastRefillAt: now - timePassed,
+      });
+
+      jest.spyOn(Date, 'now').mockReturnValue(now);
+
+      await service.tokenBucketThrottle(key, 10, maxTokens, timeWindow);
+
+      expect(cacheStorageService.set).toHaveBeenCalledWith(
+        key,
+        {
+          tokens: 20 + expectedRefill - 10, // 20 + 50 - 10 = 60
+          lastRefillAt: now,
+        },
+        timeWindow * 2,
+      );
+
+      jest.restoreAllMocks();
+    });
+
+    it('should cap tokens at maxTokens', async () => {
+      const now = Date.now();
+      const timePassed = 5000; // Long time passed, would refill beyond max
+
+      cacheStorageService.get.mockResolvedValue({
+        tokens: 80,
+        lastRefillAt: now - timePassed,
+      });
+
+      jest.spyOn(Date, 'now').mockReturnValue(now);
+
+      await service.tokenBucketThrottle(key, 10, maxTokens, timeWindow);
+
+      // Available tokens should be capped at maxTokens (100)
+      expect(cacheStorageService.set).toHaveBeenCalledWith(
+        key,
+        {
+          tokens: 90, // maxTokens (100) - tokensToConsume (10)
+          lastRefillAt: now,
+        },
+        timeWindow * 2,
+      );
+
+      jest.restoreAllMocks();
+    });
+  });
+});
@@ -27,4 +27,39 @@ export class ThrottlerService {

    await this.cacheStorage.set(key, currentCount + 1, ttl);
  }
+
+  async tokenBucketThrottle(
+    key: string,
+    tokensToConsume: number,
+    maxTokens: number,
+    timeWindow: number,
+  ): Promise<void> {
+    const now = Date.now();
+    const refillRate = maxTokens / timeWindow;
+
+    const { tokens, lastRefillAt } = (await this.cacheStorage.get<{
+      tokens: number;
+      lastRefillAt: number;
+    }>(key)) || { tokens: maxTokens, lastRefillAt: now };
+
+    const refillAmount = Math.floor((now - lastRefillAt) * refillRate);
+
+    const availableTokens = Math.min(tokens + refillAmount, maxTokens);
+
+    if (availableTokens < tokensToConsume) {
+      throw new ThrottlerException(
+        `Limit reached (${maxTokens} tokens per ${timeWindow} ms)`,
+        ThrottlerExceptionCode.LIMIT_REACHED,
+      );
+    }
+
+    await this.cacheStorage.set(
+      key,
+      {
+        tokens: availableTokens - tokensToConsume,
+        lastRefillAt: now,
+      },
+      timeWindow * 2,
+    );
+  }
 }
@@ -0,0 +1,19 @@
+import { assertUnreachable } from 'twenty-shared/utils';
+
+import { UserInputError } from 'src/engine/core-modules/graphql/utils/graphql-errors.util';
+import {
+  type ThrottlerException,
+  ThrottlerExceptionCode,
+} from 'src/engine/core-modules/throttler/throttler.exception';
+
+export const throttlerToGraphqlApiExceptionHandler = (
+  error: ThrottlerException,
+) => {
+  switch (error.code) {
+    case ThrottlerExceptionCode.LIMIT_REACHED:
+      throw new UserInputError(error);
+    default: {
+      return assertUnreachable(error.code);
+    }
+  }
+};
@@ -0,0 +1,20 @@
+import { HttpException, HttpStatus } from '@nestjs/common';
+
+import { assertUnreachable } from 'twenty-shared/utils';
+
+import {
+  type ThrottlerException,
+  ThrottlerExceptionCode,
+} from 'src/engine/core-modules/throttler/throttler.exception';
+
+export const throttlerToRestApiExceptionHandler = (
+  error: ThrottlerException,
+): never => {
+  switch (error.code) {
+    case ThrottlerExceptionCode.LIMIT_REACHED:
+      throw new HttpException(error.message, HttpStatus.TOO_MANY_REQUESTS);
+    default: {
+      return assertUnreachable(error.code);
+    }
+  }
+};
@@ -945,20 +945,37 @@ export class ConfigVariables {

  @ConfigVariablesMetadata({
    group: ConfigVariablesGroup.RATE_LIMITING,
-    description: 'Time-to-live for API rate limiting in milliseconds',
+    description: 'Time-to-live for short API rate limiting in milliseconds',
    type: ConfigVariableType.NUMBER,
  })
  @CastToPositiveNumber()
-  API_RATE_LIMITING_TTL = 100;
+  API_RATE_LIMITING_SHORT_TTL_IN_MS = 1000;

  @ConfigVariablesMetadata({
    group: ConfigVariablesGroup.RATE_LIMITING,
    description:
-      'Maximum number of requests allowed in the rate limiting window',
+      'Maximum number of requests allowed in the short rate limiting window',
    type: ConfigVariableType.NUMBER,
  })
  @CastToPositiveNumber()
-  API_RATE_LIMITING_LIMIT = 500;
+  API_RATE_LIMITING_SHORT_LIMIT = 100;
+
+  @ConfigVariablesMetadata({
+    group: ConfigVariablesGroup.RATE_LIMITING,
+    description: 'Time-to-live for long API rate limiting in milliseconds',
+    type: ConfigVariableType.NUMBER,
+  })
+  @CastToPositiveNumber()
+  API_RATE_LIMITING_LONG_TTL_IN_MS = 60000;
+
+  @ConfigVariablesMetadata({
+    group: ConfigVariablesGroup.RATE_LIMITING,
+    description:
+      'Maximum number of requests allowed in the long rate limiting window',
+    type: ConfigVariableType.NUMBER,
+  })
+  @CastToPositiveNumber()
+  API_RATE_LIMITING_LONG_LIMIT = 1000;

  @ConfigVariablesMetadata({
    group: ConfigVariablesGroup.SSL,