| 1 | import type { Injectable } from '@nestjs/common/interfaces' |
| 2 | import type { InstanceWrapper } from '@nestjs/core/injector/instance-wrapper' |
| 3 | import { Injectable as InjectableDec, Logger, OnModuleInit } from '@nestjs/common' |
| 4 | import { MetadataScanner, ModulesContainer } from '@nestjs/core' |
| 5 | import { RedlockConfig } from './redlock.config' |
| 6 | import { |
| 7 | REDLOCK_METADATA, |
| 8 | RedlockOptions, |
| 9 | } from './redlock.decorator' |
| 10 | import { RedlockService } from './redlock.service' |
| 11 | |
| 12 | @InjectableDec() |
| 13 | export class RedlockInjector implements OnModuleInit { |
| 14 | private readonly logger = new Logger(RedlockInjector.name) |
| 15 | private readonly metadataScanner: MetadataScanner = new MetadataScanner() |
| 16 | |
| 17 | constructor( |
| 18 | private readonly modulesContainer: ModulesContainer, |
| 19 | private readonly redlockService: RedlockService, |
| 20 | private readonly config: RedlockConfig, |
| 21 | ) {} |
| 22 | |
| 23 | async onModuleInit() { |
| 24 | for (const provider of this.getProviders()) { |
| 25 | this.injectToProvider(provider) |
| 26 | } |
| 27 | } |
| 28 | |
| 29 | private* getProviders(): Generator<InstanceWrapper<Injectable>> { |
| 30 | for (const module of this.modulesContainer.values()) { |
| 31 | for (const provider of module.providers.values()) { |
| 32 | if (provider && provider.metatype?.prototype) { |
| 33 | yield provider as InstanceWrapper<Injectable> |
| 34 | } |
| 35 | } |
| 36 | } |
| 37 | } |
| 38 | |
| 39 | private injectToProvider(wrapper: InstanceWrapper<Injectable>): void { |
| 40 | const { metatype } = wrapper |
| 41 | if (!metatype) |
| 42 | return |
| 43 | |
| 44 | const prototype = metatype.prototype |
| 45 | const methodNames = this.metadataScanner.getAllMethodNames(prototype) |
| 46 | |
| 47 | for (const methodName of methodNames) { |
| 48 | const method = prototype[methodName] |
| 49 | if (this.isDecorated(method)) { |
| 50 | const options = this.getDecoratorOptions(method) |
| 51 | const wrappedMethod = this.wrapMethod(method, methodName, prototype.constructor.name, options) |
| 52 | this.reDecorate(method, wrappedMethod) |
| 53 | prototype[methodName] = wrappedMethod |
| 54 | this.logger.log(`Injected distributed lock to ${prototype.constructor.name}.${methodName}`) |
| 55 | } |
| 56 | } |
| 57 | } |
| 58 | |
| 59 | private isDecorated(target: object): boolean { |
| 60 | return Reflect.hasMetadata(REDLOCK_METADATA, target) |
| 61 | } |
| 62 | |
| 63 | private getDecoratorOptions(target: object): RedlockOptions { |
| 64 | return Reflect.getMetadata(REDLOCK_METADATA, target) |
| 65 | } |
| 66 | |
| 67 | private reDecorate(source: object, destination: object): void { |
| 68 | const keys = Reflect.getMetadataKeys(source) |
| 69 | for (const key of keys) { |
| 70 | const meta = Reflect.getMetadata(key, source) |
| 71 | Reflect.defineMetadata(key, meta, destination) |
| 72 | } |
| 73 | } |
| 74 | |
| 75 | private wrapMethod( |
| 76 | originalMethod: (...args: unknown[]) => unknown, |
| 77 | methodName: string, |
| 78 | className: string, |
| 79 | options: RedlockOptions, |
| 80 | ): (...args: unknown[]) => unknown { |
| 81 | const lockService = this.redlockService |
| 82 | const logger = this.logger |
| 83 | |
| 84 | return new Proxy(originalMethod, { |
| 85 | apply: async (target, thisArg, args: unknown[]) => { |
| 86 | const baseKey = typeof options.key === 'function' ? options.key(...args) : options.key |
| 87 | const lockKey = `lock:${baseKey}` |
| 88 | const lockValue = `${Date.now()}-${Math.random()}` |
| 89 | const ttl = options.ttl ?? this.config.ttl |
| 90 | const retryDelay = options.retryDelay ?? this.config.retryDelay |
| 91 | const retryCount = options.retryCount ?? this.config.retryCount |
| 92 | const throwOnFailure = options.throwOnFailure ?? true |
| 93 | |
| 94 | let acquired = false |
| 95 | let attempts = 0 |
| 96 | |
| 97 | while (!acquired && attempts < retryCount) { |
| 98 | acquired = await lockService.acquireLock(lockKey, lockValue, ttl) |
| 99 | |
| 100 | if (!acquired) { |
| 101 | attempts++ |
| 102 | if (attempts < retryCount) { |
| 103 | logger.debug(`Failed to acquire lock ${lockKey} for ${className}.${methodName}, attempt ${attempts}/${retryCount}. Retrying in ${retryDelay}ms`) |
| 104 | await new Promise(resolve => setTimeout(resolve, retryDelay)) |
| 105 | } |
| 106 | } |
| 107 | } |
| 108 | |
| 109 | if (!acquired) { |
| 110 | const message = `Could not acquire lock ${lockKey} for ${className}.${methodName} after ${retryCount} attempts` |
| 111 | if (throwOnFailure) { |
| 112 | logger.error(message) |
| 113 | throw new Error(message) |
| 114 | } |
| 115 | else { |
| 116 | logger.debug(`${message}, skipping execution`) |
| 117 | return |
| 118 | } |
| 119 | } |
| 120 | |
| 121 | logger.debug(`Acquired lock ${lockKey} for ${className}.${methodName}, executing method`) |
| 122 | |
| 123 | const result = (async () => Reflect.apply(target, thisArg, args))() |
| 124 | return result.finally(async () => { |
| 125 | await lockService.releaseLock(lockKey, lockValue) |
| 126 | logger.debug(`Released lock ${lockKey} for ${className}.${methodName}`) |
| 127 | }) |
| 128 | }, |
| 129 | }) |
| 130 | } |
| 131 | } |
| 132 |