Add function execution throttler (#6742)

Add throttler service to limit the number of function execution
This commit is contained in:
Thomas Trompette
2024-08-27 18:56:47 +02:00
committed by GitHub
parent 81fa3f0c41
commit 9f69383aa2
13 changed files with 123 additions and 41 deletions

View File

@ -0,0 +1,12 @@
import { CustomException } from 'src/utils/custom-exception';
export class ThrottlerException extends CustomException {
code: ThrottlerExceptionCode;
constructor(message: string, code: ThrottlerExceptionCode) {
super(message, code);
}
}
export enum ThrottlerExceptionCode {
TOO_MANY_REQUESTS = 'TOO_MANY_REQUESTS',
}

View File

@ -0,0 +1,10 @@
import { Module } from '@nestjs/common';
import { ThrottlerService } from 'src/engine/core-modules/throttler/throttler.service';
@Module({
imports: [],
providers: [ThrottlerService],
exports: [ThrottlerService],
})
export class ThrottlerModule {}

View File

@ -0,0 +1,30 @@
import { Injectable } from '@nestjs/common';
import {
ThrottlerException,
ThrottlerExceptionCode,
} from 'src/engine/core-modules/throttler/throttler.exception';
import { CacheStorageService } from 'src/engine/integrations/cache-storage/cache-storage.service';
import { InjectCacheStorage } from 'src/engine/integrations/cache-storage/decorators/cache-storage.decorator';
import { CacheStorageNamespace } from 'src/engine/integrations/cache-storage/types/cache-storage-namespace.enum';
@Injectable()
export class ThrottlerService {
constructor(
@InjectCacheStorage(CacheStorageNamespace.EngineWorkspace)
private readonly cacheStorage: CacheStorageService,
) {}
async throttle(key: string, limit: number, ttl: number): Promise<void> {
const currentCount = (await this.cacheStorage.get<number>(key)) ?? 0;
if (currentCount >= limit) {
throw new ThrottlerException(
'Too many requests',
ThrottlerExceptionCode.TOO_MANY_REQUESTS,
);
}
await this.cacheStorage.set(key, currentCount + 1, ttl);
}
}