feat: support multiple lock holders (#1)

package.json

@@ -3,7 +3,7 @@
   "private": false,
   "access": "public",
   "license": "Apache-2.0",
-  "version": "1.0.0",
+  "version": "1.1.0",
   "description": "Tiny mutex helper",
   "repository": "[email protected]:livekit/ts-mutex.git",
   "exports": {

src/index.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, it } from 'vitest';
-import { Mutex } from './index';
+import { Mutex, MultiMutex } from './index';
 
 describe('Mutex', () => {
   it('should not be locked initially', () => {
@@ -40,3 +40,61 @@ describe('Mutex', () => {
     expect(mutex.isLocked()).toBe(false);
   });
 });
+
+describe('MultiMutex', () => {
+  it('should not be locked initially', () => {
+    const mutex = new MultiMutex(3);
+    expect(mutex.isLocked()).toBe(false);
+  });
+
+  it('should lock and unlock correctly', async () => {
+    const mutex = new MultiMutex(1);
+    const unlock = await mutex.lock();
+    expect(mutex.isLocked()).toBe(true);
+    unlock();
+    expect(mutex.isLocked()).toBe(false);
+  });
+
+  it('should handle multiple locks', async () => {
+    const mutex = new MultiMutex(1);
+    const unlock1 = await mutex.lock();
+    const unlock2Promise = mutex.lock();
+    expect(mutex.isLocked()).toBe(true);
+    unlock1();
+    expect(mutex.isLocked()).toBe(true);
+    const unlock3Promise = mutex.lock();
+    expect(mutex.isLocked()).toBe(true);
+    (await unlock2Promise)();
+    expect(mutex.isLocked()).toBe(true);
+    (await unlock3Promise)();
+    expect(mutex.isLocked()).toBe(false);
+  });
+
+  it('should not care about unlocking the same lock twice', async () => {
+    const mutex = new MultiMutex(1);
+    const unlock1 = await mutex.lock();
+    expect(mutex.isLocked()).toBe(true);
+    unlock1();
+    expect(mutex.isLocked()).toBe(false);
+    unlock1();
+    expect(mutex.isLocked()).toBe(false);
+  });
+
+  it('should support multiple locks being used at once', async () => {
+    const mutex = new MultiMutex(3);
+    const unlock1 = await mutex.lock();
+    expect(mutex.isLocked()).toBe(false);
+    const unlock2 = await mutex.lock();
+    expect(mutex.isLocked()).toBe(false);
+    const unlock3 = await mutex.lock();
+    expect(mutex.isLocked()).toBe(true);
+    const unlock4Promise = mutex.lock();
+    unlock1();
+    expect(mutex.isLocked()).toBe(true);
+    unlock2();
+    expect(mutex.isLocked()).toBe(false);
+    unlock3();
+    expect(mutex.isLocked()).toBe(false);
+    await unlock4Promise.then((unlock) => unlock);
+  });
+});

src/index.ts

@@ -32,3 +32,41 @@ export class Mutex {
     return willUnlock;
   }
 }
+
+export class MultiMutex {
+  private _queue: (() => void)[];
+  private _limit: number;
+  private _locks: number;
+
+  constructor(limit: number) {
+    this._queue = [];
+    this._limit = limit;
+    this._locks = 0;
+  }
+
+  isLocked() {
+    return this._locks >= this._limit;
+  }
+
+  async lock(): Promise<() => void> {
+    if (!this.isLocked()) {
+      this._locks++;
+      return this._unlock.bind(this);
+    }
+
+    return new Promise((resolve) => {
+      this._queue.push(() => {
+        this._locks++;
+        resolve(this._unlock.bind(this));
+      });
+    });
+  }
+
+  private _unlock() {
+    this._locks--;
+    if (this._queue.length && !this.isLocked()) {
+      const nextUnlock = this._queue.shift();
+      nextUnlock?.();
+    }
+  }
+}