better-lock 3.3.0 → 4.0.0

package/src/better_lock.ts

@@ -1,9 +1,11 @@
+import { captureAsyncContext, getHeldJobs, runHoldingJob } from './async_context';
 import {
   BetterLockExecutionTimeoutError,
   BetterLockInternalError,
   BetterLockInvalidArgumentError,
   BetterLockJobAbortedError,
   BetterLockQueueOverflowError,
+  BetterLockReentrancyError,
   BetterLockWaitTimeoutError,
   isAcquisitionFailureError,
 } from './errors';
@@ -22,12 +24,18 @@ export class BetterLock {
   static DEFAULT_OPTIONS: IOptions = {
     extend_stack_traces: true,
     promise_tester: p => p && typeof p.then === 'function',
+    reentrancy: 'allow',
   };
 
   private readonly options: IOptions;
   private readonly log: (msg: string) => void;
   private readonly queues: { [key: string]: KeyQueue };
 
+  /**
+   * Reentrant jobs that haven't started executing yet. They are not in any key queue, so we track them here for abort().
+   */
+  private readonly pendingReentrantJobs: LockJob<any>[] = [];
+
   /**
    * Create a better lock instance
    */
@@ -52,11 +60,49 @@ export class BetterLock {
   }
 
   /**
-   * Returns true if the caller can *immediately* acquire the lock. There is nothing holding the key and nothing in the queue.
+   * Returns true if the caller can *immediately* acquire the lock. There is nothing holding the key and nothing in the queue,
+   * or the caller already holds the key and reentrancy is allowed.
    */
   public canAcquire(key?: string): boolean {
     const queue = this.getQueue(key);
-    return !queue.active && !queue.jobs.length;
+    if (!queue.active && !queue.jobs.length) {
+      return true;
+    }
+    return this.options.reentrancy !== 'error' && this.isHeldByCaller(queue, getHeldJobs());
+  }
+
+  /**
+   * True if the given queue's key is held by a job that is executing in the caller's async context.
+   * Ended jobs don't count, so detached work that outlives its job doesn't skip the lock.
+   */
+  private isHeldByCaller(queue: KeyQueue, heldJobs: ReadonlyArray<LockJob<any>> | undefined) {
+    const holder = queue.active;
+    return !!(heldJobs && holder && !holder.ended_at && heldJobs.indexOf(holder) >= 0);
+  }
+
+  /**
+   * Find which of the given keys the caller already holds, and the jobs holding them.
+   * Returns null if the caller holds none of them.
+   */
+  private findHeldKeys(
+    keys: Array<string | null>
+  ): { holders: LockJob<any>[]; heldKeys: Array<string | null> } | null {
+    const heldJobs = getHeldJobs();
+    if (!heldJobs || !keys.length) {
+      return null;
+    }
+
+    const holders: LockJob<any>[] = [];
+    const heldKeys: Array<string | null> = [];
+    for (let i = 0; i < keys.length; i++) {
+      const queue = this.getQueue(keys[i]);
+      if (this.isHeldByCaller(queue, heldJobs)) {
+        holders.push(queue.active);
+        heldKeys.push(keys[i]);
+      }
+    }
+
+    return holders.length ? { holders, heldKeys } : null;
   }
 
   private normalizeAndValidateKey(key: ILockKey) {
@@ -114,6 +160,10 @@ export class BetterLock {
    * Acquire the lock for given key or list of keys. If waiting on a list, keys will be acquired in order,
    * as they become free. Any deadlock prevention must be handled by the caller.
    *
+   * If the caller already holds all the keys (it runs inside an executor holding them), the executor runs
+   * without waiting. If it holds only some of the keys, or the "reentrancy" option is set to "error", the
+   * callback (or the promise) gets BetterLockReentrancyError. Reentrancy requires Node.js AsyncLocalStorage.
+   *
    * Once the lock is acquired, we will call the executor function, with a "done" method.
    * You must call "done" to release the lock. Alternatively, return a Promise; lock will be released
    * once promise resolves or rejects.
@@ -187,8 +237,27 @@ export class BetterLock {
       );
     }
 
+    // Invalid arguments throw above. Reentrancy errors depend on where acquire is called from,
+    // so they are reported like other lock errors: through the callback or promise.
+    const held = this.findHeldKeys(keys);
+    if (
+      held &&
+      // Waiting for the remaining keys while holding these could deadlock
+      (held.heldKeys.length < keys.length || effectiveJobOptions.reentrancy === 'error')
+    ) {
+      const err = new BetterLockReentrancyError(this.options.name, keys, held.heldKeys);
+      this.log(err.message);
+      // Never call back synchronously from acquire
+      setTimeout(() => callback(err), 0);
+      return (callback as any).promise;
+    }
+    const reentrantHolders = held ? held.holders : null;
+
     const job = new LockJob(keys, executor, callback, effectiveJobOptions);
 
+    // Executor and callback will run in the caller's async context, not the context of whoever releases the lock
+    job.async_context = captureAsyncContext();
+
     // Set incoming stack
     if (
       effectiveJobOptions.extend_stack_traces &&
@@ -199,10 +268,26 @@ export class BetterLock {
       job.incoming_stack = tempErr.stack;
     }
 
-    // Add job to its key queues
+    if (reentrantHolders) {
+      // The caller holds all the keys. Run without entering the queues.
+      job.reentrant_holders = reentrantHolders;
+      this.pendingReentrantJobs.push(job);
+      this.log(`Reentering ${job}`);
+      this.executeJob(job);
+    } else {
+      this.enqueueJob(job);
+    }
+
+    return (callback as any).promise;
+  }
+
+  /**
+   * Add job to its key queues and start waiting for the keys
+   */
+  private enqueueJob(job: LockJob<any>) {
     const queuesToUpdate = [];
-    for (let i = 0; i < keys.length; i++) {
-      const queue = this.getQueue(keys[i]);
+    for (let i = 0; i < job.keys.length; i++) {
+      const queue = this.getQueue(job.keys[i]);
       queue.jobs.push(job);
       queuesToUpdate.push(queue);
     }
@@ -224,8 +309,6 @@ export class BetterLock {
       // Special case. We are given empty list of keys. We will immediately execute jobs like this.
       this.executeJob(job);
     }
-
-    return (callback as any).promise;
   }
 
   public acquireOr<TResult extends IValidResult>(
@@ -365,17 +448,30 @@ export class BetterLock {
         return;
       }
 
-      // Validate this job is holding all queues it needs
-      for (let i = 0; i < job.keys.length; i++) {
-        const queue = this.getQueue(job.keys[i]);
-        if (queue.active !== job) {
-          // This should never happen
-          return this.endJob(job, [
-            new BetterLockInternalError(
-              this.options.name,
-              `Corrupted wait queue state for ${queue.key}`
-            ),
-          ]);
+      if (job.reentrant_holders) {
+        this.removePendingReentrantJob(job);
+
+        if (job.reentrant_holders.some(holder => holder.ended_at)) {
+          // The holder released the keys before this job could start (eg. the reentrant acquire wasn't awaited).
+          // Wait for the keys like any other job.
+          this.log(`${job} is no longer reentrant because its holder has ended`);
+          job.reentrant_holders = null;
+          this.enqueueJob(job);
+          return;
+        }
+      } else {
+        // Validate this job is holding all queues it needs
+        for (let i = 0; i < job.keys.length; i++) {
+          const queue = this.getQueue(job.keys[i]);
+          if (queue.active !== job) {
+            // This should never happen
+            return this.endJob(job, [
+              new BetterLockInternalError(
+                this.options.name,
+                `Corrupted wait queue state for ${queue.key}`
+              ),
+            ]);
+          }
         }
       }
 
@@ -395,30 +491,43 @@ export class BetterLock {
         this.endJob(job, args);
       };
 
-      if (job.executor.length > 0) {
-        // Callback interface
-        job.executor(lockDone);
-        return;
-      }
+      const runExecutor = () => {
+        if (job.executor.length > 0) {
+          // Callback interface
+          job.executor(lockDone);
+          return;
+        }
 
-      // Promise interface
-      let executorResult;
-      try {
-        executorResult = (job.executor as any)();
-      } catch (err) {
-        // Error thrown directly from executor
-        lockDone(err);
-      }
+        // Promise interface
+        let executorResult;
+        try {
+          executorResult = (job.executor as any)();
+        } catch (err) {
+          // Error thrown directly from executor
+          lockDone(err);
+        }
 
-      // Promise sniffing
-      if (this.options.promise_tester(executorResult)) {
-        executorResult.then(
-          res => lockDone(null, res),
-          err => lockDone(err)
-        );
+        // Promise sniffing
+        if (this.options.promise_tester(executorResult)) {
+          executorResult.then(
+            res => lockDone(null, res),
+            err => lockDone(err)
+          );
+        } else {
+          // "promise" is just some random value. We don't have to wait
+          lockDone(null, executorResult);
+        }
+      };
+
+      if (!job.async_context) {
+        // No async context support in this runtime
+        runExecutor();
+      } else if (job.reentrant_holders) {
+        // The caller's context already records the holders
+        job.async_context.runInAsyncScope(runExecutor);
       } else {
-        // "promise" is just some random value. We don't have to wait
-        lockDone(null, executorResult);
+        // Restore the caller's context first, then record this job as holding its keys inside it
+        job.async_context.runInAsyncScope(() => runHoldingJob(job, runExecutor));
       }
     }, 0);
   }
@@ -459,9 +568,15 @@ export class BetterLock {
 
     job.ended_at = new Date();
 
+    // Reentrant jobs never enter key queues, so they have nothing to release
+    const keysToRelease = job.reentrant_holders ? [] : job.keys;
+    if (job.reentrant_holders) {
+      this.removePendingReentrantJob(job);
+    }
+
     const queuesToUpdate = [];
-    for (let i = 0; i < job.keys.length; i++) {
-      const queue = this.getQueue(job.keys[i]);
+    for (let i = 0; i < keysToRelease.length; i++) {
+      const queue = this.getQueue(keysToRelease[i]);
       queuesToUpdate.push(queue);
 
       if (queue.active === job) {
@@ -484,7 +599,11 @@ export class BetterLock {
 
     // Call the callback
     try {
-      job.callback.apply(null, callbackArgs);
+      if (job.async_context) {
+        job.async_context.runInAsyncScope(job.callback as any, null, ...callbackArgs);
+      } else {
+        job.callback.apply(null, callbackArgs);
+      }
     } finally {
       // Whatever happens, schedule a queue update
       setTimeout(() => this.update(queuesToUpdate), 0);
@@ -506,6 +625,19 @@ export class BetterLock {
       // If we have a job holding this queue that hasn't been executed yet, abort it as well
       this.endJob(queue.active, [new BetterLockJobAbortedError(this.options.name, queue.active)]);
     }
+    // Reentrant jobs for this key that haven't started yet
+    this.pendingReentrantJobs.slice().forEach(job => {
+      if (job.keys.some(jobKey => this.getQueue(jobKey) === queue)) {
+        this.endJob(job, [new BetterLockJobAbortedError(this.options.name, job)]);
+      }
+    });
+  }
+
+  private removePendingReentrantJob(job: LockJob<any>) {
+    const index = this.pendingReentrantJobs.indexOf(job);
+    if (index >= 0) {
+      this.pendingReentrantJobs.splice(index, 1);
+    }
   }
 
   /**

package/src/errors.ts

@@ -131,6 +131,34 @@ export class BetterLockJobAbortedError extends BetterLockError {
   }
 }
 
+export class BetterLockReentrancyError extends BetterLockError {
+  /**
+   * Keys requested by the rejected acquire call
+   */
+  public keys: string[];
+
+  /**
+   * Requested keys that the caller already holds
+   */
+  public held_keys: string[];
+
+  constructor(lockName, keys: string[], heldKeys: string[]) {
+    const describe = (list: string[]) =>
+      list.map(key => (key === null ? '<default>' : `"${key}"`)).join(', ');
+    const message =
+      heldKeys.length < keys.length
+        ? `Caller holds only some of the requested keys (${describe(
+            heldKeys
+          )} of ${describe(keys)}). Waiting for the rest while holding these could deadlock`
+        : `Reentrant acquire of ${describe(keys)} is not allowed (reentrancy: "error")`;
+
+    super(lockName, message);
+
+    this.keys = keys;
+    this.held_keys = heldKeys;
+  }
+}
+
 export function isAcquisitionFailureError(err) {
   return (
     (err && err instanceof BetterLockQueueOverflowError) ||

package/src/index.ts

@@ -6,9 +6,10 @@ import {
   BetterLockInvalidArgumentError,
   BetterLockJobAbortedError,
   BetterLockQueueOverflowError,
+  BetterLockReentrancyError,
   BetterLockWaitTimeoutError,
 } from './errors';
-import { ILockJobOptions as BetterLockJobOptions, IOptions as BetterLockOptions, IQueueEjectionStrategy as BetterLockQueueEjectionStrategy } from './options';
+import { ILockJobOptions as BetterLockJobOptions, IOptions as BetterLockOptions, IQueueEjectionStrategy as BetterLockQueueEjectionStrategy, IReentrancyMode as BetterLockReentrancyMode } from './options';
 import {
   IErrorName as BetterLockErrorName,
   IExecutor as BetterLockExecutor,
@@ -18,6 +19,7 @@ import {
 export {
   BetterLockOptions,
   BetterLockQueueEjectionStrategy,
+  BetterLockReentrancyMode,
   BetterLockJobOptions,
   BetterLockExecutor,
   BetterLockKey,
@@ -29,6 +31,7 @@ export {
   BetterLockInvalidArgumentError,
   BetterLockJobAbortedError,
   BetterLockQueueOverflowError,
+  BetterLockReentrancyError,
   BetterLockWaitTimeoutError,
 };
 
@@ -50,6 +53,7 @@ export const BetterLock = attachExports(BetterLockUnadorned, {
   ExecutionTimeoutError: BetterLockExecutionTimeoutError,
   QueueOverflowError: BetterLockQueueOverflowError,
   JobAbortedError: BetterLockJobAbortedError,
+  ReentrancyError: BetterLockReentrancyError,
 });
 
 export default BetterLock;

package/src/internals.ts

@@ -1,3 +1,4 @@
+import { IAsyncResource } from './async_context';
 import { ILockJobOptions } from './options';
 import { ICallback, IExecutor } from './types';
 
@@ -66,6 +67,17 @@ export class LockJob<TResult> {
    */
   public incoming_stack: string = null;
 
+  /**
+   * Async context of the acquire() call. Executor and callback run in it. Undefined if the runtime doesn't support it.
+   */
+  public async_context: IAsyncResource | undefined = undefined;
+
+  /**
+   * Set for reentrant jobs: the executing jobs (from the caller's async context) that hold this job's keys.
+   * A reentrant job doesn't enter key queues. It runs while these jobs hold the keys.
+   */
+  public reentrant_holders: LockJob<any>[] | null = null;
+
   constructor(
     keys: Array<string | null>,
     executor: IExecutor<TResult>,

package/src/options.ts

@@ -25,10 +25,21 @@ export interface ILockJobOptions {
    * Overrides BetterLockOptions value.
    */
   lock_condition?: (key: ILockKey | Array<ILockKey>) => boolean;
+
+  /**
+   * What to do when acquire is called from code that already holds all the requested keys. Overrides BetterLockOptions value.
+   */
+  reentrancy?: IReentrancyMode;
 }
 
 export type IQueueEjectionStrategy = 'oldest' | 'newest';
 
+/**
+ * - 'allow': run the executor without waiting for the keys the caller already holds.
+ * - 'error': fail with BetterLockReentrancyError.
+ */
+export type IReentrancyMode = 'allow' | 'error';
+
 /**
  * Options that apply on BetterLock instance
  */
@@ -60,6 +71,19 @@ export interface IOptions extends ILockJobOptions {
    */
   lock_condition?: (key: ILockKey | Array<ILockKey>) => boolean;
 
+  /**
+   * What to do when acquire is called from code that already holds all the requested keys (inside the executor, or
+   * in async work the executor started and that runs before it finishes). 'allow' runs the executor without
+   * waiting, 'error' fails with BetterLockReentrancyError. Defaults to 'allow'.
+   *
+   * Holding only some of the requested keys always fails with BetterLockReentrancyError, because waiting for the rest
+   * while holding these could deadlock.
+   *
+   * Requires Node.js AsyncLocalStorage, looked up through process.getBuiltinModule (Node.js 20.16+ or 22.3+).
+   * In other runtimes, this option has no effect and nested acquire calls wait in the queue like any other.
+   */
+  reentrancy?: IReentrancyMode;
+
   /**
    * Lock name. This will be written in all logs and error messages, to help you distinguish between different locks
    */