[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();