feat(meta): track the time spent on meta client methods

src/common/base/src/future.rs

@@ -0,0 +1,220 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::future::Future;
+use std::marker::PhantomData;
+use std::pin::Pin;
+use std::task::Context;
+use std::task::Poll;
+use std::time::Duration;
+
+use pin_project_lite::pin_project;
+use tokio::time::Instant;
+
+pin_project! {
+    /// A [`Future`] that tracks the time spent on a future.
+    /// When the future is ready, the callback will be called with the total time and busy time.
+    #[must_use = "futures do nothing unless you `.await` or poll them"]
+    pub struct TimingFuture<'a, T, F>
+    where F: Fn(Duration, Duration) ,
+    F: 'a,
+    {
+        #[pin]
+        inner: T,
+
+        start: Option<Instant>,
+        busy: Duration,
+        callback: F,
+        _p : PhantomData<&'a ()>,
+
+    }
+}
+
+impl<'a, T, F> TimingFuture<'a, T, F>
+where
+    F: Fn(Duration, Duration),
+    F: 'a,
+{
+    pub fn new(inner: T, callback: F) -> Self {
+        Self {
+            inner,
+            start: None,
+            busy: Duration::default(),
+            callback,
+            _p: Default::default(),
+        }
+    }
+}
+
+impl<'a, T: Future, F> Future for TimingFuture<'a, T, F>
+where
+    F: Fn(Duration, Duration),
+    F: 'a,
+{
+    type Output = T::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+
+        if this.start.is_none() {
+            *this.start = Some(Instant::now());
+        }
+
+        let t0 = Instant::now();
+
+        let res = this.inner.poll(cx);
+
+        *this.busy += t0.elapsed();
+
+        if res.is_ready() {
+            let total = this.start.unwrap().elapsed();
+            (this.callback)(total, *this.busy);
+        }
+
+        res
+    }
+}
+
+/// Enable timing for a future with `fu.timed(f)`.
+pub trait TimingFutureExt {
+    /// Wrap the future with a timing future.
+    fn timed<'a, F>(self, f: F) -> TimingFuture<'a, Self, F>
+    where
+        F: Fn(Duration, Duration) + 'a,
+        Self: Future + Sized;
+
+    /// Wrap the future with a timing future, and call the callback if the total time is greater than the threshold.
+    fn timed_ge<'a, F>(
+        self,
+        threshold: Duration,
+        f: F,
+    ) -> TimingFuture<'a, Self, impl Fn(Duration, Duration)>
+    where
+        F: Fn(Duration, Duration) + 'a,
+        Self: Future + Sized,
+    {
+        self.timed::<'a>(move |total, busy| {
+            if total >= threshold {
+                f(total, busy)
+            }
+        })
+    }
+}
+
+impl<T> TimingFutureExt for T
+where T: Future + Sized
+{
+    fn timed<'a, F>(self, f: F) -> TimingFuture<'a, Self, F>
+    where
+        F: Fn(Duration, Duration),
+        F: 'a,
+    {
+        TimingFuture::new(self, f)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::future::Future;
+    use std::pin::Pin;
+    use std::task::Context;
+    use std::task::Poll;
+    use std::time::Duration;
+
+    use crate::future::TimingFuture;
+    use crate::future::TimingFutureExt;
+
+    struct BlockingSleep20ms {}
+
+    impl Future for BlockingSleep20ms {
+        type Output = ();
+
+        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
+            std::thread::sleep(Duration::from_millis(20));
+            Poll::Ready(())
+        }
+    }
+
+    #[test]
+    fn test_timing_future() -> anyhow::Result<()> {
+        let rt = tokio::runtime::Builder::new_multi_thread()
+            .worker_threads(1)
+            .enable_all()
+            .build()
+            .unwrap();
+
+        // Blocking sleep
+
+        let f = BlockingSleep20ms {};
+        let f = TimingFuture::new(f, |total, busy| {
+            // println!("total: {:?}, busy: {:?}", total, busy);
+            assert!(total >= Duration::from_millis(20));
+            assert!(total <= Duration::from_millis(50));
+
+            assert!(busy >= Duration::from_millis(20));
+            assert!(busy <= Duration::from_millis(50));
+        });
+
+        rt.block_on(f);
+
+        // Async sleep
+
+        let f = async move { tokio::time::sleep(Duration::from_millis(20)).await };
+        let f = TimingFuture::new(f, |total, busy| {
+            // println!("total: {:?}, busy: {:?}", total, busy);
+            assert!(total >= Duration::from_millis(20));
+            assert!(total <= Duration::from_millis(50));
+
+            assert!(busy <= Duration::from_millis(10));
+        });
+
+        rt.block_on(f);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_time_future_ext() -> anyhow::Result<()> {
+        let rt = tokio::runtime::Builder::new_multi_thread()
+            .worker_threads(1)
+            .enable_all()
+            .build()
+            .unwrap();
+
+        // Blocking sleep
+
+        let f = BlockingSleep20ms {}.timed(|total, busy| {
+            assert!(total >= Duration::from_millis(20));
+            assert!(total <= Duration::from_millis(50));
+
+            assert!(busy >= Duration::from_millis(20));
+            assert!(busy <= Duration::from_millis(50));
+        });
+
+        rt.block_on(f);
+
+        rt.block_on(
+            BlockingSleep20ms {}.timed_ge(Duration::from_millis(10), |_total, _busy| {
+                // OK, triggered
+            }),
+        );
+        rt.block_on(
+            BlockingSleep20ms {}.timed_ge(Duration::from_millis(100), |_total, _busy| {
+                unreachable!("should not be called")
+            }),
+        );
+
+        Ok(())
+    }
+}

src/common/base/src/lib.rs

@@ -22,9 +22,11 @@
 #![feature(thread_id_value)]
 #![feature(backtrace_frames)]
 #![feature(alloc_error_hook)]
+#![feature(return_position_impl_trait_in_trait)]
 
 pub mod base;
 pub mod containers;
+pub mod future;
 pub mod mem_allocator;
 pub mod rangemap;
 pub mod runtime;

src/meta/client/src/grpc_client.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::fmt::Debug;
+use std::fmt::Display;
 use std::fmt::Formatter;
 use std::sync::Arc;
 use std::time::Duration;
@@ -30,6 +31,7 @@ use common_base::base::tokio::time::sleep;
 use common_base::containers::ItemManager;
 use common_base::containers::Pool;
 use common_base::containers::TtlHashMap;
+use common_base::future::TimingFutureExt;
 use common_base::runtime::Runtime;
 use common_base::runtime::TrySpawn;
 use common_base::runtime::UnlimitedFuture;
@@ -393,23 +395,38 @@ impl MetaGrpcClient {
 
             let resp = match req {
                 message::Request::Get(r) => {
-                    let resp = self.kv_api(r).await;
+                    let resp = self
+                        .kv_api(r)
+                        .timed_ge(threshold(), info_spent("MetaGrpcClient::kv_api"))
+                        .await;
                     message::Response::Get(resp)
                 }
                 message::Request::MGet(r) => {
-                    let resp = self.kv_api(r).await;
+                    let resp = self
+                        .kv_api(r)
+                        .timed_ge(threshold(), info_spent("MetaGrpcClient::kv_api"))
+                        .await;
                     message::Response::MGet(resp)
                 }
                 message::Request::PrefixList(r) => {
-                    let resp = self.kv_api(r).await;
+                    let resp = self
+                        .kv_api(r)
+                        .timed_ge(threshold(), info_spent("MetaGrpcClient::kv_api"))
+                        .await;
                     message::Response::PrefixList(resp)
                 }
                 message::Request::Upsert(r) => {
-                    let resp = self.kv_api(r).await;
+                    let resp = self
+                        .kv_api(r)
+                        .timed_ge(threshold(), info_spent("MetaGrpcClient::kv_api"))
+                        .await;
                     message::Response::Upsert(resp)
                 }
                 message::Request::Txn(r) => {
-                    let resp = self.transaction(r).await;
+                    let resp = self
+                        .transaction(r)
+                        .timed_ge(threshold(), info_spent("MetaGrpcClient::transaction"))
+                        .await;
                     message::Response::Txn(resp)
                 }
                 message::Request::Watch(r) => {
@@ -543,6 +560,7 @@ impl MetaGrpcClient {
                         &self.password,
                     )
                     .await;
+
                     match new_token {
                         Ok(token) => {
                             let client =
@@ -882,8 +900,15 @@ impl MetaGrpcClient {
 
         let req = common_tracing::inject_span_to_tonic_request(req);
 
-        let mut client = self.make_client().await?;
-        let result = client.kv_api(req).await;
+        let mut client = self
+            .make_client()
+            .timed_ge(threshold(), info_spent("MetaGrpcClient::make_client-1"))
+            .await?;
+
+        let result = client
+            .kv_api(req)
+            .timed_ge(threshold(), info_spent("client::kv_api-1"))
+            .await;
 
         debug!(reply = debug(&result), "MetaGrpcClient::kv_api reply");
 
@@ -892,15 +917,22 @@ impl MetaGrpcClient {
             Err(s) => {
                 if status_is_retryable(&s) {
                     self.mark_as_unhealthy().await;
-                    let mut client = self.make_client().await?;
+                    let mut client = self
+                        .make_client()
+                        .timed_ge(threshold(), info_spent("MetaGrpcClient::make_client-2"))
+                        .await?;
                     let req: Request<RaftRequest> = read_req.try_into().map_err(|e| {
                         MetaNetworkError::InvalidArgument(InvalidArgument::new(
                             e,
                             "fail to encode request",
                         ))
                     })?;
                     let req = common_tracing::inject_span_to_tonic_request(req);
-                    Ok(client.kv_api(req).await?.into_inner())
+                    Ok(client
+                        .kv_api(req)
+                        .timed_ge(threshold(), info_spent("client::kv_api-2"))
+                        .await?
+                        .into_inner())
                 } else {
                     Err(s)
                 }
@@ -972,3 +1004,17 @@ impl Interceptor for AuthInterceptor {
         Ok(req)
     }
 }
+
+fn threshold() -> Duration {
+    Duration::from_millis(300)
+}
+
+fn info_spent(msg: impl Display) -> impl Fn(Duration, Duration) {
+    move |total, busy| {
+        info!(
+            total = ?total,
+            busy = ?busy,
+            "{} spent", msg
+        );
+    }
+}