[WIP] *: trace scan requests

Cargo.toml

@@ -32,7 +32,8 @@ either = "1.6"
 fail = "0.4"
 futures = { version = "0.3" }
 lazy_static = "1"
-log = "0.4"
+log = { version = "0.4", features = ["kv_unstable"] }
+minitrace = "0.6.2"
 pin-project = "1"
 prometheus = { version = "0.13", default-features = false }
 prost = "0.12"

src/kv/mod.rs

@@ -14,7 +14,7 @@ pub use key::Key;
 pub use kvpair::KvPair;
 pub use value::Value;
 
-struct HexRepr<'a>(pub &'a [u8]);
+pub struct HexRepr<'a>(pub &'a [u8]);
 
 impl<'a> fmt::Display for HexRepr<'a> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {

src/pd/cluster.rs

@@ -103,6 +103,7 @@ impl Connection {
         Connection { security_mgr }
     }
 
+    #[minitrace::trace]
     pub async fn connect_cluster(
         &self,
         endpoints: &[String],
@@ -122,6 +123,7 @@ impl Connection {
     }
 
     // Re-establish connection with PD leader in asynchronous fashion.
+    #[minitrace::trace]
     pub async fn reconnect(&self, cluster: &mut Cluster, timeout: Duration) -> Result<()> {
         warn!("updating pd client");
         let start = Instant::now();

src/pd/retry.rs

@@ -8,6 +8,8 @@ use std::time::Duration;
 use std::time::Instant;
 
 use async_trait::async_trait;
+use log::debug;
+use minitrace::prelude::*;
 use tokio::sync::RwLock;
 use tokio::time::sleep;
 
@@ -74,14 +76,17 @@ macro_rules! retry_core {
     ($self: ident, $tag: literal, $call: expr) => {{
         let stats = pd_stats($tag);
         let mut last_err = Ok(());
-        for _ in 0..LEADER_CHANGE_RETRY {
+        for retry in 0..LEADER_CHANGE_RETRY {
+            let _span = LocalSpan::enter_with_local_parent("RetryClient::retry");
+
             let res = $call;
 
             match stats.done(res) {
                 Ok(r) => return Ok(r),
                 Err(e) => last_err = Err(e),
             }
 
+            debug!("retry {} on last_err: {:?}", retry, last_err);
             let mut reconnect_count = MAX_REQUEST_COUNT;
             while let Err(e) = $self.reconnect(RECONNECT_INTERVAL_SEC).await {
                 reconnect_count -= 1;
@@ -142,6 +147,7 @@ impl RetryClient<Cluster> {
 impl RetryClientTrait for RetryClient<Cluster> {
     // These get_* functions will try multiple times to make a request, reconnecting as necessary.
     // It does not know about encoding. Caller should take care of it.
+    #[minitrace::trace]
     async fn get_region(self: Arc<Self>, key: Vec<u8>) -> Result<RegionWithLeader> {
         retry_mut!(self, "get_region", |cluster| {
             let key = key.clone();
@@ -156,6 +162,7 @@ impl RetryClientTrait for RetryClient<Cluster> {
         })
     }
 
+    #[minitrace::trace]
     async fn get_region_by_id(self: Arc<Self>, region_id: RegionId) -> Result<RegionWithLeader> {
         retry_mut!(self, "get_region_by_id", |cluster| async {
             cluster
@@ -167,6 +174,7 @@ impl RetryClientTrait for RetryClient<Cluster> {
         })
     }
 
+    #[minitrace::trace]
     async fn get_store(self: Arc<Self>, id: StoreId) -> Result<metapb::Store> {
         retry_mut!(self, "get_store", |cluster| async {
             cluster
@@ -176,6 +184,7 @@ impl RetryClientTrait for RetryClient<Cluster> {
         })
     }
 
+    #[minitrace::trace]
     async fn get_all_stores(self: Arc<Self>) -> Result<Vec<metapb::Store>> {
         retry_mut!(self, "get_all_stores", |cluster| async {
             cluster
@@ -185,10 +194,12 @@ impl RetryClientTrait for RetryClient<Cluster> {
         })
     }
 
+    #[minitrace::trace]
     async fn get_timestamp(self: Arc<Self>) -> Result<Timestamp> {
         retry!(self, "get_timestamp", |cluster| cluster.get_timestamp())
     }
 
+    #[minitrace::trace]
     async fn update_safepoint(self: Arc<Self>, safepoint: u64) -> Result<bool> {
         retry_mut!(self, "update_gc_safepoint", |cluster| async {
             cluster
@@ -277,7 +288,7 @@ mod test {
         }
 
         async fn retry_ok(client: Arc<MockClient>) -> Result<()> {
-            retry!(client, "test", |_c| ready(Ok::<_, Error>(())))
+            retry_mut!(client, "test", |_c| ready(Ok::<_, Error>(())))
         }
 
         executor::block_on(async {
@@ -342,7 +353,7 @@ mod test {
             client: Arc<MockClient>,
             max_retries: Arc<AtomicUsize>,
         ) -> Result<()> {
-            retry!(client, "test", |c| {
+            retry_mut!(client, "test", |c| {
                 c.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
 
                 let max_retries = max_retries.fetch_sub(1, Ordering::SeqCst) - 1;

src/pd/timestamp.rs

@@ -22,6 +22,7 @@ use futures::task::Context;
 use futures::task::Poll;
 use log::debug;
 use log::info;
+use minitrace::prelude::*;
 use pin_project::pin_project;
 use tokio::sync::mpsc;
 use tokio::sync::oneshot;
@@ -63,6 +64,7 @@ impl TimestampOracle {
         Ok(TimestampOracle { request_tx })
     }
 
+    #[minitrace::trace]
     pub(crate) async fn get_timestamp(self) -> Result<Timestamp> {
         debug!("getting current timestamp");
         let (request, response) = oneshot::channel();
@@ -74,6 +76,7 @@ impl TimestampOracle {
     }
 }
 
+#[minitrace::trace]
 async fn run_tso(
     cluster_id: u64,
     mut pd_client: PdClient<Channel>,
@@ -98,6 +101,9 @@ async fn run_tso(
     let mut responses = pd_client.tso(request_stream).await?.into_inner();
 
     while let Some(Ok(resp)) = responses.next().await {
+        let _span = LocalSpan::enter_with_local_parent("handle_response");
+        debug!("got response: {:?}", resp);
+
         {
             let mut pending_requests = pending_requests.lock().await;
             allocate_timestamps(&resp, &mut pending_requests)?;
@@ -128,6 +134,7 @@ struct TsoRequestStream {
 impl Stream for TsoRequestStream {
     type Item = TsoRequest;
 
+    #[minitrace::trace]
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
         let mut this = self.project();
 
@@ -152,6 +159,12 @@ impl Stream for TsoRequestStream {
             }
         }
 
+        debug!(
+            "got requests: len {}, pending_requests {}",
+            requests.len(),
+            pending_requests.len()
+        );
+
         if !requests.is_empty() {
             let req = TsoRequest {
                 header: Some(RequestHeader {
@@ -168,6 +181,12 @@ impl Stream for TsoRequestStream {
             };
             pending_requests.push_back(request_group);
 
+            debug!(
+                "sending request to PD: {:?}, pending_requests {}",
+                req,
+                pending_requests.len()
+            );
+
             Poll::Ready(Some(req))
         } else {
             // Set the waker to the context, then the stream can be waked up after the pending queue

src/raw/requests.rs

@@ -1,6 +1,7 @@
 // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
 
 use std::any::Any;
+use std::fmt::Formatter;
 use std::ops::Range;
 use std::sync::Arc;
 use std::time::Duration;
@@ -404,6 +405,14 @@ impl Request for RawCoprocessorRequest {
     }
 }
 
+impl std::fmt::Debug for RawCoprocessorRequest {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("RawCoprocessorRequest")
+            .field("inner", &self.inner)
+            .finish()
+    }
+}
+
 impl KvRequest for RawCoprocessorRequest {
     type Response = kvrpcpb::RawCoprocessorResponse;
 }

src/request/mod.rs

@@ -129,7 +129,7 @@ mod test {
 
         impl HasLocks for MockRpcResponse {}
 
-        #[derive(Clone)]
+        #[derive(Debug, Clone)]
         struct MockKvRequest {
             test_invoking_count: Arc<AtomicUsize>,
         }

src/request/plan.rs

@@ -9,6 +9,8 @@ use futures::future::try_join_all;
 use futures::prelude::*;
 use log::debug;
 use log::info;
+use minitrace::future::FutureExt;
+use minitrace::prelude::*;
 use tokio::sync::Semaphore;
 use tokio::time::sleep;
 
@@ -57,6 +59,7 @@ pub struct Dispatch<Req: KvRequest> {
 impl<Req: KvRequest> Plan for Dispatch<Req> {
     type Result = Req::Response;
 
+    #[minitrace::trace]
     async fn execute(&self) -> Result<Self::Result> {
         let stats = tikv_stats(self.request.label());
         let result = self
@@ -104,6 +107,7 @@ where
 {
     // A plan may involve multiple shards
     #[async_recursion]
+    #[minitrace::trace]
     async fn single_plan_handler(
         pd_client: Arc<PdC>,
         current_plan: P,
@@ -117,14 +121,17 @@ where
             let (shard, region_store) = shard?;
             let mut clone = current_plan.clone();
             clone.apply_shard(shard, &region_store)?;
-            let handle = tokio::spawn(Self::single_shard_handler(
-                pd_client.clone(),
-                clone,
-                region_store,
-                backoff.clone(),
-                permits.clone(),
-                preserve_region_results,
-            ));
+            let handle = tokio::spawn(
+                Self::single_shard_handler(
+                    pd_client.clone(),
+                    clone,
+                    region_store,
+                    backoff.clone(),
+                    permits.clone(),
+                    preserve_region_results,
+                )
+                .in_span(Span::enter_with_local_parent("single_shard_handler")),
+            );
             handles.push(handle);
         }
 
@@ -149,6 +156,7 @@ where
     }
 
     #[async_recursion]
+    #[minitrace::trace]
     async fn single_shard_handler(
         pd_client: Arc<PdC>,
         plan: P,
@@ -210,11 +218,17 @@ where
     // 1. Ok(true): error has been resolved, retry immediately
     // 2. Ok(false): backoff, and then retry
     // 3. Err(Error): can't be resolved, return the error to upper level
+    #[minitrace::trace]
     async fn handle_region_error(
         pd_client: Arc<PdC>,
         e: errorpb::Error,
         region_store: RegionStore,
     ) -> Result<bool> {
+        debug!(
+            "handle_region_error, error:{:?}, region_store:{:?}",
+            e, region_store
+        );
+
         let ver_id = region_store.region_with_leader.ver_id();
         if let Some(not_leader) = e.not_leader {
             if let Some(leader) = not_leader.leader {
@@ -266,6 +280,7 @@ where
     // 1. Ok(true): error has been resolved, retry immediately
     // 2. Ok(false): backoff, and then retry
     // 3. Err(Error): can't be resolved, return the error to upper level
+    #[minitrace::trace]
     async fn on_region_epoch_not_match(
         pd_client: Arc<PdC>,
         region_store: RegionStore,
@@ -302,6 +317,7 @@ where
         Ok(false)
     }
 
+    #[minitrace::trace]
     async fn handle_grpc_error(
         pd_client: Arc<PdC>,
         plan: P,
@@ -349,6 +365,7 @@ where
 {
     type Result = Vec<Result<P::Result>>;
 
+    #[minitrace::trace]
     async fn execute(&self) -> Result<Self::Result> {
         // Limit the maximum concurrency of multi-region request. If there are
         // too many concurrent requests, TiKV is more likely to return a "TiKV
@@ -469,6 +486,7 @@ impl<In: Clone + Send + Sync + 'static, P: Plan<Result = Vec<Result<In>>>, M: Me
 {
     type Result = M::Out;
 
+    #[minitrace::trace]
     async fn execute(&self) -> Result<Self::Result> {
         self.merge.merge(self.inner.execute().await?)
     }
@@ -565,27 +583,43 @@ where
 {
     type Result = P::Result;
 
+    #[minitrace::trace]
     async fn execute(&self) -> Result<Self::Result> {
         let mut result = self.inner.execute().await?;
         let mut clone = self.clone();
+        let mut retry_cnt = 0;
         loop {
+            retry_cnt += 1;
+            let _span = LocalSpan::enter_with_local_parent("ResolveLock::execute::retry")
+                .with_property(|| ("retry_count", retry_cnt.to_string()));
+
             let locks = result.take_locks();
             if locks.is_empty() {
+                debug!("ResolveLock::execute ok");
                 return Ok(result);
             }
 
             if self.backoff.is_none() {
+                debug!("ResolveLock::execute lock error");
                 return Err(Error::ResolveLockError(locks));
             }
 
             let pd_client = self.pd_client.clone();
             let live_locks = resolve_locks(locks, pd_client.clone()).await?;
             if live_locks.is_empty() {
+                debug!("ResolveLock::execute lock error retry (resolved)",);
                 result = self.inner.execute().await?;
             } else {
                 match clone.backoff.next_delay_duration() {
-                    None => return Err(Error::ResolveLockError(live_locks)),
+                    None => {
+                        debug!("ResolveLock::execute lock error");
+                        return Err(Error::ResolveLockError(live_locks));
+                    }
                     Some(delay_duration) => {
+                        debug!(
+                            "ResolveLock::execute lock error retry (delay {:?})",
+                            delay_duration
+                        );
                         sleep(delay_duration).await;
                         result = clone.inner.execute().await?;
                     }
@@ -595,7 +629,7 @@ where
     }
 }
 
-#[derive(Default)]
+#[derive(Default, Debug)]
 pub struct CleanupLocksResult {
     pub region_error: Option<errorpb::Error>,
     pub key_error: Option<Vec<Error>>,
@@ -667,6 +701,7 @@ where
 {
     type Result = CleanupLocksResult;
 
+    #[minitrace::trace]
     async fn execute(&self) -> Result<Self::Result> {
         let mut result = CleanupLocksResult::default();
         let mut inner = self.inner.clone();