Merge pull request #1120 from openmina/fix/core/channels/wasm/panic

Fix WebNode: tokio channels sometimes panic when used in wasm main browser context

Author: binier

Cargo.lock

@@ -17,6 +17,7 @@ binprot_derive = { git = "https://github.com/openmina/binprot-rs", rev = "400b52
 rand = "0.8.0"
 redux = { workspace = true }
 tokio = { version = "1.26", features = ["sync"] }
+flume = { version = "0.11.1", features = ["async", "spin"] }
 time = { version = "0.3", features = ["formatting", "macros", "parsing"] }
 md5 = "0.7.0"
 multihash = { version = "0.18.1", features = ["blake2b"] }

core/Cargo.toml

@@ -1,15 +1,17 @@
 pub use tokio::sync::oneshot;
 
 pub mod mpsc {
-    use std::{
-        sync::{Arc, Weak},
-        task::{Context, Poll},
-    };
-    use tokio::sync::mpsc::error::*;
-    pub use tokio::sync::mpsc::{self, *};
+    use std::sync::{Arc, Weak};
 
-    pub struct UnboundedSender<T>(mpsc::UnboundedSender<T>, Arc<()>);
-    pub struct UnboundedReceiver<T>(mpsc::UnboundedReceiver<T>);
+    pub use flume::{SendError, TryRecvError, TrySendError};
+
+    pub type RecvStream<T> = flume::r#async::RecvStream<'static, T>;
+
+    pub struct Sender<T>(flume::Sender<T>);
+    pub struct Receiver<T>(flume::Receiver<T>);
+
+    pub struct UnboundedSender<T>(flume::Sender<T>, Arc<()>);
+    pub struct UnboundedReceiver<T>(flume::Receiver<T>);
 
     pub type TrackedUnboundedSender<T> = UnboundedSender<Tracked<T>>;
     pub type TrackedUnboundedReceiver<T> = UnboundedReceiver<Tracked<T>>;
@@ -31,6 +33,12 @@ pub mod mpsc {
         }
     }
 
+    impl<T> Clone for Sender<T> {
+        fn clone(&self) -> Self {
+            Self(self.0.clone())
+        }
+    }
+
     impl<T> Clone for UnboundedSender<T> {
         fn clone(&self) -> Self {
             Self(self.0.clone(), self.1.clone())
@@ -51,6 +59,34 @@ pub mod mpsc {
         }
     }
 
+    impl<T> Sender<T> {
+        pub async fn send(&self, message: T) -> Result<(), SendError<T>> {
+            self.0.send_async(message).await
+        }
+
+        pub fn try_send(&self, message: T) -> Result<(), TrySendError<T>> {
+            self.0.try_send(message)
+        }
+    }
+
+    impl<T> Receiver<T> {
+        pub fn is_empty(&self) -> bool {
+            self.0.is_empty()
+        }
+
+        pub fn len(&self) -> usize {
+            self.0.len()
+        }
+
+        pub async fn recv(&mut self) -> Option<T> {
+            self.0.recv_async().await.ok()
+        }
+
+        pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
+            self.0.try_recv()
+        }
+    }
+
     impl<T> UnboundedSender<T> {
         pub fn is_empty(&self) -> bool {
             self.len() == 0
@@ -82,31 +118,37 @@ pub mod mpsc {
         }
 
         pub async fn recv(&mut self) -> Option<T> {
-            self.0.recv().await
+            self.0.recv_async().await.ok()
         }
 
         pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
             self.0.try_recv()
         }
 
-        pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> {
-            self.0.poll_recv(cx)
+        pub fn stream(&self) -> RecvStream<T> {
+            self.0.clone().into_stream()
         }
 
         pub fn blocking_recv(&mut self) -> Option<T> {
-            self.0.blocking_recv()
+            self.0.recv().ok()
         }
     }
 
+    pub fn channel<T>(bound: usize) -> (Sender<T>, Receiver<T>) {
+        let (tx, rx) = flume::bounded(bound);
+
+        (Sender(tx), Receiver(rx))
+    }
+
     pub fn unbounded_channel<T>() -> (UnboundedSender<T>, UnboundedReceiver<T>) {
-        let (tx, rx) = mpsc::unbounded_channel();
+        let (tx, rx) = flume::unbounded();
 
         (UnboundedSender(tx, Arc::new(())), UnboundedReceiver(rx))
     }
 
     pub fn tracked_unbounded_channel<T>(
     ) -> (UnboundedSender<Tracked<T>>, UnboundedReceiver<Tracked<T>>) {
-        let (tx, rx) = mpsc::unbounded_channel();
+        let (tx, rx) = flume::unbounded();
 
         (UnboundedSender(tx, Arc::new(())), UnboundedReceiver(rx))
     }
@@ -133,21 +175,21 @@ pub mod watch {
 }
 
 #[allow(dead_code)]
-pub struct Aborter(mpsc::Receiver<()>, mpsc::Sender<()>);
+pub struct Aborter(flume::Receiver<()>, flume::Sender<()>);
 
 #[derive(Clone)]
-pub struct Aborted(mpsc::Sender<()>);
+pub struct Aborted(flume::Sender<()>);
 
 impl Default for Aborter {
     fn default() -> Self {
-        let (tx, rx) = mpsc::channel(1);
+        let (tx, rx) = flume::bounded(0);
         Self(rx, tx)
     }
 }
 
 impl Aborter {
     pub fn listener_count(&self) -> usize {
-        self.1.strong_count().saturating_sub(1)
+        self.0.sender_count().saturating_sub(1)
     }
 
     /// Simply drops the object. No need to call manually, unless you
@@ -163,6 +205,7 @@ impl Aborter {
 
 impl Aborted {
     pub async fn wait(&self) {
-        self.0.closed().await;
+        // it returning an error means receiver was dropped
+        while self.0.send_async(()).await.is_ok() {}
     }
 }

core/src/channels.rs

@@ -5,6 +5,7 @@ mod p2p_task_spawner;
 
 mod node_id;
 pub use node_id::{ClusterNodeId, ClusterOcamlNodeId};
+use openmina_core::channels::Aborter;
 
 pub mod runner;
 
@@ -19,7 +20,6 @@ use libp2p::futures::{stream::FuturesUnordered, StreamExt};
 
 use ledger::proofs::provers::BlockProver;
 use node::account::{AccountPublicKey, AccountSecretKey};
-use node::core::channels::mpsc;
 use node::core::consensus::ConsensusConstants;
 use node::core::constants::constraint_constants;
 use node::core::invariants::InvariantsState;
@@ -208,7 +208,8 @@ impl Cluster {
         let node_config = testing_config.clone();
         let node_id = ClusterNodeId::new_unchecked(self.nodes.len());
         let work_dir = TempDir::new().unwrap();
-        let (shutdown_tx, shutdown_rx) = mpsc::channel(1);
+        let shutdown_initiator = Aborter::default();
+        let shutdown_listener = shutdown_initiator.aborted();
         let p2p_sec_key = match testing_config.peer_id {
             TestPeerId::Derived => P2pSecretKey::deterministic(node_id.index()),
             TestPeerId::Bytes(bytes) => P2pSecretKey::from_bytes(bytes),
@@ -306,7 +307,7 @@ impl Cluster {
             .ledger_init()
             .p2p_init_with_custom_task_spawner(
                 p2p_sec_key.clone(),
-                p2p_task_spawner::P2pTaskSpawner::new(shutdown_tx.clone()),
+                p2p_task_spawner::P2pTaskSpawner::new(shutdown_listener.clone()),
             )
             .gather_stats()
             .record(match testing_config.recorder {
@@ -331,15 +332,15 @@ impl Cluster {
             .enable_all()
             .build()
             .unwrap();
-        let shutdown = shutdown_tx.clone();
+        let shutdown = shutdown_listener.clone();
         let rpc_sender = real_service.rpc_sender();
         thread::Builder::new()
             .name("openmina_http_server".to_owned())
             .spawn(move || {
                 let local_set = tokio::task::LocalSet::new();
                 let task = async {
                     tokio::select! {
-                        _ = shutdown.closed() => {}
+                        _ = shutdown.wait() => {}
                         _ = http_server::run(http_port, rpc_sender) => {}
                     }
                 };
@@ -349,7 +350,7 @@ impl Cluster {
 
         let invariants_state = self.invariants_state.clone();
         let mut service =
-            NodeTestingService::new(real_service, node_id, invariants_state, shutdown_rx);
+            NodeTestingService::new(real_service, node_id, invariants_state, shutdown_initiator);
 
         service.set_proof_kind(self.config.proof_kind());
         if self.config.all_rust_to_rust_use_webrtc() {

node/testing/src/cluster/mod.rs

@@ -1,14 +1,14 @@
-use node::core::channels::mpsc;
 use node::core::thread;
 use node::p2p::service_impl::TaskSpawner;
+use openmina_core::channels::Aborted;
 
 #[derive(Clone)]
 pub struct P2pTaskSpawner {
-    shutdown: mpsc::Sender<()>,
+    shutdown: Aborted,
 }
 
 impl P2pTaskSpawner {
-    pub fn new(shutdown: mpsc::Sender<()>) -> Self {
+    pub fn new(shutdown: Aborted) -> Self {
         Self { shutdown }
     }
 }
@@ -28,7 +28,7 @@ impl TaskSpawner for P2pTaskSpawner {
             .spawn(move || {
                 let fut = async {
                     tokio::select! {
-                        _ = shutdown.closed() => {}
+                        _ = shutdown.wait() => {}
                         _ = fut => {}
                     }
                 };

node/testing/src/cluster/p2p_task_spawner.rs

@@ -53,6 +53,7 @@ use node::{
     },
 };
 use node::{ActionWithMeta, State};
+use openmina_core::channels::Aborter;
 use openmina_node_native::NodeService;
 use redux::Instant;
 
@@ -141,15 +142,15 @@ pub struct NodeTestingService {
 
     cluster_invariants_state: Arc<StdMutex<InvariantsState>>,
     /// Once dropped, it will cause all threads associated to shutdown.
-    _shutdown: mpsc::Receiver<()>,
+    _shutdown: Aborter,
 }
 
 impl NodeTestingService {
     pub fn new(
         real: NodeService,
         id: ClusterNodeId,
         cluster_invariants_state: Arc<StdMutex<InvariantsState>>,
-        _shutdown: mpsc::Receiver<()>,
+        _shutdown: Aborter,
     ) -> Self {
         Self {
             real,

node/testing/src/service/mod.rs

@@ -73,14 +73,14 @@ impl RustNodeConfig {
 
 pub struct RustNode {
     store: Store,
-    event_receiver: mpsc::UnboundedReceiver<P2pEvent>,
+    event_receiver: mpsc::RecvStream<P2pEvent>,
 }
 
 impl RustNode {
     pub(super) fn new(store: Store, event_receiver: mpsc::UnboundedReceiver<P2pEvent>) -> Self {
         RustNode {
             store,
-            event_receiver,
+            event_receiver: event_receiver.stream(),
         }
     }
 
@@ -109,7 +109,7 @@ impl RustNode {
     }
 
     fn poll_event_receiver(&mut self, cx: &mut Context<'_>) -> Poll<Option<RustNodeEvent>> {
-        let event = ready!(Pin::new(&mut self.event_receiver).poll_recv(cx));
+        let event = ready!(Pin::new(&mut self.event_receiver).poll_next(cx));
         Poll::Ready(event.map(|event| {
             self.dispatch_event(event.clone());
             RustNodeEvent::P2p { event }