From c62611788b466f7abe574858cc081faedba04e1c Mon Sep 17 00:00:00 2001 From: Brandon Williams Date: Wed, 24 Jul 2024 20:30:42 -0500 Subject: [PATCH] m --- Cargo.lock | 1 + consensus/core/src/network/tonic_network.rs | 5 +- crates/sui-storage/Cargo.toml | 1 + crates/sui-storage/tests/key_value_tests.rs | 60 ++++++++------------- 4 files changed, 28 insertions(+), 39 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7677dcc65eb04b..4bdcfed4d98805 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -14405,6 +14405,7 @@ version = "0.1.0" dependencies = [ "anyhow", "async-trait", + "axum 0.7.5", "backoff", "base64-url", "bcs", diff --git a/consensus/core/src/network/tonic_network.rs b/consensus/core/src/network/tonic_network.rs index 9b4ca009d52af1..486cd54c210a4e 100644 --- a/consensus/core/src/network/tonic_network.rs +++ b/consensus/core/src/network/tonic_network.rs @@ -895,10 +895,13 @@ impl NetworkManager 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 index 77a34512daf46c..9536a8d64746a2 100644 --- 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 index 8d3fb1b051917c..5e2957c290d148 100644 --- 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>>>>, + request: Request, + ) -> 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>>>) { 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| { - 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();