Add Semaphore utility to help implement connection pools (#906)

Author: simolus3

.changeset/fair-toes-dance.md

@@ -0,0 +1,6 @@
+---
+'@powersync/node': patch
+'@powersync/op-sqlite': patch
+---
+
+Internal: Refactor connection pool to use `Semaphore` from `@powersync/common`.

.changeset/short-moons-hug.md

@@ -0,0 +1,5 @@
+---
+'@powersync/common': minor
+---
+
+Internal: Add `Semaphore` utility to manage leases in a connection pool.

packages/common/src/utils/mutex.ts

@@ -1,21 +1,32 @@
+import { Queue } from './queue.js';
+
 export type UnlockFn = () => void;
 
 /**
- * An asynchronous mutex implementation.
+ * An asynchronous semaphore implementation with associated items per lease.
  *
  * @internal This class is meant to be used in PowerSync SDKs only, and is not part of the public API.
  */
-export class Mutex {
-  private inCriticalSection = false;
+export class Semaphore<T> {
+  // Available items that are not currently assigned to a waiter.
+  private readonly available: Queue<T>;
 
+  readonly size: number;
   // Linked list of waiters. We don't expect the wait list to become particularly large, and this allows removing
   // aborted waiters from the middle of the list efficiently.
-  private firstWaiter?: MutexWaitNode;
-  private lastWaiter?: MutexWaitNode;
+  private firstWaiter?: SemaphoreWaitNode<T>;
+  private lastWaiter?: SemaphoreWaitNode<T>;
 
-  private addWaiter(onAcquire: () => void): MutexWaitNode {
-    const node: MutexWaitNode = {
+  constructor(elements: Iterable<T>) {
+    this.available = new Queue(elements);
+    this.size = this.available.length;
+  }
+
+  private addWaiter(requestedItems: number, onAcquire: () => void): SemaphoreWaitNode<T> {
+    const node: SemaphoreWaitNode<T> = {
       isActive: true,
+      acquiredItems: [],
+      remainingItems: requestedItems,
       onAcquire,
       prev: this.lastWaiter
     };
@@ -30,7 +41,7 @@ export class Mutex {
     return node;
   }
 
-  private deactivateWaiter(waiter: MutexWaitNode) {
+  private deactivateWaiter(waiter: SemaphoreWaitNode<T>) {
     const { prev, next } = waiter;
     waiter.isActive = false;
 
@@ -40,77 +51,129 @@ export class Mutex {
     if (waiter == this.lastWaiter) this.lastWaiter = prev;
   }
 
-  acquire(abort?: AbortSignal): Promise<UnlockFn> {
+  private requestPermits(amount: number, abort?: AbortSignal): Promise<{ items: T[]; release: UnlockFn }> {
+    if (amount <= 0 || amount > this.size) {
+      throw new Error(`Invalid amount of items requested (${amount}), must be between 1 and ${this.size}`);
+    }
+
     return new Promise((resolve, reject) => {
       function rejectAborted() {
-        reject(abort?.reason ?? new Error('Mutex acquire aborted'));
+        reject(abort?.reason ?? new Error('Semaphore acquire aborted'));
       }
       if (abort?.aborted) {
         return rejectAborted();
       }
 
-      let holdsMutex = false;
+      let waiter: SemaphoreWaitNode<T>;
 
       const markCompleted = () => {
-        if (!holdsMutex) return;
-        holdsMutex = false;
+        const items = waiter.acquiredItems;
+        waiter.acquiredItems = []; // Avoid releasing items twice.
+
+        for (const element of items) {
+          // Give to next waiter, if possible.
+          const nextWaiter = this.firstWaiter;
+          if (nextWaiter) {
+            nextWaiter.acquiredItems.push(element);
+            nextWaiter.remainingItems--;
+            if (nextWaiter.remainingItems == 0) {
+              nextWaiter.onAcquire();
+            }
+          } else {
+            // No pending waiter, return lease into pool.
+            this.available.addLast(element);
+          }
+        }
+      };
+
+      const onAbort = () => {
+        abort?.removeEventListener('abort', onAbort);
 
-        const waiter = this.firstWaiter;
-        if (waiter) {
+        if (waiter.isActive) {
           this.deactivateWaiter(waiter);
-          // Still in critical section, but owned by next waiter now.
-          waiter.onAcquire();
-        } else {
-          this.inCriticalSection = false;
+          rejectAborted();
         }
       };
 
-      if (!this.inCriticalSection) {
-        this.inCriticalSection = true;
-        holdsMutex = true;
-        return resolve(markCompleted);
-      } else {
-        let node: MutexWaitNode;
+      const resolvePromise = () => {
+        this.deactivateWaiter(waiter);
+        abort?.removeEventListener('abort', onAbort);
 
-        const onAbort = () => {
-          abort?.removeEventListener('abort', onAbort);
+        const items = waiter.acquiredItems;
+        resolve({ items, release: markCompleted });
+      };
 
-          if (node.isActive) {
-            this.deactivateWaiter(node);
-            rejectAborted();
-          }
-        };
+      waiter = this.addWaiter(amount, resolvePromise);
 
-        node = this.addWaiter(() => {
-          abort?.removeEventListener('abort', onAbort);
-          holdsMutex = true;
-          resolve(markCompleted);
-        });
+      // If there are items in the pool that haven't been assigned, we can pull them into this waiter. Note that this is
+      // only the case if we're the first waiter (otherwise, items would have been assigned to an earlier waiter).
+      while (!this.available.isEmpty && waiter.remainingItems > 0) {
+        waiter.acquiredItems.push(this.available.removeFirst());
+        waiter.remainingItems--;
+      }
 
-        abort?.addEventListener('abort', onAbort);
+      if (waiter.remainingItems == 0) {
+        return resolvePromise();
       }
+
+      abort?.addEventListener('abort', onAbort);
     });
   }
 
-  async runExclusive<T>(fn: () => PromiseLike<T> | T, abort?: AbortSignal): Promise<T> {
-    const returnMutex = await this.acquire(abort);
+  /**
+   * Requests a single item from the pool.
+   *
+   * The returned `release` callback must be invoked to return the item into the pool.
+   */
+  async requestOne(abort?: AbortSignal): Promise<{ item: T; release: UnlockFn }> {
+    const { items, release } = await this.requestPermits(1, abort);
+    return { release, item: items[0] };
+  }
 
-    try {
-      return await fn();
-    } finally {
-      returnMutex();
-    }
+  /**
+   * Requests access to all items from the pool.
+   *
+   * The returned `release` callback must be invoked to return items into the pool.
+   */
+  requestAll(abort?: AbortSignal): Promise<{ items: T[]; release: UnlockFn }> {
+    return this.requestPermits(this.size, abort);
   }
 }
 
-interface MutexWaitNode {
+interface SemaphoreWaitNode<T> {
   /**
    * Whether the waiter is currently active (not aborted and not fullfilled).
    */
   isActive: boolean;
+  acquiredItems: T[];
+  remainingItems: number;
   onAcquire: () => void;
-  prev?: MutexWaitNode;
-  next?: MutexWaitNode;
+  prev?: SemaphoreWaitNode<T>;
+  next?: SemaphoreWaitNode<T>;
+}
+
+/**
+ * An asynchronous mutex implementation.
+ *
+ * @internal This class is meant to be used in PowerSync SDKs only, and is not part of the public API.
+ */
+export class Mutex {
+  private inner = new Semaphore([null]);
+
+  async acquire(abort?: AbortSignal): Promise<UnlockFn> {
+    const { release } = await this.inner.requestOne(abort);
+    return release;
+  }
+
+  async runExclusive<T>(fn: () => PromiseLike<T> | T, abort?: AbortSignal): Promise<T> {
+    const returnMutex = await this.acquire(abort);
+
+    try {
+      return await fn();
+    } finally {
+      returnMutex();
+    }
+  }
 }
 
 /**

packages/common/src/utils/queue.ts

@@ -0,0 +1,48 @@
+/**
+ * A simple fixed-capacity queue implementation.
+ *
+ * Unlike a naive queue implemented by `array.push()` and `array.shift()`, this avoids moving array elements around
+ * and is `O(1)` for {@link addLast} and {@link removeFirst}.
+ */
+export class Queue<T> {
+  private table: (T | undefined)[];
+  // Index of the first element in the table.
+  private head: number;
+  // Amount of items currently in the queue.
+  private _length: number;
+
+  constructor(initialItems: Iterable<T>) {
+    this.table = [...initialItems];
+    this.head = 0;
+    this._length = this.table.length;
+  }
+
+  get isEmpty(): boolean {
+    return this.length == 0;
+  }
+
+  get length(): number {
+    return this._length;
+  }
+
+  removeFirst(): T {
+    if (this.isEmpty) {
+      throw new Error('Queue is empty');
+    }
+
+    const result = this.table[this.head] as T;
+    this._length--;
+    this.table[this.head] = undefined;
+    this.head = (this.head + 1) % this.table.length;
+    return result;
+  }
+
+  addLast(element: T) {
+    if (this.length == this.table.length) {
+      throw new Error('Queue is full');
+    }
+
+    this.table[(this.head + this._length) % this.table.length] = element;
+    this._length++;
+  }
+}