From a9822ec80b0915c42d78d88adfe3625219254bd0 Mon Sep 17 00:00:00 2001
From: David Pedersen <david.pdrsn@gmail.com>
Date: Sun, 17 Sep 2023 18:02:22 +0200
Subject: [PATCH] Use tower-http's `TimeoutLayer` (#2231)

---
 axum/src/routing/method_routing.rs | 16 ++++++----------
 axum/src/routing/tests/merge.rs    | 10 ++++------
 axum/src/routing/tests/mod.rs      | 12 ++++++------
 3 files changed, 16 insertions(+), 22 deletions(-)

diff --git a/axum/src/routing/method_routing.rs b/axum/src/routing/method_routing.rs
index ae1b2fdca7..0feb8db953 100644
--- a/axum/src/routing/method_routing.rs
+++ b/axum/src/routing/method_routing.rs
@@ -1247,15 +1247,14 @@ const _: () = {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{
-        body::Body, error_handling::HandleErrorLayer, extract::State,
-        handler::HandlerWithoutStateExt,
-    };
+    use crate::{body::Body, extract::State, handler::HandlerWithoutStateExt};
     use axum_core::response::IntoResponse;
     use http::{header::ALLOW, HeaderMap};
     use std::time::Duration;
-    use tower::{timeout::TimeoutLayer, Service, ServiceExt};
-    use tower_http::{services::fs::ServeDir, validate_request::ValidateRequestHeaderLayer};
+    use tower::{Service, ServiceExt};
+    use tower_http::{
+        services::fs::ServeDir, timeout::TimeoutLayer, validate_request::ValidateRequestHeaderLayer,
+    };
 
     #[crate::test]
     async fn method_not_allowed_by_default() {
@@ -1354,10 +1353,7 @@ mod tests {
                 .merge(delete_service(ServeDir::new(".")))
                 .fallback(|| async { StatusCode::NOT_FOUND })
                 .put(ok)
-                .layer((
-                    HandleErrorLayer::new(|_| async { StatusCode::REQUEST_TIMEOUT }),
-                    TimeoutLayer::new(Duration::from_secs(10)),
-                )),
+                .layer(TimeoutLayer::new(Duration::from_secs(10))),
         );
 
         let listener = tokio::net::TcpListener::bind("0.0.0.0:0").await.unwrap();
diff --git a/axum/src/routing/tests/merge.rs b/axum/src/routing/tests/merge.rs
index 74eb24bea5..cc65628999 100644
--- a/axum/src/routing/tests/merge.rs
+++ b/axum/src/routing/tests/merge.rs
@@ -1,7 +1,8 @@
 use super::*;
-use crate::{error_handling::HandleErrorLayer, extract::OriginalUri, response::IntoResponse, Json};
+use crate::{extract::OriginalUri, response::IntoResponse, Json};
 use serde_json::{json, Value};
-use tower::{limit::ConcurrencyLimitLayer, timeout::TimeoutLayer};
+use tower::limit::ConcurrencyLimitLayer;
+use tower_http::timeout::TimeoutLayer;
 
 #[crate::test]
 async fn basic() {
@@ -127,10 +128,7 @@ async fn layer_and_handle_error() {
     let one = Router::new().route("/foo", get(|| async {}));
     let two = Router::new()
         .route("/timeout", get(std::future::pending::<()>))
-        .layer((
-            HandleErrorLayer::new(|_| async { StatusCode::REQUEST_TIMEOUT }),
-            TimeoutLayer::new(Duration::from_millis(10)),
-        ));
+        .layer(TimeoutLayer::new(Duration::from_millis(10)));
     let app = one.merge(two);
 
     let client = TestClient::new(app);
diff --git a/axum/src/routing/tests/mod.rs b/axum/src/routing/tests/mod.rs
index ab7c825b9f..1513cf6964 100644
--- a/axum/src/routing/tests/mod.rs
+++ b/axum/src/routing/tests/mod.rs
@@ -30,8 +30,11 @@ use std::{
     task::{Context, Poll},
     time::Duration,
 };
-use tower::{service_fn, timeout::TimeoutLayer, util::MapResponseLayer, ServiceExt};
-use tower_http::{limit::RequestBodyLimitLayer, validate_request::ValidateRequestHeaderLayer};
+use tower::{service_fn, util::MapResponseLayer, ServiceExt};
+use tower_http::{
+    limit::RequestBodyLimitLayer, timeout::TimeoutLayer,
+    validate_request::ValidateRequestHeaderLayer,
+};
 use tower_service::Service;
 
 mod fallback;
@@ -301,10 +304,7 @@ async fn wildcard_sees_whole_url() {
 async fn middleware_applies_to_routes_above() {
     let app = Router::new()
         .route("/one", get(std::future::pending::<()>))
-        .layer((
-            HandleErrorLayer::new(|_: BoxError| async move { StatusCode::REQUEST_TIMEOUT }),
-            TimeoutLayer::new(Duration::new(0, 0)),
-        ))
+        .layer(TimeoutLayer::new(Duration::ZERO))
         .route("/two", get(|| async {}));
 
     let client = TestClient::new(app);