Fix simtest determinism (MystenLabs#16029)

crates/sui-e2e-tests/tests/full_node_tests.rs

@@ -528,7 +528,16 @@ async fn test_full_node_cold_sync() -> Result<(), anyhow::Error> {
 }
 
 #[sim_test]
-async fn test_full_node_sync_flood() -> Result<(), anyhow::Error> {
+async fn test_full_node_sync_flood() {
+    do_test_full_node_sync_flood().await
+}
+
+#[sim_test(check_determinism)]
+async fn test_full_node_sync_flood_determinism() {
+    do_test_full_node_sync_flood().await
+}
+
+async fn do_test_full_node_sync_flood() {
     let mut test_cluster = TestClusterBuilder::new().build().await;
 
     // Start a new fullnode that is not on the write path
@@ -624,8 +633,6 @@ async fn test_full_node_sync_flood() -> Result<(), anyhow::Error> {
         .notify_read_executed_effects(digests)
         .await
         .unwrap();
-
-    Ok(())
 }
 
 #[sim_test]

narwhal/worker/src/client.rs

@@ -14,9 +14,39 @@ use tokio::time::{sleep, timeout, Duration};
 use tracing::info;
 use types::{Transaction, TxResponse};
 
-/// Uses a map to allow running multiple Narwhal instances in the same process.
-static LOCAL_NARWHAL_CLIENTS: Mutex<BTreeMap<Multiaddr, Arc<ArcSwap<LocalNarwhalClient>>>> =
-    Mutex::new(BTreeMap::new());
+#[cfg(msim)]
+mod static_client_cache {
+    use super::*;
+    thread_local! {
+        /// Uses a map to allow running multiple Narwhal instances in the same process.
+        static LOCAL_NARWHAL_CLIENTS: Mutex<BTreeMap<Multiaddr, Arc<ArcSwap<LocalNarwhalClient>>>> =
+            Mutex::new(BTreeMap::new());
+    }
+
+    pub(super) fn with_clients<T>(
+        f: impl FnOnce(&mut BTreeMap<Multiaddr, Arc<ArcSwap<LocalNarwhalClient>>>) -> T,
+    ) -> T {
+        LOCAL_NARWHAL_CLIENTS.with(|clients| {
+            let mut clients = clients.lock().unwrap();
+            f(&mut clients)
+        })
+    }
+}
+
+#[cfg(not(msim))]
+mod static_client_cache {
+    use super::*;
+    /// Uses a map to allow running multiple Narwhal instances in the same process.
+    static LOCAL_NARWHAL_CLIENTS: Mutex<BTreeMap<Multiaddr, Arc<ArcSwap<LocalNarwhalClient>>>> =
+        Mutex::new(BTreeMap::new());
+
+    pub(super) fn with_clients<T>(
+        f: impl FnOnce(&mut BTreeMap<Multiaddr, Arc<ArcSwap<LocalNarwhalClient>>>) -> T,
+    ) -> T {
+        let mut clients = LOCAL_NARWHAL_CLIENTS.lock().unwrap();
+        f(&mut clients)
+    }
+}
 
 /// The maximum allowed size of transactions into Narwhal.
 /// TODO: maybe move to TxValidator?
@@ -96,23 +126,23 @@ impl LocalNarwhalClient {
     pub fn set_global(addr: Multiaddr, instance: Arc<Self>) {
         info!("Narwhal worker client added ({})", addr);
         let addr = Self::canonicalize_address_key(addr);
-        let mut clients = LOCAL_NARWHAL_CLIENTS.lock().unwrap();
-        match clients.entry(addr) {
-            Entry::Vacant(entry) => {
-                entry.insert(Arc::new(ArcSwap::from(instance)));
-            }
-            Entry::Occupied(mut entry) => {
-                entry.get_mut().store(instance);
-            }
-        };
+        static_client_cache::with_clients(|clients| {
+            match clients.entry(addr) {
+                Entry::Vacant(entry) => {
+                    entry.insert(Arc::new(ArcSwap::from(instance)));
+                }
+                Entry::Occupied(mut entry) => {
+                    entry.get_mut().store(instance);
+                }
+            };
+        });
     }
 
     /// Gets the instance of LocalNarwhalClient for the local address.
     /// Address is only used as the key.
     pub fn get_global(addr: &Multiaddr) -> Option<Arc<ArcSwap<Self>>> {
         let addr = Self::canonicalize_address_key(addr.clone());
-        let clients = LOCAL_NARWHAL_CLIENTS.lock().unwrap();
-        clients.get(&addr).cloned()
+        static_client_cache::with_clients(|clients| clients.get(&addr).cloned())
     }
 
     /// Submits a transaction to the local Narwhal worker.