From 95104291198709a86ffbc1f85b002284f6d18be3 Mon Sep 17 00:00:00 2001 From: dreamgeneX Date: Mon, 29 Jun 2026 14:19:27 -0700 Subject: [PATCH] Add atomic Redis lock acquisition to DistributedLockService MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit acquireLock already used a single SET key value NX PX command, but it returned a boolean and a hardcoded value, so any caller could release a lock it never held once the TTL expired. acquireLock now returns a unique token (or null if already held), releaseLock only deletes the key if the token still matches (atomic GET+DEL via Lua), and a new withLock(key, ttl, fn) helper acquires, runs fn, and always releases — even if fn throws. Closes #888 --- .../locks/distributed-lock.service.spec.ts | 118 ++++++++++++++++++ .../locks/distributed-lock.service.ts | 48 +++++-- 2 files changed, 157 insertions(+), 9 deletions(-) create mode 100644 src/orchestration/locks/distributed-lock.service.spec.ts diff --git a/src/orchestration/locks/distributed-lock.service.spec.ts b/src/orchestration/locks/distributed-lock.service.spec.ts new file mode 100644 index 00000000..a3dbde62 --- /dev/null +++ b/src/orchestration/locks/distributed-lock.service.spec.ts @@ -0,0 +1,118 @@ +jest.mock('ioredis', () => { + return jest.fn(); +}); + +import Redis from 'ioredis'; +import { DistributedLockService } from './distributed-lock.service'; + +describe('DistributedLockService', () => { + let store: Map; + let setSpy: jest.Mock; + let evalSpy: jest.Mock; + + beforeEach(() => { + store = new Map(); + + setSpy = jest.fn( + async (key: string, value: string, mode: string, ttl: number, flag: string) => { + if (flag === 'NX' && store.has(key)) { + return null; + } + store.set(key, value); + return 'OK'; + }, + ); + + evalSpy = jest.fn(async (_script: string, _numKeys: number, key: string, token: string) => { + if (store.get(key) === token) { + store.delete(key); + return 1; + } + return 0; + }); + + (Redis as unknown as jest.Mock).mockImplementation(() => ({ + on: jest.fn(), + set: setSpy, + del: jest.fn(async (key: string) => { + store.delete(key); + return 1; + }), + eval: evalSpy, + })); + }); + + it('acquires the lock using a single SET key value NX PX command', async () => { + const service = new DistributedLockService(); + + const token = await service.acquireLock('lock:test', 5000); + + expect(token).not.toBeNull(); + expect(setSpy).toHaveBeenCalledTimes(1); + expect(setSpy).toHaveBeenCalledWith('lock:test', token, 'PX', 5000, 'NX'); + }); + + it('returns null when the lock is already held', async () => { + const service = new DistributedLockService(); + + const first = await service.acquireLock('lock:test', 5000); + const second = await service.acquireLock('lock:test', 5000); + + expect(first).not.toBeNull(); + expect(second).toBeNull(); + }); + + it('allows exactly one caller to acquire the lock under 100 concurrent attempts', async () => { + const service = new DistributedLockService(); + + const results = await Promise.all( + Array.from({ length: 100 }, () => service.acquireLock('lock:contended', 5000)), + ); + + const winners = results.filter((token) => token !== null); + expect(winners).toHaveLength(1); + }); + + it('releaseLock only removes the lock if the token matches', async () => { + const service = new DistributedLockService(); + const token = await service.acquireLock('lock:test', 5000); + + await service.releaseLock('lock:test', 'wrong-token'); + expect(await service.acquireLock('lock:test', 5000)).toBeNull(); + + await service.releaseLock('lock:test', token as string); + expect(await service.acquireLock('lock:test', 5000)).not.toBeNull(); + }); + + it('withLock releases the lock even if fn throws', async () => { + const service = new DistributedLockService(); + + await expect( + service.withLock('lock:test', 5000, async () => { + throw new Error('boom'); + }), + ).rejects.toThrow('boom'); + + // Lock must be free again — a fresh acquire should succeed immediately. + const token = await service.acquireLock('lock:test', 5000); + expect(token).not.toBeNull(); + }); + + it('withLock returns the value produced by fn and releases the lock', async () => { + const service = new DistributedLockService(); + + const result = await service.withLock('lock:test', 5000, async () => 'done'); + + expect(result).toBe('done'); + expect(await service.acquireLock('lock:test', 5000)).not.toBeNull(); + }); + + it('withLock throws if the lock cannot be acquired', async () => { + const service = new DistributedLockService(); + await service.acquireLock('lock:test', 5000); + + await expect(service.withLock('lock:test', 5000, async () => 'unreachable')).rejects.toThrow( + 'Could not acquire lock', + ); + }); +}); diff --git a/src/orchestration/locks/distributed-lock.service.ts b/src/orchestration/locks/distributed-lock.service.ts index 273a2903..81d18f1e 100644 --- a/src/orchestration/locks/distributed-lock.service.ts +++ b/src/orchestration/locks/distributed-lock.service.ts @@ -1,5 +1,13 @@ import { Injectable } from '@nestjs/common'; import Redis from 'ioredis'; +import { randomUUID } from 'crypto'; + +const RELEASE_SCRIPT = ` + if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('DEL', KEYS[1]) + end + return 0 +`; /** * Provides distributed Lock operations. @@ -15,21 +23,43 @@ export class DistributedLockService { } /** - * Executes acquire Lock. + * Atomically acquires a lock via a single SET key value NX PX command, + * so two concurrent callers can never both believe they hold it. * @param key The key. - * @param ttl The ttl. - * @returns Whether the operation succeeded. + * @param ttl The ttl in milliseconds. + * @returns A unique token if acquired, or null if the lock is already held. */ - async acquireLock(key: string, ttl = 5000): Promise { - const result = await this.redis.set(key, 'locked', 'PX', ttl, 'NX'); - return result === 'OK'; + async acquireLock(key: string, ttl = 5000): Promise { + const token = randomUUID(); + const result = await this.redis.set(key, token, 'PX', ttl, 'NX'); + return result === 'OK' ? token : null; } /** - * Executes release Lock. + * Releases the lock, but only if it is still held by this token. This + * prevents releasing a lock that has since expired and been acquired by + * another caller. * @param key The key. + * @param token The token returned by acquireLock. + */ + async releaseLock(key: string, token: string): Promise { + await this.redis.eval(RELEASE_SCRIPT, 1, key, token); + } + + /** + * Acquires the lock, runs fn, and always releases the lock afterwards + * (even if fn throws). Throws if the lock cannot be acquired. */ - async releaseLock(key: string): Promise { - await this.redis.del(key); + async withLock(key: string, ttl: number, fn: () => Promise): Promise { + const token = await this.acquireLock(key, ttl); + if (!token) { + throw new Error(`Could not acquire lock: ${key}`); + } + + try { + return await fn(); + } finally { + await this.releaseLock(key, token); + } } }