Merge pull request #141 from storycraft/fix/dashmap-deadlock

fix: replace to sccmap

Author: storycraft

Cargo.lock

@@ -19,7 +19,7 @@ asdf-overlay-event = { workspace = true, features = ["bincode"] }
 tokio = { workspace = true, features = ["macros", "time", "net"] }
 
 bincode = "2.0.1"
-dashmap = "6.1.0"
+scc = "3.3.3"
 ntapi = "0.4.1"
 anyhow = "1.0.97"
 scopeguard = "1.2.0"

crates/client/Cargo.toml

@@ -11,7 +11,6 @@ use asdf_overlay_common::{
 };
 use asdf_overlay_event::OverlayEvent;
 use bincode::Decode;
-use dashmap::DashMap;
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt, WriteHalf, split},
     net::windows::named_pipe::NamedPipeClient,
@@ -24,7 +23,7 @@ pub struct IpcClientConn {
     next_id: u32,
     tx: WriteHalf<NamedPipeClient>,
     buf: Vec<u8>,
-    map: Weak<DashMap<u32, oneshot::Sender<Vec<u8>>>>,
+    map: Weak<scc::HashMap<u32, oneshot::Sender<Vec<u8>>>>,
     read_task: JoinHandle<anyhow::Result<()>>,
 }
 
