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);
}
}

View File

@ -421,6 +421,13 @@ export class EnvironmentVariables {
AUTH_GOOGLE_APIS_CALLBACK_URL: string;
CHROME_EXTENSION_ID: string;
@CastToPositiveNumber()
SERVERLESS_FUNCTION_EXEC_THROTTLE_LIMIT = 10;
// milliseconds
@CastToPositiveNumber()
SERVERLESS_FUNCTION_EXEC_THROTTLE_TTL = 1000;
}
export const validate = (

View File

@ -2,24 +2,24 @@ import { Module } from '@nestjs/common';
import { HttpAdapterHost } from '@nestjs/core';
import { EventEmitterModule } from '@nestjs/event-emitter';
import { ExceptionHandlerModule } from 'src/engine/integrations/exception-handler/exception-handler.module';
import { exceptionHandlerModuleFactory } from 'src/engine/integrations/exception-handler/exception-handler.module-factory';
import { fileStorageModuleFactory } from 'src/engine/integrations/file-storage/file-storage.module-factory';
import { loggerModuleFactory } from 'src/engine/integrations/logger/logger.module-factory';
import { messageQueueModuleFactory } from 'src/engine/integrations/message-queue/message-queue.module-factory';
import { EmailModule } from 'src/engine/integrations/email/email.module';
import { emailModuleFactory } from 'src/engine/integrations/email/email.module-factory';
import { CacheStorageModule } from 'src/engine/integrations/cache-storage/cache-storage.module';
import { CaptchaModule } from 'src/engine/integrations/captcha/captcha.module';
import { captchaModuleFactory } from 'src/engine/integrations/captcha/captcha.module-factory';
import { EmailModule } from 'src/engine/integrations/email/email.module';
import { emailModuleFactory } from 'src/engine/integrations/email/email.module-factory';
import { ExceptionHandlerModule } from 'src/engine/integrations/exception-handler/exception-handler.module';
import { exceptionHandlerModuleFactory } from 'src/engine/integrations/exception-handler/exception-handler.module-factory';
import { fileStorageModuleFactory } from 'src/engine/integrations/file-storage/file-storage.module-factory';
import { FileStorageService } from 'src/engine/integrations/file-storage/file-storage.service';
import { LLMChatModelModule } from 'src/engine/integrations/llm-chat-model/llm-chat-model.module';
import { llmChatModelModuleFactory } from 'src/engine/integrations/llm-chat-model/llm-chat-model.module-factory';
import { LLMTracingModule } from 'src/engine/integrations/llm-tracing/llm-tracing.module';
import { llmTracingModuleFactory } from 'src/engine/integrations/llm-tracing/llm-tracing.module-factory';
import { ServerlessModule } from 'src/engine/integrations/serverless/serverless.module';
import { serverlessModuleFactory } from 'src/engine/integrations/serverless/serverless-module.factory';
import { FileStorageService } from 'src/engine/integrations/file-storage/file-storage.service';
import { loggerModuleFactory } from 'src/engine/integrations/logger/logger.module-factory';
import { messageQueueModuleFactory } from 'src/engine/integrations/message-queue/message-queue.module-factory';
import { BuildDirectoryManagerService } from 'src/engine/integrations/serverless/drivers/services/build-directory-manager.service';
import { serverlessModuleFactory } from 'src/engine/integrations/serverless/serverless-module.factory';
import { ServerlessModule } from 'src/engine/integrations/serverless/serverless.module';
import { EnvironmentModule } from './environment/environment.module';
import { EnvironmentService } from './environment/environment.service';

View File

@ -2,16 +2,16 @@ import fs from 'fs';
import {
CreateFunctionCommand,
DeleteFunctionCommand,
GetFunctionCommand,
InvokeCommand,
Lambda,
LambdaClientConfig,
InvokeCommand,
GetFunctionCommand,
UpdateFunctionCodeCommand,
DeleteFunctionCommand,
ResourceNotFoundException,
waitUntilFunctionUpdatedV2,
PublishVersionCommandInput,
PublishVersionCommand,
PublishVersionCommandInput,
ResourceNotFoundException,
UpdateFunctionCodeCommand,
waitUntilFunctionUpdatedV2,
} from '@aws-sdk/client-lambda';
import { CreateFunctionCommandInput } from '@aws-sdk/client-lambda/dist-types/commands/CreateFunctionCommand';
import { UpdateFunctionCodeCommandInput } from '@aws-sdk/client-lambda/dist-types/commands/UpdateFunctionCodeCommand';
@ -21,12 +21,12 @@ import {
ServerlessExecuteResult,
} from 'src/engine/integrations/serverless/drivers/interfaces/serverless-driver.interface';
import { createZipFile } from 'src/engine/integrations/serverless/drivers/utils/create-zip-file';
import { ServerlessFunctionEntity } from 'src/engine/metadata-modules/serverless-function/serverless-function.entity';
import { FileStorageService } from 'src/engine/integrations/file-storage/file-storage.service';
import { BaseServerlessDriver } from 'src/engine/integrations/serverless/drivers/base-serverless.driver';
import { BuildDirectoryManagerService } from 'src/engine/integrations/serverless/drivers/services/build-directory-manager.service';
import { createZipFile } from 'src/engine/integrations/serverless/drivers/utils/create-zip-file';
import { ServerlessFunctionExecutionStatus } from 'src/engine/metadata-modules/serverless-function/dtos/serverless-function-execution-result.dto';
import { ServerlessFunctionEntity } from 'src/engine/metadata-modules/serverless-function/serverless-function.entity';
import {
ServerlessFunctionException,
ServerlessFunctionExceptionCode,

View File

@ -1,29 +1,28 @@
/* eslint-disable no-console */
import { join } from 'path';
import { tmpdir } from 'os';
import { promises as fs } from 'fs';
import { fork } from 'child_process';
import { promises as fs } from 'fs';
import { tmpdir } from 'os';
import { join } from 'path';
import { v4 } from 'uuid';
import { FileStorageExceptionCode } from 'src/engine/integrations/file-storage/interfaces/file-storage-exception';
import {
ServerlessDriver,
ServerlessExecuteError,
ServerlessExecuteResult,
} from 'src/engine/integrations/serverless/drivers/interfaces/serverless-driver.interface';
import { FileStorageExceptionCode } from 'src/engine/integrations/file-storage/interfaces/file-storage-exception';
import { FileStorageService } from 'src/engine/integrations/file-storage/file-storage.service';
import { readFileContent } from 'src/engine/integrations/file-storage/utils/read-file-content';
import { ServerlessFunctionEntity } from 'src/engine/metadata-modules/serverless-function/serverless-function.entity';
import { BUILD_FILE_NAME } from 'src/engine/integrations/serverless/drivers/constants/build-file-name';
import { BaseServerlessDriver } from 'src/engine/integrations/serverless/drivers/base-serverless.driver';
import { BUILD_FILE_NAME } from 'src/engine/integrations/serverless/drivers/constants/build-file-name';
import { getServerlessFolder } from 'src/engine/integrations/serverless/utils/serverless-get-folder.utils';
import { ServerlessFunctionExecutionStatus } from 'src/engine/metadata-modules/serverless-function/dtos/serverless-function-execution-result.dto';
import { ServerlessFunctionEntity } from 'src/engine/metadata-modules/serverless-function/serverless-function.entity';
import {
ServerlessFunctionException,
ServerlessFunctionExceptionCode,
} from 'src/engine/metadata-modules/serverless-function/serverless-function.exception';
import { getServerlessFolder } from 'src/engine/integrations/serverless/utils/serverless-get-folder.utils';
export interface LocalDriverOptions {
fileStorageService: FileStorageService;

View File

@ -1,12 +1,12 @@
import { fromNodeProviderChain } from '@aws-sdk/credential-providers';
import {
ServerlessModuleOptions,
ServerlessDriverType,
} from 'src/engine/integrations/serverless/serverless.interface';
import { EnvironmentService } from 'src/engine/integrations/environment/environment.service';
import { FileStorageService } from 'src/engine/integrations/file-storage/file-storage.service';
import { BuildDirectoryManagerService } from 'src/engine/integrations/serverless/drivers/services/build-directory-manager.service';
import {
ServerlessDriverType,
ServerlessModuleOptions,
} from 'src/engine/integrations/serverless/serverless.interface';
export const serverlessModuleFactory = async (
environmentService: EnvironmentService,

View File

@ -1,14 +1,14 @@
import { DynamicModule, Global } from '@nestjs/common';
import { LambdaDriver } from 'src/engine/integrations/serverless/drivers/lambda.driver';
import { LocalDriver } from 'src/engine/integrations/serverless/drivers/local.driver';
import { BuildDirectoryManagerService } from 'src/engine/integrations/serverless/drivers/services/build-directory-manager.service';
import { SERVERLESS_DRIVER } from 'src/engine/integrations/serverless/serverless.constants';
import {
ServerlessDriverType,
ServerlessModuleAsyncOptions,
} from 'src/engine/integrations/serverless/serverless.interface';
import { ServerlessService } from 'src/engine/integrations/serverless/serverless.service';
import { SERVERLESS_DRIVER } from 'src/engine/integrations/serverless/serverless.constants';
import { LocalDriver } from 'src/engine/integrations/serverless/drivers/local.driver';
import { LambdaDriver } from 'src/engine/integrations/serverless/drivers/lambda.driver';
import { BuildDirectoryManagerService } from 'src/engine/integrations/serverless/drivers/services/build-directory-manager.service';
@Global()
export class ServerlessModule {

View File

@ -13,4 +13,5 @@ export enum ServerlessFunctionExceptionCode {
FEATURE_FLAG_INVALID = 'FEATURE_FLAG_INVALID',
SERVERLESS_FUNCTION_ALREADY_EXIST = 'SERVERLESS_FUNCTION_ALREADY_EXIST',
SERVERLESS_FUNCTION_NOT_READY = 'SERVERLESS_FUNCTION_NOT_READY',
SERVERLESS_FUNCTION_EXECUTION_LIMIT_REACHED = 'SERVERLESS_FUNCTION_EXECUTION_LIMIT_REACHED',
}

View File

@ -11,6 +11,7 @@ import { NestjsQueryTypeOrmModule } from '@ptc-org/nestjs-query-typeorm';
import { FeatureFlagEntity } from 'src/engine/core-modules/feature-flag/feature-flag.entity';
import { FileUploadModule } from 'src/engine/core-modules/file/file-upload/file-upload.module';
import { FileModule } from 'src/engine/core-modules/file/file.module';
import { ThrottlerModule } from 'src/engine/core-modules/throttler/throttler.module';
import { JwtAuthGuard } from 'src/engine/guards/jwt.auth.guard';
import { ServerlessModule } from 'src/engine/integrations/serverless/serverless.module';
import { ServerlessFunctionDTO } from 'src/engine/metadata-modules/serverless-function/dtos/serverless-function.dto';
@ -29,6 +30,7 @@ import { ServerlessFunctionService } from 'src/engine/metadata-modules/serverles
),
TypeOrmModule.forFeature([FeatureFlagEntity], 'core'),
FileModule,
ThrottlerModule,
],
services: [ServerlessFunctionService],
resolvers: [

View File

@ -5,13 +5,16 @@ import { TypeOrmQueryService } from '@ptc-org/nestjs-query-typeorm';
import { FileUpload } from 'graphql-upload';
import { Repository } from 'typeorm';
import { ServerlessExecuteResult } from 'src/engine/integrations/serverless/drivers/interfaces/serverless-driver.interface';
import { FileStorageExceptionCode } from 'src/engine/integrations/file-storage/interfaces/file-storage-exception';
import { ServerlessExecuteResult } from 'src/engine/integrations/serverless/drivers/interfaces/serverless-driver.interface';
import { ThrottlerService } from 'src/engine/core-modules/throttler/throttler.service';
import { EnvironmentService } from 'src/engine/integrations/environment/environment.service';
import { FileStorageService } from 'src/engine/integrations/file-storage/file-storage.service';
import { readFileContent } from 'src/engine/integrations/file-storage/utils/read-file-content';
import { SOURCE_FILE_NAME } from 'src/engine/integrations/serverless/drivers/constants/source-file-name';
import { ServerlessService } from 'src/engine/integrations/serverless/serverless.service';
import { getServerlessFolder } from 'src/engine/integrations/serverless/utils/serverless-get-folder.utils';
import { CreateServerlessFunctionFromFileInput } from 'src/engine/metadata-modules/serverless-function/dtos/create-serverless-function-from-file.input';
import { UpdateServerlessFunctionInput } from 'src/engine/metadata-modules/serverless-function/dtos/update-serverless-function.input';
import {
@ -24,7 +27,6 @@ import {
} from 'src/engine/metadata-modules/serverless-function/serverless-function.exception';
import { serverlessFunctionCreateHash } from 'src/engine/metadata-modules/serverless-function/utils/serverless-function-create-hash.utils';
import { isDefined } from 'src/utils/is-defined';
import { getServerlessFolder } from 'src/engine/integrations/serverless/utils/serverless-get-folder.utils';
@Injectable()
export class ServerlessFunctionService extends TypeOrmQueryService<ServerlessFunctionEntity> {
@ -33,6 +35,8 @@ export class ServerlessFunctionService extends TypeOrmQueryService<ServerlessFun
private readonly serverlessService: ServerlessService,
@InjectRepository(ServerlessFunctionEntity, 'metadata')
private readonly serverlessFunctionRepository: Repository<ServerlessFunctionEntity>,
private readonly throttlerService: ThrottlerService,
private readonly environmentService: EnvironmentService,
) {
super(serverlessFunctionRepository);
}
@ -86,6 +90,8 @@ export class ServerlessFunctionService extends TypeOrmQueryService<ServerlessFun
payload: object | undefined = undefined,
version = 'latest',
): Promise<ServerlessExecuteResult> {
await this.throttleExecution(workspaceId);
const functionToExecute = await this.serverlessFunctionRepository.findOne({
where: {
id,
@ -268,4 +274,19 @@ export class ServerlessFunctionService extends TypeOrmQueryService<ServerlessFun
return await this.findById(createdServerlessFunction.id);
}
private async throttleExecution(workspaceId: string) {
try {
await this.throttlerService.throttle(
`${workspaceId}-serverless-function-execution`,
this.environmentService.get('SERVERLESS_FUNCTION_EXEC_THROTTLE_LIMIT'),
this.environmentService.get('SERVERLESS_FUNCTION_EXEC_THROTTLE_TTL'),
);
} catch (error) {
throw new ServerlessFunctionException(
'Serverless function execution rate limit exceeded',
ServerlessFunctionExceptionCode.SERVERLESS_FUNCTION_EXECUTION_LIMIT_REACHED,
);
}
}
}

View File

@ -1,13 +1,13 @@
import {
ServerlessFunctionException,
ServerlessFunctionExceptionCode,
} from 'src/engine/metadata-modules/serverless-function/serverless-function.exception';
import {
ConflictError,
ForbiddenError,
InternalServerError,
NotFoundError,
} from 'src/engine/core-modules/graphql/utils/graphql-errors.util';
import {
ServerlessFunctionException,
ServerlessFunctionExceptionCode,
} from 'src/engine/metadata-modules/serverless-function/serverless-function.exception';
export const serverlessFunctionGraphQLApiExceptionHandler = (error: any) => {
if (error instanceof ServerlessFunctionException) {