feat(server): make a singleton global mutex service (#7900)

This commit is contained in:
forehalo
2024-08-21 05:30:19 +00:00
parent 6b0c398ae5
commit 682a01e441
6 changed files with 50 additions and 42 deletions
@@ -5,7 +5,7 @@ import {
ActionForbidden, ActionForbidden,
EventEmitter, EventEmitter,
InternalServerError, InternalServerError,
MutexService, Mutex,
PasswordRequired, PasswordRequired,
} from '../../fundamentals'; } from '../../fundamentals';
import { AuthService, Public } from '../auth'; import { AuthService, Public } from '../auth';
@@ -23,7 +23,7 @@ export class CustomSetupController {
private readonly user: UserService, private readonly user: UserService,
private readonly auth: AuthService, private readonly auth: AuthService,
private readonly event: EventEmitter, private readonly event: EventEmitter,
private readonly mutex: MutexService, private readonly mutex: Mutex,
private readonly server: ServerService private readonly server: ServerService
) {} ) {}
@@ -20,7 +20,7 @@ import {
InternalServerError, InternalServerError,
MailService, MailService,
MemberQuotaExceeded, MemberQuotaExceeded,
MutexService, RequestMutex,
Throttle, Throttle,
TooManyRequest, TooManyRequest,
UserNotFound, UserNotFound,
@@ -57,7 +57,7 @@ export class WorkspaceResolver {
private readonly users: UserService, private readonly users: UserService,
private readonly event: EventEmitter, private readonly event: EventEmitter,
private readonly blobStorage: WorkspaceBlobStorage, private readonly blobStorage: WorkspaceBlobStorage,
private readonly mutex: MutexService private readonly mutex: RequestMutex
) {} ) {}
@ResolveField(() => Permission, { @ResolveField(() => Permission, {
@@ -19,7 +19,7 @@ export type { GraphqlContext } from './graphql';
export { CryptoHelper, URLHelper } from './helpers'; export { CryptoHelper, URLHelper } from './helpers';
export { MailService } from './mailer'; export { MailService } from './mailer';
export { CallCounter, CallTimer, metrics } from './metrics'; export { CallCounter, CallTimer, metrics } from './metrics';
export { type ILocker, Lock, Locker, MutexService } from './mutex'; export { type ILocker, Lock, Locker, Mutex, RequestMutex } from './mutex';
export { export {
GatewayErrorWrapper, GatewayErrorWrapper,
getOptionalModuleMetadata, getOptionalModuleMetadata,
@@ -1,14 +1,14 @@
import { Global, Module } from '@nestjs/common'; import { Global, Module } from '@nestjs/common';
import { Locker } from './local-lock'; import { Locker } from './local-lock';
import { MutexService } from './mutex'; import { Mutex, RequestMutex } from './mutex';
@Global() @Global()
@Module({ @Module({
providers: [MutexService, Locker], providers: [Mutex, RequestMutex, Locker],
exports: [MutexService], exports: [Mutex, RequestMutex],
}) })
export class MutexModule {} export class MutexModule {}
export { Locker, MutexService }; export { Locker, Mutex, RequestMutex };
export { type Locker as ILocker, Lock } from './lock'; export { type Locker as ILocker, Lock } from './lock';
@@ -11,36 +11,11 @@ import { Locker } from './local-lock';
export const MUTEX_RETRY = 5; export const MUTEX_RETRY = 5;
export const MUTEX_WAIT = 100; export const MUTEX_WAIT = 100;
@Injectable({ scope: Scope.REQUEST }) @Injectable()
export class MutexService { export class Mutex {
protected logger = new Logger(MutexService.name); protected logger = new Logger(Mutex.name);
private readonly locker: Locker;
constructor( constructor(protected readonly locker: Locker) {}
@Inject(REQUEST) private readonly request: Request | GraphqlContext,
private readonly ref: ModuleRef
) {
// nestjs will always find and injecting the locker from local module
// so the RedisLocker implemented by the plugin mechanism will not be able to overwrite the internal locker
// we need to use find and get the locker from the `ModuleRef` manually
//
// NOTE: when a `constructor` execute in normal service, the Locker module we expect may not have been initialized
// but in the Service with `Scope.REQUEST`, we will create a separate Service instance for each request
// at this time, all modules have been initialized, so we able to get the correct Locker instance in `constructor`
this.locker = this.ref.get(Locker, { strict: false });
}
protected getId() {
const req = 'req' in this.request ? this.request.req : this.request;
let id = req.headers['x-transaction-id'] as string;
if (!id) {
id = randomUUID();
req.headers['x-transaction-id'] = id;
}
return id;
}
/** /**
* lock an resource and return a lock guard, which will release the lock when disposed * lock an resource and return a lock guard, which will release the lock when disposed
@@ -63,10 +38,10 @@ export class MutexService {
* @param key resource key * @param key resource key
* @returns LockGuard * @returns LockGuard
*/ */
async lock(key: string) { async lock(key: string, owner: string = 'global') {
try { try {
return await retryable( return await retryable(
() => this.locker.lock(this.getId(), key), () => this.locker.lock(owner, key),
MUTEX_RETRY, MUTEX_RETRY,
MUTEX_WAIT MUTEX_WAIT
); );
@@ -79,3 +54,36 @@ export class MutexService {
} }
} }
} }
@Injectable({ scope: Scope.REQUEST })
export class RequestMutex extends Mutex {
constructor(
@Inject(REQUEST) private readonly request: Request | GraphqlContext,
ref: ModuleRef
) {
// nestjs will always find and injecting the locker from local module
// so the RedisLocker implemented by the plugin mechanism will not be able to overwrite the internal locker
// we need to use find and get the locker from the `ModuleRef` manually
//
// NOTE: when a `constructor` execute in normal service, the Locker module we expect may not have been initialized
// but in the Service with `Scope.REQUEST`, we will create a separate Service instance for each request
// at this time, all modules have been initialized, so we able to get the correct Locker instance in `constructor`
super(ref.get(Locker));
}
protected getId() {
const req = 'req' in this.request ? this.request.req : this.request;
let id = req.headers['x-transaction-id'] as string;
if (!id) {
id = randomUUID();
req.headers['x-transaction-id'] = id;
}
return id;
}
override lock(key: string) {
return super.lock(key, this.getId());
}
}
@@ -26,7 +26,7 @@ import { UserType } from '../../core/user';
import { import {
CopilotFailedToCreateMessage, CopilotFailedToCreateMessage,
FileUpload, FileUpload,
MutexService, RequestMutex,
Throttle, Throttle,
TooManyRequest, TooManyRequest,
} from '../../fundamentals'; } from '../../fundamentals';
@@ -265,7 +265,7 @@ export class CopilotType {
export class CopilotResolver { export class CopilotResolver {
constructor( constructor(
private readonly permissions: PermissionService, private readonly permissions: PermissionService,
private readonly mutex: MutexService, private readonly mutex: RequestMutex,
private readonly chatSession: ChatSessionService, private readonly chatSession: ChatSessionService,
private readonly storage: CopilotStorage private readonly storage: CopilotStorage
) {} ) {}