fix: grant no locks once the application starts shutting down

lib/errors/lock-not-acquired.error.ts

@@ -1,17 +1,23 @@
 import { LocksError } from './locks.error.js';
 
-/** `withLock()` couldn't take the lock within `wait` (another holder kept it). */
+/**
+ * `withLock()` couldn't take the lock within `wait`: another holder kept it (`'taken'`), or the
+ * application is shutting down and grants no new locks (`'shutdown'`).
+ */
 export class LockNotAcquiredError extends LocksError {
   override name = 'LockNotAcquiredError';
   constructor(
     readonly key: string,
     /** How long it waited, in ms (0: it tried once). */
     readonly waitMs: number,
+    readonly reason: 'taken' | 'shutdown' = 'taken',
   ) {
     super(
-      waitMs > 0
-        ? `Could not acquire the lock "${key}" within ${waitMs}ms: another holder has it`
-        : `Could not acquire the lock "${key}": another holder has it`,
+      reason === 'shutdown'
+        ? `Could not acquire the lock "${key}": the application is shutting down`
+        : waitMs > 0
+          ? `Could not acquire the lock "${key}" within ${waitMs}ms: another holder has it`
+          : `Could not acquire the lock "${key}": another holder has it`,
     );
   }
 }

lib/interfaces/lock-options.interface.ts

@@ -12,13 +12,16 @@ export interface LockOptions {
    * How long to wait for a lock another holder has, trying again with backoff (and at once
    * when a holder in this process releases it). Default `0`: try once. A store call already
    * sent is answered past `wait`: `wait` bounds the waiting for another holder, and the
-   * module's `storeTimeout` bounds the store.
+   * module's `storeTimeout` bounds the store. The waiting stops when the application shuts
+   * down: `acquire()` then resolves `null`, without calling the store again.
    */
   wait?: Duration;
   /**
    * Stops waiting, also for a store call in flight: `acquire()` rejects with the signal's
    * reason (and a lock the store grants after that goes back to it), and `withLock()` doesn't
-   * run its callback.
+   * run its callback. Not needed for the application's own shutdown: from
+   * `beforeApplicationShutdown()` on, `acquire()` resolves `null` without a store call, and a
+   * grant that lands after that goes back too.
    */
   signal?: AbortSignal;
 }

lib/locks.service.ts

