Add atomic Redis lock acquisition to DistributedLockService

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<string, string>;
+  let setSpy: jest.Mock;
+  let evalSpy: jest.Mock;
+
+  beforeEach(() => {
+    store = new Map<string, string>();
+
+    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',
+    );
+  });
+});

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<boolean> {
-    const result = await this.redis.set(key, 'locked', 'PX', ttl, 'NX');
-    return result === 'OK';
+  async acquireLock(key: string, ttl = 5000): Promise<string | null> {
+    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<void> {
+    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<void> {
-    await this.redis.del(key);
+  async withLock<T>(key: string, ttl: number, fn: () => Promise<T>): Promise<T> {
+    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);
+    }
   }
 }