@@ -33,7 +32,7 @@ impl IpcClientConn {
     pub async fn new(client: NamedPipeClient) -> anyhow::Result<(Self, IpcClientEventStream)> {
         let (mut rx, tx) = split(client);
 
-        let map = Arc::new(DashMap::<u32, oneshot::Sender<Vec<u8>>>::new());
+        let map = Arc::new(scc::HashMap::<u32, oneshot::Sender<Vec<u8>>>::new());
         let (event_tx, event_rx) = mpsc::unbounded_channel();
 
         let read_task = tokio::spawn({
@@ -51,7 +50,7 @@ impl IpcClientConn {
 
                     match packet {
                         ServerToClientPacket::Response(res) => {
-                            if let Some((_, sender)) = map.remove(&res.id) {
+                            if let Some((_, sender)) = map.remove_async(&res.id).await {
                                 _ = sender.send(res.data);
                             }
                         }
@@ -129,7 +128,7 @@ impl IpcClientConn {
         .await?;
 
         let (tx, rx) = oneshot::channel();
-        map.insert(id, tx);
+        map.upsert_async(id, tx).await;
         self.tx.write_all(&self.buf).await?;
 
         self.tx.flush().await?;

crates/client/src/client.rs

@@ -17,15 +17,14 @@ asdf-overlay-client = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
 
 neon = "1"
-parking_lot = "0.12.3"
 anyhow = "1.0.97"
 once_cell = "1.21.3"
-dashmap = "6.1.0"
 rustc-hash = "2.1.1"
 mimalloc = "0.1.46"
 bytemuck = "1.22.0"
 futures = "0.3.31"
 num = "0.4.3"
+scc = "3.3.3"
 
 [dependencies.windows]
 workspace = true

packages/core/native/Cargo.toml

@@ -3,12 +3,11 @@ mod overlay;
 mod surface;
 mod util;
 
-use dashmap::DashMap;
 use mimalloc::MiMalloc;
 use neon::prelude::*;
 use rustc_hash::FxBuildHasher;
 
-type FxDashMap<K, V> = DashMap<K, V, FxBuildHasher>;
+type FxSccMap<K, V> = scc::HashMap<K, V, FxBuildHasher>;
 
 #[global_allocator]
 static GLOBAL: MiMalloc = MiMalloc;

packages/core/native/src/lib.rs

@@ -1,12 +1,14 @@
 use core::{
+    pin::pin,
     sync::atomic::{AtomicU32, Ordering},
+    task::Poll,
     time::Duration,
 };
 use std::{path::PathBuf, sync::LazyLock};
 
 use super::conv::{deserialize_percent_length, emit_event};
 use super::util::with_rt;
-use crate::{conv::deserialize_handle_update, util::runtime};
+use crate::{FxSccMap, conv::deserialize_handle_update, util::runtime};
 use anyhow::Context as AnyhowContext;
 use asdf_overlay_client::{
     OverlayDll,
@@ -18,12 +20,9 @@ use asdf_overlay_client::{
     event::OverlayEvent,
     inject,
 };
-use dashmap::DashMap;
+use futures::future::poll_fn;
 use neon::prelude::*;
 use num::FromPrimitive;
-use rustc_hash::FxBuildHasher;
-
-type FxDashMap<K, V> = DashMap<K, V, FxBuildHasher>;
 
 struct Overlay {
     ipc: IpcClientConn,
@@ -32,7 +31,7 @@ struct Overlay {
 
 struct OverlayStore {
     next_id: AtomicU32,
-    overlay_map: FxDashMap<u32, Overlay>,
+    overlay_map: FxSccMap<u32, Overlay>,
 }
 
 impl OverlayStore {
@@ -55,7 +54,9 @@ impl OverlayStore {
         .context("cannot inject to the process")?;
 
         let id = self.next_id.fetch_add(1, Ordering::AcqRel);
-        self.overlay_map.insert(id, Overlay { ipc, event });
+        self.overlay_map
+            .upsert_async(id, Overlay { ipc, event })
+            .await;
 
         Ok(id)
     }
@@ -65,20 +66,35 @@ impl OverlayStore {
         id: u32,
         f: impl AsyncFnOnce(&mut Overlay) -> R,
     ) -> anyhow::Result<R> {
-        let mut overlay = self.overlay_map.get_mut(&id).context("invalid id")?;
+        let mut overlay = self
+            .overlay_map
+            .get_async(&id)
+            .await
+            .context("invalid id")?;
         Ok(f(&mut *overlay).await)
     }
 
+    async fn next_event(&self, id: u32) -> anyhow::Result<Option<OverlayEvent>> {
+        poll_fn(|cx| {
+            let Some(mut overlay) = self.overlay_map.get_sync(&id) else {
+                return Poll::Ready(Err(anyhow::anyhow!("invalid id")));
+            };
+
+            pin!(overlay.event.recv()).poll(cx).map(Ok)
+        })
+        .await
+    }
+
     fn destroy(&self, id: u32) -> anyhow::Result<()> {
-        self.overlay_map.remove(&id).context("invalid id")?;
+        self.overlay_map.remove_sync(&id).context("invalid id")?;
 
         Ok(())
     }
 }
 
 static STORE: LazyLock<OverlayStore> = LazyLock::new(|| OverlayStore {
     next_id: AtomicU32::new(0),
-    overlay_map: FxDashMap::default(),
+    overlay_map: FxSccMap::default(),
 });
 
 pub async fn try_with_ipc<T>(
@@ -214,15 +230,8 @@ fn overlay_call_next_event(mut cx: FunctionContext) -> JsResult<JsPromise> {
 
     let (deferred, promise) = cx.promise();
     rt.spawn(async move {
-        let res = async move {
-            let event: Option<OverlayEvent> = STORE
-                .with_mut(id, async move |overlay| overlay.event.recv().await)
-                .await?;
-            Ok::<_, anyhow::Error>(event)
-        }
-        .await;
-
-        deferred.settle_with(&channel, move |mut cx| match res {
+        let event = STORE.next_event(id).await;
+        deferred.settle_with(&channel, move |mut cx| match event {
             Ok(Some(event)) => {
                 let emitter = emitter.into_inner(&mut cx);
                 let emit = emit.into_inner(&mut cx);

packages/core/native/src/overlay.rs

@@ -11,14 +11,14 @@ use neon::{
 use once_cell::sync::Lazy;
 
 use crate::{
-    FxDashMap,
+    FxSccMap,
     conv::{deserialize_copy_rect, deserialize_gpu_luid, serialize_handle_update},
     util::create_adapter_by_luid,
 };
 
 struct SurfaceStore {
     next_id: AtomicU32,
-    overlay_map: FxDashMap<u32, OverlaySurface>,
+    overlay_map: FxSccMap<u32, OverlaySurface>,
 }
 
 impl SurfaceStore {
@@ -27,7 +27,7 @@ impl SurfaceStore {
         let surface = OverlaySurface::new(adapter.as_ref())?;
 
         let id = self.next_id.fetch_add(1, Ordering::AcqRel);
-        self.overlay_map.insert(id, surface);
+        self.overlay_map.upsert_sync(id, surface);
         Ok(id)
     }
 
@@ -38,19 +38,19 @@ impl SurfaceStore {
     ) -> anyhow::Result<R> {
         let mut surface = self
             .overlay_map
-            .get_mut(&id)
+            .get_sync(&id)
             .context("Invalid surface id.")?;
         f(&mut surface)
     }
 
     fn destroy(&self, id: u32) -> bool {
-        self.overlay_map.remove(&id).is_some()
+        self.overlay_map.remove_sync(&id).is_some()
     }
 }
 
 static STORE: Lazy<SurfaceStore> = Lazy::new(|| SurfaceStore {
     next_id: AtomicU32::new(0),
-    overlay_map: FxDashMap::default(),
+    overlay_map: FxSccMap::default(),
 });
 
 fn surface_create(mut cx: FunctionContext) -> JsResult<JsNumber> {