move to MultiMutex

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', () => {
@@ -39,9 +39,49 @@ describe('Mutex', () => {
     unlock1();
     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 Mutex(3);
+    const mutex = new MultiMutex(3);
     const unlock1 = await mutex.lock();
     expect(mutex.isLocked()).toBe(false);
     const unlock2 = await mutex.lock();

src/index.ts

@@ -1,9 +1,44 @@
 export class Mutex {
+  private _locking: Promise<void>;
+
+  private _locks: number;
+
+  constructor() {
+    this._locking = Promise.resolve();
+    this._locks = 0;
+  }
+
+  isLocked() {
+    return this._locks > 0;
+  }
+
+  lock() {
+    this._locks += 1;
+
+    let unlockNext: () => void;
+
+    const willLock = new Promise<void>(
+      (resolve) =>
+        (unlockNext = () => {
+          this._locks -= 1;
+          resolve();
+        }),
+    );
+
+    const willUnlock = this._locking.then(() => unlockNext);
+
+    this._locking = this._locking.then(() => willLock);
+
+    return willUnlock;
+  }
+}
+
+export class MultiMutex {
   private _queue: (() => void)[];
   private _limit: number;
   private _locks: number;
 
-  constructor(limit = 1) {
+  constructor(limit: number) {
     this._queue = [];
     this._limit = limit;
     this._locks = 0;