Merge pull request #9404 from weseek/feat/157512-set-a-rate-limit-for-vector-store-rebuild

feat(ai): Set a rate limit for vector store rebuild

apps/app/src/features/rate-limiter/config/index.ts

@@ -56,6 +56,11 @@ export const defaultConfig: IApiRateLimitEndpointMap = {
     method: 'GET',
     maxRequests: MAX_REQUESTS_TIER_3,
   },
+  '/_api/v3/openai/rebuild-vector-store': {
+    method: 'POST',
+    maxRequests: 1,
+    usersPerIpProspection: 1,
+  },
 };
 
 const isDev = process.env.NODE_ENV === 'development';

apps/app/src/features/rate-limiter/middleware/consume-points.integ.ts

@@ -0,0 +1,60 @@
+import { faker } from '@faker-js/faker';
+
+const testRateLimitErrorWhenExceedingMaxRequests = async(method: string, key: string, maxRequests: number): Promise<void> => {
+  // dynamic import is used because rateLimiterMongo needs to be initialized after connecting to DB
+  // Issue: https://github.com/animir/node-rate-limiter-flexible/issues/216
+  const { consumePoints } = await import('./consume-points');
+  let count = 0;
+  try {
+    for (let i = 1; i <= maxRequests + 1; i++) {
+      count += 1;
+      // eslint-disable-next-line no-await-in-loop
+      const res = await consumePoints(method, key, { method, maxRequests });
+      if (count === maxRequests) {
+        // Expect consumedPoints to be equal to maxRequest when maxRequest is reached
+        expect(res?.consumedPoints).toBe(maxRequests);
+        // Expect remainingPoints to be 0 when maxRequest is reached
+        expect(res?.remainingPoints).toBe(0);
+      }
+      if (count > maxRequests) {
+        throw new Error('Exception occurred');
+      }
+    }
+  }
+  catch (err) {
+    // Expect rate limit error to be called
+    expect(err.message).not.toBe('Exception occurred');
+    // Expect rate limit error at maxRequest + 1
+    expect(count).toBe(maxRequests + 1);
+  }
+};
+
+
+describe('consume-points.ts', async() => {
+  it('Should trigger a rate limit error when maxRequest is exceeded (maxRequest: 1)', async() => {
+    // setup
+    const method = 'GET';
+    const key = 'test-key-1';
+    const maxRequests = 1;
+
+    await testRateLimitErrorWhenExceedingMaxRequests(method, key, maxRequests);
+  });
+
+  it('Should trigger a rate limit error when maxRequest is exceeded (maxRequest: 500)', async() => {
+    // setup
+    const method = 'GET';
+    const key = 'test-key-2';
+    const maxRequests = 500;
+
+    await testRateLimitErrorWhenExceedingMaxRequests(method, key, maxRequests);
+  });
+
+  it('Should trigger a rate limit error when maxRequest is exceeded (maxRequest: {random integer between 1 and 1000})', async() => {
+    // setup
+    const method = 'GET';
+    const key = 'test-key-3';
+    const maxRequests = faker.number.int({ min: 1, max: 1000 });
+
+    await testRateLimitErrorWhenExceedingMaxRequests(method, key, maxRequests);
+  });
+});

apps/app/src/features/rate-limiter/middleware/consume-points.ts

@@ -0,0 +1,31 @@
+import { type RateLimiterRes } from 'rate-limiter-flexible';
+
+import { DEFAULT_MAX_REQUESTS, type IApiRateLimitConfig } from '../config';
+
+import { rateLimiterFactory } from './rate-limiter-factory';
+
+export const consumePoints = async(
+    method: string, key: string | null, customizedConfig?: IApiRateLimitConfig, maxRequestsMultiplier?: number,
+): Promise<RateLimiterRes | undefined> => {
+  if (key == null) {
+    return;
+  }
+
+  let maxRequests = DEFAULT_MAX_REQUESTS;
+
+  // use customizedConfig
+  if (customizedConfig != null && (customizedConfig.method.includes(method) || customizedConfig.method === 'ALL')) {
+    maxRequests = customizedConfig.maxRequests;
+  }
+
+  // multiply
+  if (maxRequestsMultiplier != null) {
+    maxRequests *= maxRequestsMultiplier;
+  }
+
+  const rateLimiter = rateLimiterFactory.getOrCreateRateLimiter(key, maxRequests);
+
+  const pointsToConsume = 1;
+  const rateLimiterRes = await rateLimiter.consume(key, pointsToConsume);
+  return rateLimiterRes;
+};

apps/app/src/features/rate-limiter/middleware/factory.ts

@@ -1,16 +1,14 @@
 import type { IUserHasId } from '@growi/core';
 import type { Handler, Request } from 'express';
 import md5 from 'md5';
-import { connection } from 'mongoose';
-import { type IRateLimiterMongoOptions, RateLimiterMongo } from 'rate-limiter-flexible';
+import { type RateLimiterRes } from 'rate-limiter-flexible';
 
 import loggerFactory from '~/utils/logger';
 
