From 37341529cafdc424f011348ab01eadebc2dff6c5 Mon Sep 17 00:00:00 2001 From: Jonas Platte Date: Sun, 23 Jun 2024 01:34:49 +0200 Subject: [PATCH 1/5] Remove duplication in serving with and without graceful shutdown --- axum/src/serve.rs | 57 ++--------------------------------------------- 1 file changed, 2 insertions(+), 55 deletions(-) diff --git a/axum/src/serve.rs b/axum/src/serve.rs index e2b974cfbe..1ba9a1452c 100644 --- a/axum/src/serve.rs +++ b/axum/src/serve.rs @@ -213,61 +213,8 @@ where type IntoFuture = private::ServeFuture; fn into_future(self) -> Self::IntoFuture { - private::ServeFuture(Box::pin(async move { - let Self { - tcp_listener, - mut make_service, - tcp_nodelay, - _marker: _, - } = self; - - loop { - let (tcp_stream, remote_addr) = match tcp_accept(&tcp_listener).await { - Some(conn) => conn, - None => continue, - }; - - if let Some(nodelay) = tcp_nodelay { - if let Err(err) = tcp_stream.set_nodelay(nodelay) { - trace!("failed to set TCP_NODELAY on incoming connection: {err:#}"); - } - } - - let tcp_stream = TokioIo::new(tcp_stream); - - poll_fn(|cx| make_service.poll_ready(cx)) - .await - .unwrap_or_else(|err| match err {}); - - let tower_service = make_service - .call(IncomingStream { - tcp_stream: &tcp_stream, - remote_addr, - }) - .await - .unwrap_or_else(|err| match err {}) - .map_request(|req: Request| req.map(Body::new)); - - let hyper_service = TowerToHyperService::new(tower_service); - - tokio::spawn(async move { - match Builder::new(TokioExecutor::new()) - // upgrades needed for websockets - .serve_connection_with_upgrades(tcp_stream, hyper_service) - .await - { - Ok(()) => {} - Err(_err) => { - // This error only appears when the client doesn't send a request and - // terminate the connection. - // - // If client sends one request then terminate connection whenever, it doesn't - // appear. - } - } - }); - } - })) + self.with_graceful_shutdown(std::future::pending()) + .into_future() } } From dacce049bd559976c418d9b1938db599b7245d66 Mon Sep 17 00:00:00 2001 From: Jonas Platte Date: Fri, 27 Sep 2024 22:28:13 +0200 Subject: [PATCH 2/5] Add a named IntoFuture type for capture_tracing --- axum/src/routing/tests/mod.rs | 2 +- axum/src/test_helpers/tracing_helpers.rs | 79 +++++++++++++++--------- 2 files changed, 52 insertions(+), 29 deletions(-) diff --git a/axum/src/routing/tests/mod.rs b/axum/src/routing/tests/mod.rs index 144c870dfa..20a5fc0bb8 100644 --- a/axum/src/routing/tests/mod.rs +++ b/axum/src/routing/tests/mod.rs @@ -966,7 +966,7 @@ async fn logging_rejections() { rejection_type: String, } - let events = capture_tracing::(|| async { + let events = capture_tracing::(|| async { let app = Router::new() .route("/extension", get(|_: Extension| async {})) .route("/string", post(|_: String| async {})); diff --git a/axum/src/test_helpers/tracing_helpers.rs b/axum/src/test_helpers/tracing_helpers.rs index 2240717ee4..62e22af2aa 100644 --- a/axum/src/test_helpers/tracing_helpers.rs +++ b/axum/src/test_helpers/tracing_helpers.rs @@ -1,5 +1,11 @@ use crate::util::AxumMutex; -use std::{future::Future, io, sync::Arc}; +use std::{ + future::{Future, IntoFuture}, + io, + marker::PhantomData, + pin::Pin, + sync::Arc, +}; use serde::{de::DeserializeOwned, Deserialize}; use tracing_subscriber::prelude::*; @@ -14,36 +20,53 @@ pub(crate) struct TracingEvent { } /// Run an async closure and capture the tracing output it produces. -pub(crate) async fn capture_tracing(f: F) -> Vec> +pub(crate) fn capture_tracing(f: F) -> CaptureTracing where - F: Fn() -> Fut, - Fut: Future, T: DeserializeOwned, { - let (make_writer, handle) = TestMakeWriter::new(); - - let subscriber = tracing_subscriber::registry().with( - tracing_subscriber::fmt::layer() - .with_writer(make_writer) - .with_target(true) - .without_time() - .with_ansi(false) - .json() - .flatten_event(false) - .with_filter("axum=trace".parse::().unwrap()), - ); - - let guard = tracing::subscriber::set_default(subscriber); - - f().await; - - drop(guard); - - handle - .take() - .lines() - .map(|line| serde_json::from_str(line).unwrap()) - .collect() + CaptureTracing(f, PhantomData) +} + +pub(crate) struct CaptureTracing(F, PhantomData T>); + +impl IntoFuture for CaptureTracing +where + F: Fn() -> Fut + Send + 'static, + Fut: Future + Send, + T: DeserializeOwned, +{ + type Output = Vec>; + type IntoFuture = Pin + Send>>; + + fn into_future(self) -> Self::IntoFuture { + let Self(f, _) = self; + Box::pin(async move { + let (make_writer, handle) = TestMakeWriter::new(); + + let subscriber = tracing_subscriber::registry().with( + tracing_subscriber::fmt::layer() + .with_writer(make_writer) + .with_target(true) + .without_time() + .with_ansi(false) + .json() + .flatten_event(false) + .with_filter("axum=trace".parse::().unwrap()), + ); + + let guard = tracing::subscriber::set_default(subscriber); + + f().await; + + drop(guard); + + handle + .take() + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect() + }) + } } struct TestMakeWriter { From 85d14bd7d86bbfbe2b3e5f1599e4c6b60e568bd8 Mon Sep 17 00:00:00 2001 From: Jonas Platte Date: Fri, 27 Sep 2024 22:29:50 +0200 Subject: [PATCH 3/5] Make CaptureTracing work more reliably The previous code would have not worked with the inner future migrating to a different OS thread within a work-stealing async executor. --- axum/src/test_helpers/tracing_helpers.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/axum/src/test_helpers/tracing_helpers.rs b/axum/src/test_helpers/tracing_helpers.rs index 62e22af2aa..16b8a82193 100644 --- a/axum/src/test_helpers/tracing_helpers.rs +++ b/axum/src/test_helpers/tracing_helpers.rs @@ -8,6 +8,7 @@ use std::{ }; use serde::{de::DeserializeOwned, Deserialize}; +use tracing::instrument::WithSubscriber; use tracing_subscriber::prelude::*; use tracing_subscriber::{filter::Targets, fmt::MakeWriter}; @@ -56,7 +57,7 @@ where let guard = tracing::subscriber::set_default(subscriber); - f().await; + f().with_current_subscriber().await; drop(guard); From 4882a2222da49b8d873a0a950e5a7a474fb6ea76 Mon Sep 17 00:00:00 2001 From: Jonas Platte Date: Fri, 27 Sep 2024 22:34:31 +0200 Subject: [PATCH 4/5] Only read rejection traces in logging_rejections test --- axum/src/routing/tests/mod.rs | 1 + axum/src/test_helpers/tracing_helpers.rs | 24 ++++++++++++++++++++---- 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/axum/src/routing/tests/mod.rs b/axum/src/routing/tests/mod.rs index 20a5fc0bb8..e9ab56e11c 100644 --- a/axum/src/routing/tests/mod.rs +++ b/axum/src/routing/tests/mod.rs @@ -987,6 +987,7 @@ async fn logging_rejections() { StatusCode::BAD_REQUEST, ); }) + .with_filter("axum::rejection=trace") .await; assert_eq!( diff --git a/axum/src/test_helpers/tracing_helpers.rs b/axum/src/test_helpers/tracing_helpers.rs index 16b8a82193..0becb7948c 100644 --- a/axum/src/test_helpers/tracing_helpers.rs +++ b/axum/src/test_helpers/tracing_helpers.rs @@ -25,10 +25,25 @@ pub(crate) fn capture_tracing(f: F) -> CaptureTracing where T: DeserializeOwned, { - CaptureTracing(f, PhantomData) + CaptureTracing { + f, + filter: None, + _phantom: PhantomData, + } +} + +pub(crate) struct CaptureTracing { + f: F, + filter: Option, + _phantom: PhantomData T>, } -pub(crate) struct CaptureTracing(F, PhantomData T>); +impl CaptureTracing { + pub(crate) fn with_filter(mut self, filter_string: &str) -> Self { + self.filter = Some(filter_string.parse().unwrap()); + self + } +} impl IntoFuture for CaptureTracing where @@ -40,10 +55,11 @@ where type IntoFuture = Pin + Send>>; fn into_future(self) -> Self::IntoFuture { - let Self(f, _) = self; + let Self { f, filter, .. } = self; Box::pin(async move { let (make_writer, handle) = TestMakeWriter::new(); + let filter = filter.unwrap_or_else(|| "axum=trace".parse().unwrap()); let subscriber = tracing_subscriber::registry().with( tracing_subscriber::fmt::layer() .with_writer(make_writer) @@ -52,7 +68,7 @@ where .with_ansi(false) .json() .flatten_event(false) - .with_filter("axum=trace".parse::().unwrap()), + .with_filter(filter), ); let guard = tracing::subscriber::set_default(subscriber); From 4c71b654dfb0ae5250308aa98e352ae03a8077f2 Mon Sep 17 00:00:00 2001 From: Jonas Platte Date: Fri, 27 Sep 2024 23:01:51 +0200 Subject: [PATCH 5/5] Try to fix CI for Rust 1.70 --- axum/src/test_helpers/tracing_helpers.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/axum/src/test_helpers/tracing_helpers.rs b/axum/src/test_helpers/tracing_helpers.rs index 0becb7948c..96abe5102e 100644 --- a/axum/src/test_helpers/tracing_helpers.rs +++ b/axum/src/test_helpers/tracing_helpers.rs @@ -47,7 +47,7 @@ impl CaptureTracing { impl IntoFuture for CaptureTracing where - F: Fn() -> Fut + Send + 'static, + F: Fn() -> Fut + Send + Sync + 'static, Fut: Future + Send, T: DeserializeOwned, {