@@ -24,6 +24,8 @@ export interface LocksInternals {
   /**
    * One attempt at `key`, with `onLost` called if the lock is lost later. It rejects with
    * `signal`'s reason when it aborts first (shutdown), and a lock granted after that goes back.
+   * Once `beforeApplicationShutdown()` has begun, it resolves `null` without calling the store,
+   * and a grant that lands after that goes back too.
    */
   tryAcquire(
     key: string,
@@ -58,6 +60,8 @@ export class Locks implements BeforeApplicationShutdown {
   private readonly storeTimeout: number;
   private readonly held = new Set<Lock>();
   private readonly waiters = new Map<string, Set<() => void>>();
+  /** Set when `beforeApplicationShutdown()` begins: from then on, the store grants no new locks. */
+  private closing = false;
 
   constructor(
     private readonly storage: LocksStorage,
@@ -82,7 +86,8 @@ export class Locks implements BeforeApplicationShutdown {
   /**
    * Takes the lock on `key`, waiting up to `wait` for another holder to give it back.
    * Resolves the `Lock` (it renews itself until you release it), or `null` if another holder
-   * kept it, or the store only granted it after `ttl` (the lease ran out on the way). Rejects
+   * kept it, the store only granted it after `ttl` (the lease ran out on the way), or the
+   * application is shutting down (a lock the store grants after that goes back to it). Rejects
    * when the store fails or doesn't answer within `storeTimeout` (`LockStoreTimeoutError`), and
    * with `signal`'s reason when it aborts, also while the store has yet to answer (a lock it
    * grants after that goes back to it).
@@ -102,10 +107,14 @@ export class Locks implements BeforeApplicationShutdown {
           void lock.release().catch(() => undefined);
           signal.throwIfAborted();
         }
+        if (this.closing) {
+          // The shutdown began between the store's answer and here, and released it with the rest.
+          return null;
+        }
         return lock;
       }
       const remaining = deadline - this.clock.now();
-      if (remaining <= 0) {
+      if (remaining <= 0 || this.closing) {
         return null;
       }
       signal?.throwIfAborted();
@@ -118,14 +127,15 @@ export class Locks implements BeforeApplicationShutdown {
   /**
    * Runs `fn` under the lock on `key`, then releases it, whether `fn` resolves or throws.
    * Rejects with `LockNotAcquiredError` if another holder kept the lock for `wait` (default:
-   * don't wait). Inside `fn`, `LocksContext` has the lock too. If the lock is lost while `fn`
-   * runs, `lock.signal` aborts with a `LockLostError`; `fn` decides whether to stop.
+   * don't wait), or the application is shutting down (its `reason` says which). Inside `fn`,
+   * `LocksContext` has the lock too. If the lock is lost while `fn` runs, `lock.signal` aborts
+   * with a `LockLostError`; `fn` decides whether to stop.
    */
   async withLock<T>(key: string, fn: (lock: Lock) => T | Promise<T>, options: LockOptions = {}): Promise<T> {
     const { wait } = this.resolve(key, options);
     const lock = await this.acquire(key, options);
     if (!lock) {
-      throw new LockNotAcquiredError(key, wait);
+      throw new LockNotAcquiredError(key, wait, this.closing ? 'shutdown' : 'taken');
     }
 
     try {
@@ -155,9 +165,20 @@ export class Locks implements BeforeApplicationShutdown {
    * jobs finish and elections step down), and before any `onApplicationShutdown`, where
    * the app closes the store's connection. Nest runs each phase for the whole application,
    * and inside a phase a module imported before `LocksModule` runs first, so a release in
-   * `onApplicationShutdown` could find the pool already ended.
+   * `onApplicationShutdown` could find the pool already ended. From here on, the store grants
+   * no new locks: `acquire()` resolves `null` without a store call, a caller waiting for a key
+   * stops waiting, and a grant that lands after that goes back to the store.
    */
   async beforeApplicationShutdown() {
+    // Before the releases below wake the callers waiting for these keys: a lock granted from
+    // here on would miss the snapshot, and renew after onApplicationShutdown() closed the store.
+    this.closing = true;
+    // A wake removes itself from its set, and an empty set from the map: safe mid-iteration.
+    for (const waiters of this.waiters.values()) {
+      for (const wake of waiters) {
+        wake();
+      }
+    }
     await Promise.all(
       [...this.held].map((lock) =>
         lock.release().catch((error: unknown) => {
@@ -177,6 +198,9 @@ export class Locks implements BeforeApplicationShutdown {
     { onLost, signal }: { onLost?: (lock: Lock, error: LockLostError) => void; signal?: AbortSignal } = {},
   ): Promise<Lock | null> {
     signal?.throwIfAborted();
+    if (this.closing) {
+      return null;
+    }
     const store = this.storage.source;
     const owner = randomUUID();
     const ask = <T>(method: keyof LockStore, answer: Promise<T>, late?: (answer: T) => void) =>
@@ -194,6 +218,11 @@ export class Locks implements BeforeApplicationShutdown {
     if (!result?.acquired) {
       return null;
     }
+    if (this.closing) {
+      // Granted after the shutdown began: nothing would release it, so give it back now.
+      await giveBack();
+      return null;
+    }
 
     const { fencingToken } = result;
 

tests/lock-lifecycle.spec.ts

@@ -436,6 +436,67 @@ describe('Locks: edge cases', () => {
       expect(logger.matching('ERROR Could not release the lock "stuck" at shutdown')).toHaveLength(1);
     });
 
+    it("doesn't hand a released lock to a caller waiting for it, which would renew after the store closed", async () => {
+      const { app, locks } = await instance();
+      await locks.acquire('k');
+      const waiting = locks.acquire('k', { wait: '1m' });
+      await clock.advance(0);
+      const acquire = vi.spyOn(store, 'acquire');
+      const renew = vi.spyOn(store, 'renew');
+      apps = [];
+      await app.close();
+
+      expect(await waiting).toBeNull();
+      expect(acquire).not.toHaveBeenCalled();
+      expect(clock.pendingTimers).toBe(0);
+      await clock.advance('1m');
+      expect(renew).not.toHaveBeenCalled();
+      expect(await store.acquire('k', 'next', 1_000)).toMatchObject({ acquired: true });
+    });
+
+    it('stops a wait for a key another instance holds at once', async () => {
+      const other = await instance();
+      const { app, locks } = await instance();
+      await other.locks.acquire('k');
+      const waiting = locks.acquire('k', { wait: '1m' });
+      await clock.advance(0);
+      const acquire = vi.spyOn(store, 'acquire');
+      apps.splice(apps.indexOf(app), 1);
+      await app.close();
+
+      expect(await waiting).toBeNull();
+      expect(acquire).not.toHaveBeenCalled();
+    });
+
+    it('gives back a lock the store grants after the shutdown began, and grants none after', async () => {
+      const { app, locks } = await instance();
+      const answer = deferred();
+      const grant = store.acquire.bind(store);
+      vi.spyOn(store, 'acquire').mockImplementationOnce(async (key, owner, ttl) => {
+        await answer.promise;
+        return grant(key, owner, ttl);
+      });
+      const acquiring = locks.acquire('k');
+      await clock.advance(0);
+      apps = [];
+      const closing = app.close();
+      await clock.advance(0);
+      const release = vi.spyOn(store, 'release');
+      answer.resolve();
+      expect(await acquiring).toBeNull();
+      await closing;
+      expect(release.mock.calls.map(([key]) => key)).toEqual(['k']);
+      expect(clock.pendingTimers).toBe(0);
+
+      await expect(locks.acquire('k')).resolves.toBeNull();
+      await expect(locks.withLock('k', () => 'ran', { wait: '1m' })).rejects.toMatchObject({
+        name: 'LockNotAcquiredError',
+        reason: 'shutdown',
+        message: 'Could not acquire the lock "k": the application is shutting down',
+      });
+      expect(store.acquire).toHaveBeenCalledTimes(1);
+    });
+
     it('completes LocksEvents.events$', async () => {
       const { app } = await instance();
       const complete = vi.fn();

tests/locks.spec.ts

@@ -153,7 +153,7 @@ describe('Locks', () => {
       await rejected;
       const error = await withLock.catch((e: unknown) => e);
       expect(error).toBeInstanceOf(LocksError);
-      expect(error).toMatchObject({ name: 'LockNotAcquiredError', key: 'k', waitMs: 2_000 });
+      expect(error).toMatchObject({ name: 'LockNotAcquiredError', key: 'k', waitMs: 2_000, reason: 'taken' });
       expect((error as Error).message).toBe('Could not acquire the lock "k" within 2000ms: another holder has it');
       expect(run).not.toHaveBeenCalled();
     });