feat: databend-metabench: benchmark list (#18745)

drmingdrmer · web-flow · commit 5a68f144ea13 · 2025-09-19T00:48:01.000+08:00
diff --git a/src/meta/binaries/metabench/main.rs b/src/meta/binaries/metabench/main.rs
@@ -17,6 +17,8 @@
 use std::collections::BTreeMap;
 use std::fmt::Debug;
 use std::fmt::Display;
+use std::sync::atomic::AtomicU64;
+use std::sync::atomic::Ordering;
 use std::sync::Arc;
 use std::time::Duration;
 use std::time::Instant;
@@ -55,6 +57,7 @@ use databend_common_tracing::LogFormat;
 use databend_common_tracing::StderrConfig;
 use databend_common_version::BUILD_INFO;
 use databend_common_version::METASRV_COMMIT_VERSION;
+use futures::TryStreamExt;
 use serde::Deserialize;
 use serde::Serialize;
 use tokio::time::sleep;
@@ -84,6 +87,12 @@ struct Config {
     /// "get_table": single get_table() rpc;
     /// "table_copy_file": upsert table with copy file.
     /// "table_copy_file:{"file_cnt":100}": upsert table with 100 copy files. After ":" is a json config string
+    /// "list": list all keys with default prefix pattern;
+    /// "list:{"limit":50}": list up to 50 keys with default prefix pattern;
+    /// "list:{"prefix":"custom_prefix"}": list all keys with custom prefix;
+    /// "list:{"prefix":"custom_prefix","limit":50}": list up to 50 keys with custom prefix;
+    /// "list:{"interval_ms":100}": add 100ms delay between reading each item (slow client simulation);
+    /// "list:{"prefix":"custom_prefix","limit":50,"interval_ms":100}": combine all options;
     #[clap(long, default_value = "upsert_kv")]
     pub rpc: String,
 }
@@ -117,39 +126,33 @@ async fn main() {
         return;
     }
 
+    let client = MetaGrpcClient::try_create_with_features(
+        vec![config.grpc_api_address.clone()],
+        &BUILD_INFO,
+        "root",
+        "xxx",
+        None,
+        None,
+        None,
+        required::read_write(),
+    )
+    .unwrap();
+
     let start = Instant::now();
     let mut client_num = 0;
     let mut handles = Vec::new();
     while client_num < config.client {
         client_num += 1;
-        let addr = config.grpc_api_address.clone();
         let rpc = config.rpc.clone();
         let prefix = config.prefix;
 
         let cmd_and_param = rpc.splitn(2, ':').collect::<Vec<_>>();
         let cmd = cmd_and_param[0].to_string();
         let param = cmd_and_param.get(1).unwrap_or(&"").to_string();
 
-        let handle = runtime::spawn(async move {
-            let client = MetaGrpcClient::try_create_with_features(
-                vec![addr.to_string()],
-                &BUILD_INFO,
-                "root",
-                "xxx",
-                None,
-                None,
-                None,
-                required::read_write(),
-            );
-
-            let client = match client {
-                Ok(client) => client,
-                Err(e) => {
-                    eprintln!("Failed to create client: {}", e);
-                    return;
-                }
-            };
+        let client = client.clone();
 
+        let handle = runtime::spawn(async move {
             for i in 0..config.number {
                 if cmd == "upsert_kv" {
                     benchmark_upsert(&client, prefix, client_num, i).await;
@@ -161,6 +164,8 @@ async fn main() {
                     benchmark_table_copy_file(&client, prefix, client_num, i, &param).await;
                 } else if cmd == "semaphore" {
                     benchmark_semaphore(&client, prefix, client_num, i, &param).await;
+                } else if cmd == "list" {
+                    benchmark_list(&client, prefix, client_num, i, &param).await;
                 } else {
                     unreachable!("Invalid config.rpc: {}", rpc);
                 }
@@ -447,6 +452,93 @@ async fn benchmark_semaphore(
     }
 }
 
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Default)]
+struct ListConfig {
+    /// Maximum number of keys to return in the list operation.
+    /// If None, all matching keys are returned.
+    limit: Option<usize>,
+    /// The prefix to search for. If None, uses the default pattern "{prefix}-{client_num}".
+    prefix: Option<String>,
+    /// Interval in milliseconds to wait between reading each item from the stream.
+    /// This simulates a slow client. If None or 0, no delay is added.
+    interval_ms: Option<u64>,
+}
+
+/// Benchmark listing keys with a prefix.
+async fn benchmark_list(
+    client: &Arc<ClientHandle>,
+    prefix: u64,
+    client_num: u64,
+    i: u64,
+    param: &str,
+) {
+    let name = format!("client[{:>05}]-{}th", client_num, i);
+
+    let config = if param.is_empty() {
+        ListConfig::default()
+    } else {
+        serde_json::from_str(param).unwrap()
+    };
+
+    let key_prefix = config
+        .prefix
+        .clone()
+        .unwrap_or_else(|| format!("{}-{}", prefix, client_num));
+
+    if i % 100 == 0 {
+        println!("{:>10}-th list using prefix: '{}'", name, key_prefix);
+    }
+
+    let start_time = Instant::now();
+    let stream_res = client.list(&key_prefix).await;
+    let stream_returned_time = Instant::now();
+
+    static TOTAL: AtomicU64 = AtomicU64::new(0);
+    TOTAL.fetch_add(1, Ordering::Relaxed);
+
+    println!(
+        "{:>10}-th list stream returned in {:?}, err: {:?}, total streams: {}",
+        name,
+        stream_returned_time.duration_since(start_time),
+        stream_res.as_ref().err(),
+        TOTAL.load(Ordering::Relaxed)
+    );
+
+    let mut strm = match stream_res {
+        Ok(stream) => stream,
+        Err(e) => {
+            println!("{:>10}-th list error: {:?}", name, e);
+            return;
+        }
+    };
+
+    let mut count = 0;
+
+    while let Ok(Some(_item)) = strm.try_next().await {
+        count += 1;
+
+        // Apply interval delay if specified (simulate slow client)
+        if let Some(interval_ms) = config.interval_ms {
+            if interval_ms > 0 {
+                sleep(Duration::from_millis(interval_ms)).await;
+            }
+        }
+
+        // Apply limit if specified
+        if let Some(limit) = config.limit {
+            if count >= limit {
+                break;
+            }
+        }
+
+        if count % 10 == 0 {
+            println!("{:>10}-th list found {} keys", name, count);
+        }
+    }
+
+    println!("{:>10}-th list found {} keys", name, count);
+}
+
 fn print_res<D: Debug>(i: u64, typ: impl Display, res: &D) {
     if i % 100 == 0 {
         println!("{:>10}-th {} result: {:?}", i, typ, res);
diff --git a/src/meta/raft-store/src/leveled_store/leveled_map/impl_scoped_seq_bounded_range.rs b/src/meta/raft-store/src/leveled_store/leveled_map/impl_scoped_seq_bounded_range.rs
@@ -14,8 +14,10 @@
 
 use std::io::Error;
 use std::ops::RangeBounds;
+use std::time::Instant;
 
 use futures_util::StreamExt;
+use log::debug;
 use log::warn;
 use map_api::mvcc;
 use map_api::mvcc::ViewKey;
@@ -60,7 +62,24 @@ where
         // writable level
 
         let (vec, immutable) = {
+            let start = Instant::now();
+            debug!(
+                "Level.writable::range(start={:?}, end={:?})",
+                range.start_bound(),
+                range.end_bound()
+            );
+
             let inner = self.data.lock().unwrap();
+
+            debug!(
+                "Level.writable::range(start={:?}, end={:?}) acquired lock, took {:?}, writable: kv.len={}, expire.len={}",
+                range.start_bound(),
+                range.end_bound(),
+                start.elapsed(),
+                inner.writable.kv.inner.len(),
+                inner.writable.expire.inner.len()
+            );
+
             let it = inner
                 .writable
                 .get_sub_table()
diff --git a/src/meta/service/src/api/grpc/grpc_service.rs b/src/meta/service/src/api/grpc/grpc_service.rs
@@ -333,19 +333,24 @@ impl MetaService for MetaServiceImpl {
         let req_str = format!("ReadRequest: {:?}", req);
 
         ThreadTracker::tracking_future(async move {
-            let _guard = InFlightRead::guard();
+            let guard = InFlightRead::guard();
             let start = Instant::now();
             let (endpoint, strm) = self.handle_kv_read_v1(req).in_span(root).await?;
 
             // Counter to track total items sent
             let count = Arc::new(AtomicU64::new(0));
             let count2 = count.clone();
 
-            let strm = strm.map(move |item| {
-                network_metrics::incr_stream_sent_item(req_typ);
-                count2.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
-                item
-            });
+            let strm = strm
+                .map(move |x| {
+                    let _g = &guard; // hold the guard until the stream is done.
+                    x
+                })
+                .map(move |item| {
+                    network_metrics::incr_stream_sent_item(req_typ);
+                    count2.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+                    item
+                });
 
             // Log the total time and item count when the stream is finished.
             let strm = OnCompleteStream::new(strm, move || {
@@ -404,7 +409,7 @@ impl MetaService for MetaServiceImpl {
         &self,
         _request: Request<databend_common_meta_types::protobuf::Empty>,
     ) -> Result<Response<Self::ExportStream>, Status> {
-        let _guard = InFlightRead::guard();
+        let guard = InFlightRead::guard();
 
         let meta_node = self.try_get_meta_node()?;
 
@@ -415,6 +420,10 @@ impl MetaService for MetaServiceImpl {
         // - Convert Vec<String> to ExportedChunk;
         // - Convert TryChunkError<_, io::Error> to Status;
         let s = strm
+            .map(move |x| {
+                let _g = &guard; // hold the guard until the stream is done.
+                x
+            })
             .try_chunks(chunk_size)
             .map_ok(|chunk: Vec<String>| ExportedChunk { data: chunk })
             .map_err(|e: TryChunksError<_, io::Error>| Status::internal(e.1.to_string()));
@@ -433,7 +442,7 @@ impl MetaService for MetaServiceImpl {
         &self,
         request: Request<pb::ExportRequest>,
     ) -> Result<Response<Self::ExportV1Stream>, Status> {
-        let _guard = InFlightRead::guard();
+        let guard = InFlightRead::guard();
 
         let meta_node = self.try_get_meta_node()?;
 
@@ -444,6 +453,10 @@ impl MetaService for MetaServiceImpl {
         // - Convert Vec<String> to ExportedChunk;
         // - Convert TryChunkError<_, io::Error> to Status;
         let s = strm
+            .map(move |x| {
+                let _g = &guard; // hold the guard until the stream is done.
+                x
+            })
             .try_chunks(chunk_size)
             .map_ok(|chunk: Vec<String>| ExportedChunk { data: chunk })
             .map_err(|e: TryChunksError<_, io::Error>| Status::internal(e.1.to_string()));
diff --git a/src/meta/service/src/meta_service/meta_node.rs b/src/meta/service/src/meta_service/meta_node.rs
@@ -337,7 +337,10 @@ impl MetaNode {
         let socket_addr = ip_port.parse::<std::net::SocketAddr>()?;
         let node_id = meta_node.raft_store.id;
 
-        let srv = tonic::transport::Server::builder().add_service(raft_server);
+        let srv = tonic::transport::Server::builder()
+            // .concurrency_limit_per_connection()
+            // .timeout(Duration::from_secs(60))
+            .add_service(raft_server);
 
         let h = databend_common_base::runtime::spawn(async move {
             srv.serve_with_shutdown(socket_addr, async move {
diff --git a/src/meta/service/src/store/store_inner.rs b/src/meta/service/src/store/store_inner.rs
@@ -202,7 +202,7 @@ impl RaftStoreInner {
 
         let mut compactor = self.state_machine().new_compactor(compactor_permit);
 
-        info!("do_build_snapshot compactor created: {:?}", compactor);
+        info!("do_build_snapshot compactor created");
 
         let (mut sys_data, mut strm) = compactor
             .compact_into_stream()
