From bc00a3b103e46bcc0895e3e8df7f003c1781b1b9 Mon Sep 17 00:00:00 2001
From: Jonathan LEI <me@xjonathan.dev>
Date: Mon, 16 Dec 2024 16:14:54 +0800
Subject: [PATCH] feat: Promise-based mutex type for wasm

---
 packages/account-wasm/src/lib.rs        |  1 +
 packages/account-wasm/src/sync.rs       | 70 +++++++++++++++++++++++++
 packages/account-wasm/src/wasm-mutex.js | 13 +++++
 3 files changed, 84 insertions(+)
 create mode 100644 packages/account-wasm/src/sync.rs
 create mode 100644 packages/account-wasm/src/wasm-mutex.js

diff --git a/packages/account-wasm/src/lib.rs b/packages/account-wasm/src/lib.rs
index dff12ab66..aa3753294 100644
--- a/packages/account-wasm/src/lib.rs
+++ b/packages/account-wasm/src/lib.rs
@@ -5,5 +5,6 @@ pub mod account;
 pub mod session;
 
 mod errors;
+mod sync;
 mod types;
 mod utils;
diff --git a/packages/account-wasm/src/sync.rs b/packages/account-wasm/src/sync.rs
new file mode 100644
index 000000000..a2f69dde2
--- /dev/null
+++ b/packages/account-wasm/src/sync.rs
@@ -0,0 +1,70 @@
+use std::{
+    ops::{Deref, DerefMut},
+    sync::{Mutex as StdMutex, MutexGuard as StdMutexGuard},
+};
+
+use wasm_bindgen::prelude::*;
+use wasm_bindgen_futures::js_sys::Function;
+
+/// A mutex implementation backed by JavaScript `Promise`s.
+///
+/// This type wraps a simple JavaScript `Mutex` implementation but exposes an idiomatic Rust API.
+pub struct WasmMutex<T> {
+    js_lock: Mutex,
+    rs_lock: StdMutex<T>,
+}
+
+impl<T> WasmMutex<T> {
+    pub fn new(value: T) -> Self {
+        Self {
+            js_lock: Mutex::new(),
+            rs_lock: std::sync::Mutex::new(value),
+        }
+    }
+
+    pub async fn lock(&self) -> WasmMutexGuard<T> {
+        WasmMutexGuard {
+            js_release: self.js_lock.obtain().await,
+            // This never actually blocks as it's guarded by the JS lock. This field exists only to
+            // provide internal mutability for the underlying value.
+            rs_guard: self.rs_lock.lock().unwrap(),
+        }
+    }
+}
+
+/// A handle to the underlying guarded value. The lock is released when the instance is dropped.
+pub struct WasmMutexGuard<'a, T> {
+    js_release: Function,
+    rs_guard: StdMutexGuard<'a, T>,
+}
+
+impl<T> Deref for WasmMutexGuard<'_, T> {
+    type Target = T;
+
+    fn deref(&self) -> &T {
+        std::sync::MutexGuard::deref(&self.rs_guard)
+    }
+}
+
+impl<T> DerefMut for WasmMutexGuard<'_, T> {
+    fn deref_mut(&mut self) -> &mut T {
+        std::sync::MutexGuard::deref_mut(&mut self.rs_guard)
+    }
+}
+
+impl<T> Drop for WasmMutexGuard<'_, T> {
+    fn drop(&mut self) {
+        self.js_release.call0(&JsValue::null()).unwrap();
+    }
+}
+
+#[wasm_bindgen(module = "/src/wasm-mutex.js")]
+extern "C" {
+    type Mutex;
+
+    #[wasm_bindgen(constructor)]
+    fn new() -> Mutex;
+
+    #[wasm_bindgen(method)]
+    async fn obtain(this: &Mutex) -> Function;
+}
diff --git a/packages/account-wasm/src/wasm-mutex.js b/packages/account-wasm/src/wasm-mutex.js
new file mode 100644
index 000000000..d3babff76
--- /dev/null
+++ b/packages/account-wasm/src/wasm-mutex.js
@@ -0,0 +1,13 @@
+function releaseStub() {}
+
+export class Mutex {
+  lastPromise = Promise.resolve();
+
+  async obtain() {
+    let release = releaseStub;
+    const lastPromise = this.lastPromise;
+    this.lastPromise = new Promise((resolve) => (release = resolve));
+    await lastPromise;
+    return release;
+  }
+}