-import {
-  DEFAULT_DURATION_SEC, DEFAULT_MAX_REQUESTS, DEFAULT_USERS_PER_IP_PROSPECTION, type IApiRateLimitConfig,
-} from '../config';
+import { DEFAULT_USERS_PER_IP_PROSPECTION, type IApiRateLimitConfig } from '../config';
 import { generateApiRateLimitConfig } from '../utils/config-generator';
 
+import { consumePoints } from './consume-points';
 
 const logger = loggerFactory('growi:middleware:api-rate-limit');
 
@@ -19,15 +17,6 @@ const logger = loggerFactory('growi:middleware:api-rate-limit');
 // API_RATE_LIMIT_010_FOO_METHODS=GET,POST
 // API_RATE_LIMIT_010_FOO_MAX_REQUESTS=10
 
-const POINTS_THRESHOLD = 100;
-
-const opts: IRateLimiterMongoOptions = {
-  storeClient: connection,
-  points: POINTS_THRESHOLD, // set default value
-  duration: DEFAULT_DURATION_SEC, // set default value
-};
-const rateLimiter = new RateLimiterMongo(opts);
-
 // generate ApiRateLimitConfig for api rate limiter
 const apiRateLimitConfig = generateApiRateLimitConfig();
 const configWithoutRegExp = apiRateLimitConfig.withoutRegExp;
@@ -37,31 +26,6 @@ const keysWithRegExp = Object.keys(configWithRegExp).map(key => new RegExp(`^${k
 const valuesWithRegExp = Object.values(configWithRegExp);
 
 
-const _consumePoints = async(
-    method: string, key: string | null, customizedConfig?: IApiRateLimitConfig, maxRequestsMultiplier?: number,
-) => {
-  if (key == null) {
-    return;
-  }
-
-  let maxRequests = DEFAULT_MAX_REQUESTS;
-
-  // use customizedConfig
-  if (customizedConfig != null && (customizedConfig.method.includes(method) || customizedConfig.method === 'ALL')) {
-    maxRequests = customizedConfig.maxRequests;
-  }
-
-  // multiply
-  if (maxRequestsMultiplier != null) {
-    maxRequests *= maxRequestsMultiplier;
-  }
-
-  // because the maximum request is reduced by 1 if it is divisible by
-  // https://github.com/weseek/growi/pull/6225
-  const consumePoints = (POINTS_THRESHOLD + 0.0001) / maxRequests;
-  await rateLimiter.consume(key, consumePoints);
-};
-
 /**
  * consume per user per endpoint
  * @param method
@@ -69,8 +33,10 @@ const _consumePoints = async(
  * @param customizedConfig
  * @returns
  */
-const consumePointsByUser = async(method: string, key: string | null, customizedConfig?: IApiRateLimitConfig) => {
-  return _consumePoints(method, key, customizedConfig);
+const consumePointsByUser = async(
+    method: string, key: string | null, customizedConfig?: IApiRateLimitConfig,
+): Promise<RateLimiterRes | undefined> => {
+  return consumePoints(method, key, customizedConfig);
 };
 
 /**
@@ -80,9 +46,11 @@ const consumePointsByUser = async(method: string, key: string | null, customized
  * @param customizedConfig
  * @returns
  */
-const consumePointsByIp = async(method: string, key: string | null, customizedConfig?: IApiRateLimitConfig) => {
+const consumePointsByIp = async(
+    method: string, key: string | null, customizedConfig?: IApiRateLimitConfig,
+): Promise<RateLimiterRes | undefined> => {
   const maxRequestsMultiplier = customizedConfig?.usersPerIpProspection ?? DEFAULT_USERS_PER_IP_PROSPECTION;
-  return _consumePoints(method, key, customizedConfig, maxRequestsMultiplier);
+  return consumePoints(method, key, customizedConfig, maxRequestsMultiplier);
 };
 
 

apps/app/src/features/rate-limiter/middleware/rate-limiter-factory.ts

@@ -0,0 +1,30 @@
+import { connection } from 'mongoose';
+import { type IRateLimiterMongoOptions, RateLimiterMongo } from 'rate-limiter-flexible';
+
+import { DEFAULT_DURATION_SEC } from '../config';
+
+class RateLimiterFactory {
+
+  private rateLimiters: Map<string, RateLimiterMongo> = new Map();
+
+  getOrCreateRateLimiter(key: string, maxRequests: number): RateLimiterMongo {
+    const cachedRateLimiter = this.rateLimiters.get(key);
+    if (cachedRateLimiter != null) {
+      return cachedRateLimiter;
+    }
+
+    const opts: IRateLimiterMongoOptions = {
+      storeClient: connection,
+      duration: DEFAULT_DURATION_SEC,
+      points: maxRequests,
+    };
+
+    const rateLimiter = new RateLimiterMongo(opts);
+    this.rateLimiters.set(key, rateLimiter);
+
+    return rateLimiter;
+  }
+
+}
+
+export const rateLimiterFactory = new RateLimiterFactory();