From 79f484499d9a839ced9a9127676a643b6812dcf5 Mon Sep 17 00:00:00 2001
From: kazk <kazk.dev@gmail.com>
Date: Tue, 1 Jun 2021 02:15:24 -0700
Subject: [PATCH] Replace `LogRequest` layer with `TraceLayer`

Keeping this simple for now by default, but it's fully customizable.
---
 examples/pod_api.rs     |  2 +-
 kube/Cargo.toml         |  2 +-
 kube/src/client/mod.rs  | 19 ++++++++++++++---
 kube/src/service/log.rs | 45 -----------------------------------------
 kube/src/service/mod.rs |  2 --
 5 files changed, 18 insertions(+), 52 deletions(-)
 delete mode 100644 kube/src/service/log.rs

diff --git a/examples/pod_api.rs b/examples/pod_api.rs
index c3c48ce3c..b915dd5c9 100644
--- a/examples/pod_api.rs
+++ b/examples/pod_api.rs
@@ -10,7 +10,7 @@ use kube::{
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    std::env::set_var("RUST_LOG", "info,kube=debug");
+    std::env::set_var("RUST_LOG", "info,kube=debug,tower_http=debug");
     env_logger::init();
     let client = Client::try_default().await?;
     let namespace = std::env::var("NAMESPACE").unwrap_or("default".into());
diff --git a/kube/Cargo.toml b/kube/Cargo.toml
index d181dce02..5b4565c03 100644
--- a/kube/Cargo.toml
+++ b/kube/Cargo.toml
@@ -65,7 +65,7 @@ hyper-tls = { version = "0.5.0", optional = true }
 hyper-rustls = { version = "0.22.1", optional = true }
 tokio-tungstenite = { version = "0.14.0", optional = true }
 tower = { version = "0.4.6", optional = true, features = ["buffer", "util"] }
-tower-http = { version = "0.1.0", optional = true, features = ["map-response-body"] }
+tower-http = { version = "0.1.0", optional = true, features = ["map-response-body", "trace"] }
 hyper-timeout = {version = "0.4.1", optional = true }
 tame-oauth = { version = "0.4.7", features = ["gcp"], optional = true }
 pin-project = { version = "1.0.4", optional = true }
diff --git a/kube/src/client/mod.rs b/kube/src/client/mod.rs
index 24de5ede8..f73e7154b 100644
--- a/kube/src/client/mod.rs
+++ b/kube/src/client/mod.rs
@@ -27,12 +27,17 @@ use tokio_util::{
     io::StreamReader,
 };
 use tower::{buffer::Buffer, util::BoxService, BoxError, Layer, Service, ServiceBuilder, ServiceExt};
-use tower_http::map_response_body::MapResponseBodyLayer;
+use tower_http::{
+    map_response_body::MapResponseBodyLayer,
+    trace::{DefaultOnRequest, DefaultOnResponse, TraceLayer},
+    LatencyUnit,
+};
+use tracing::Level;
 
 use crate::{
     api::WatchEvent,
     error::{ConfigError, ErrorResponse},
-    service::{set_default_headers, AuthLayer, Authentication, LogRequest, SetBaseUriLayer},
+    service::{set_default_headers, AuthLayer, Authentication, SetBaseUriLayer},
     Config, Error, Result,
 };
 
@@ -464,7 +469,15 @@ impl TryFrom<Config> for Client {
         let inner = ServiceBuilder::new()
             .layer(common)
             .option_layer(maybe_auth)
-            .layer(tower::layer::layer_fn(LogRequest::new))
+            .layer(
+                TraceLayer::new_for_http()
+                    .on_request(DefaultOnRequest::new().level(Level::DEBUG))
+                    .on_response(
+                        DefaultOnResponse::new()
+                            .level(Level::DEBUG)
+                            .latency_unit(LatencyUnit::Millis),
+                    ),
+            )
             .service(client);
         Ok(Self::new_with_default_ns(inner, default_ns))
     }
diff --git a/kube/src/service/log.rs b/kube/src/service/log.rs
deleted file mode 100644
index 8981fa469..000000000
--- a/kube/src/service/log.rs
+++ /dev/null
@@ -1,45 +0,0 @@
-use std::task::{Context, Poll};
-
-use http::Request;
-use tower::Service;
-
-// `Clone` so that it can be composed with `AuthLayer`.
-/// Example service to log complete request before sending.
-/// Can be used to support better logging of API calls.
-/// https://github.com/clux/kube-rs/issues/26
-#[derive(Clone)]
-pub struct LogRequest<S>
-where
-    S: Clone,
-{
-    service: S,
-}
-
-impl<S> LogRequest<S>
-where
-    S: Clone,
-{
-    /// Create `LogRequest` service wrapping `service`.
-    pub fn new(service: S) -> Self {
-        Self { service }
-    }
-}
-
-impl<S, B> Service<Request<B>> for LogRequest<S>
-where
-    S: Service<Request<B>> + Clone,
-    B: http_body::Body,
-{
-    type Error = S::Error;
-    type Future = S::Future;
-    type Response = S::Response;
-
-    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        self.service.poll_ready(cx)
-    }
-
-    fn call(&mut self, req: Request<B>) -> Self::Future {
-        tracing::debug!("{} {}", req.method(), req.uri());
-        self.service.call(req)
-    }
-}
diff --git a/kube/src/service/mod.rs b/kube/src/service/mod.rs
index 8e25bae95..087796e4c 100644
--- a/kube/src/service/mod.rs
+++ b/kube/src/service/mod.rs
@@ -3,11 +3,9 @@
 mod auth;
 mod base_uri;
 mod headers;
-mod log;
 
 pub(crate) use self::{
     auth::{AuthLayer, Authentication},
     headers::set_default_headers,
-    log::LogRequest,
 };
 pub use base_uri::{SetBaseUri, SetBaseUriLayer};