From 69e4a5e9b852c75117078f0417734e36987f165b Mon Sep 17 00:00:00 2001
From: pengx17 <pengxiao@outlook.com>
Date: Mon, 20 Jan 2025 17:03:15 +0000
Subject: [PATCH] fix(infra): asynclock deadlock (#9808)

---
 .../src/utils/__tests__/async-lock.spec.ts    | 113 ++++++++++++++++++
 packages/common/infra/src/utils/async-lock.ts |  30 +++--
 2 files changed, 136 insertions(+), 7 deletions(-)
 create mode 100644 packages/common/infra/src/utils/__tests__/async-lock.spec.ts

diff --git a/packages/common/infra/src/utils/__tests__/async-lock.spec.ts b/packages/common/infra/src/utils/__tests__/async-lock.spec.ts
new file mode 100644
index 000000000..c3195789e
--- /dev/null
+++ b/packages/common/infra/src/utils/__tests__/async-lock.spec.ts
@@ -0,0 +1,113 @@
+import { describe, expect, test } from 'vitest';
+
+import { AsyncLock } from '../async-lock';
+
+describe('AsyncLock', () => {
+  test('should acquire and release lock', async () => {
+    const lock = new AsyncLock();
+    const lock1 = await lock.acquire();
+    lock1.release();
+  });
+
+  test('should wait for previous lock to be released', async () => {
+    const lock = new AsyncLock();
+    const order: number[] = [];
+
+    const task1 = async () => {
+      const lock1 = await lock.acquire();
+      order.push(1);
+      await new Promise(resolve => setTimeout(resolve, 10));
+      order.push(2);
+      lock1.release();
+    };
+
+    const task2 = async () => {
+      const lock2 = await lock.acquire();
+      order.push(3);
+      lock2.release();
+    };
+
+    await Promise.all([task1(), task2()]);
+    expect(order).toEqual([1, 2, 3]);
+  });
+
+  test('should work with using statement', async () => {
+    const lock = new AsyncLock();
+    const order: number[] = [];
+
+    const task1 = async () => {
+      using _lock1 = await lock.acquire();
+      order.push(1);
+      await new Promise(resolve => setTimeout(resolve, 10));
+      order.push(2);
+    };
+
+    const task2 = async () => {
+      using _lock2 = await lock.acquire();
+      order.push(3);
+    };
+
+    await Promise.all([task1(), task2()]);
+    expect(order).toEqual([1, 2, 3]);
+  });
+
+  test('should handle multiple concurrent locks', async () => {
+    const lock = new AsyncLock();
+    const results: number[] = [];
+
+    const createTask = (id: number, delay: number) => async () => {
+      using _lockHandle = await lock.acquire();
+      results.push(id);
+      await new Promise(resolve => setTimeout(resolve, delay));
+    };
+
+    await Promise.all([
+      createTask(1, 20)(),
+      createTask(2, 10)(),
+      createTask(3, 5)(),
+    ]);
+
+    expect(results).toEqual([1, 2, 3]);
+  });
+
+  test('should properly block after a release', async () => {
+    const lock = new AsyncLock();
+    const order: number[] = [];
+
+    // First acquisition
+    const lock1 = await lock.acquire();
+    order.push(1);
+    lock1.release();
+
+    // These two should be properly serialized
+    const task1 = async () => {
+      const lock2 = await lock.acquire();
+      order.push(2);
+      await new Promise(resolve => setTimeout(resolve, 10));
+      order.push(3);
+      lock2.release();
+    };
+
+    const task2 = async () => {
+      const lock3 = await lock.acquire();
+      order.push(4);
+      lock3.release();
+    };
+
+    await Promise.all([task1(), task2()]);
+    expect(order).toEqual([1, 2, 3, 4]); // This might fail due to the bug
+  });
+
+  test('should prevent multiple releases', async () => {
+    const lock = new AsyncLock();
+    const handle = await lock.acquire();
+    handle.release();
+
+    // This second release should either throw or be a no-op
+    handle.release();
+
+    // The lock should still be usable
+    const handle2 = await lock.acquire();
+    handle2.release();
+  });
+});
diff --git a/packages/common/infra/src/utils/async-lock.ts b/packages/common/infra/src/utils/async-lock.ts
index c7e5c9c86..0afb4e985 100644
--- a/packages/common/infra/src/utils/async-lock.ts
+++ b/packages/common/infra/src/utils/async-lock.ts
@@ -1,19 +1,35 @@
 export class AsyncLock {
-  private _lock = Promise.resolve();
+  private _lock: Promise<void> | null = null;
 
   async acquire() {
-    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-    let release: () => void = null!;
+    let release: (() => void) | null = null;
     const nextLock = new Promise<void>(resolve => {
-      release = resolve;
+      release = () => {
+        this._lock = null;
+        resolve();
+      };
     });
 
-    await this._lock;
+    // Atomic check and set of lock state
+    const currentLock = this._lock;
     this._lock = nextLock;
+
+    if (currentLock) {
+      await currentLock;
+    }
+
     return {
-      release,
+      release: () => {
+        if (release) {
+          release();
+          release = null;
+        }
+      },
       [Symbol.dispose]: () => {
-        release();
+        if (release) {
+          release();
+          release = null;
+        }
       },
     };
   }