Add function execution throttler (#6742)
Add throttler service to limit the number of function execution
This commit is contained in:
@ -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',
|
||||
}
|
||||
@ -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 {}
|
||||
@ -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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user