m

MystenLabs · Jul 25, 2024 · c626117 · c626117
1 parent 89c1e36
commit c626117
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 39 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/consensus/core/src/network/tonic_network.rs b/consensus/core/src/network/tonic_network.rs
@@ -895,10 +895,13 @@ impl<S: NetworkService> NetworkManager<S> for TonicManager {
                             },
                             _ = shutdown_notif.wait(), if !has_shutdown => {
                                 trace!("Received shutdown. Stopping connection for {peer_addr:?}");
-                                connection.as_mut().graceful_shutdown();
                                 has_shutdown = true;
                             },
                         }
+
+                        if has_shutdown {
+                            break;
+                        }
                     }
 
                     Ok(())

diff --git a/crates/sui-storage/Cargo.toml b/crates/sui-storage/Cargo.toml
@@ -55,6 +55,7 @@ move-binary-format.workspace = true
 move-bytecode-utils.workspace = true
 
 [dev-dependencies]
+axum.workspace = true
 anyhow.workspace = true
 criterion.workspace = true
 tempfile.workspace = true

diff --git a/crates/sui-storage/tests/key_value_tests.rs b/crates/sui-storage/tests/key_value_tests.rs
@@ -428,13 +428,9 @@ async fn test_get_tx_from_fallback() {
 
 #[cfg(msim)]
 mod simtests {
-
     use super::*;
-    use hyper::{
-        service::{make_service_fn, service_fn},
-        Body, Request, Response, Server,
-    };
-    use std::convert::Infallible;
+    use axum::routing::get;
+    use axum::{body::Body, extract::Request, extract::State, response::Response};
     use std::net::SocketAddr;
     use std::sync::Mutex;
     use std::time::{Duration, Instant};
@@ -443,6 +439,23 @@ mod simtests {
     use sui_storage::http_key_value_store::*;
     use tracing::info;
 
+    async fn svc(
+        State(state): State<Arc<Mutex<HashMap<String, Vec<u8>>>>>,
+        request: Request<Body>,
+    ) -> Response {
+        let path = request.uri().path().to_string();
+        let key = path.trim_start_matches('/');
+        let value = state.lock().unwrap().get(key).cloned();
+        info!("Got request for key: {:?}, value: {:?}", key, value);
+        match value {
+            Some(v) => Response::new(Body::from(v)),
+            None => Response::builder()
+                .status(hyper::StatusCode::NOT_FOUND)
+                .body(Body::empty())
+                .unwrap(),
+        }
+    }
+
     async fn test_server(data: Arc<Mutex<HashMap<String, Vec<u8>>>>) {
         let handle = sui_simulator::runtime::Handle::current();
         let builder = handle.create_node();
@@ -456,41 +469,12 @@ mod simtests {
                 let data = data.clone();
                 let startup_sender = startup_sender.clone();
                 async move {
-                    let make_svc = make_service_fn(move |_| {
-                        let data = data.clone();
-                        async {
-                            Ok::<_, Infallible>(service_fn(move |req: Request<Body>| {
-                                let data = data.clone();
-                                async move {
-                                    let path = req.uri().path().to_string();
-                                    let key = path.trim_start_matches('/');
-                                    let value = data.lock().unwrap().get(key).cloned();
-                                    info!("Got request for key: {:?}, value: {:?}", key, value);
-                                    match value {
-                                        Some(v) => {
-                                            Ok::<_, Infallible>(Response::new(Body::from(v)))
-                                        }
-                                        None => Ok::<_, Infallible>(
-                                            Response::builder()
-                                                .status(hyper::StatusCode::NOT_FOUND)
-                                                .body(Body::empty())
-                                                .unwrap(),
-                                        ),
-                                    }
-                                }
-                            }))
-                        }
-                    });
-
+                    let router = get(svc).with_state(data);
                     let addr = SocketAddr::from(([10, 10, 10, 10], 8080));
-                    let server = Server::bind(&addr).serve(make_svc);
-
-                    let graceful = server.with_graceful_shutdown(async {
-                        tokio::time::sleep(Duration::from_secs(86400)).await;
-                    });
+                    let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
 
                     tokio::spawn(async {
-                        let _ = graceful.await;
+                        axum::serve(listener, router).await.unwrap();
                     });
 
                     startup_sender.send(true).ok();