mz_persist_client/cli/
bench.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! CLI benchmarking tools for persist
11
12use futures_util::stream::StreamExt;
13use futures_util::{TryStreamExt, stream};
14use std::sync::Arc;
15use std::time::Instant;
16
17use mz_persist::indexed::encoding::BlobTraceBatchPart;
18
19use crate::cli::args::StateArgs;
20use crate::cli::bench::blob_read::BlobReadArgs;
21use crate::internal::state::BatchPart;
22
23mod blob_read;
24
25/// Commands for read-only inspection of persist state
26#[derive(Debug, clap::Args)]
27pub struct BenchArgs {
28    #[clap(subcommand)]
29    command: Command,
30}
31
32/// Individual subcommands of bench
33#[derive(Debug, clap::Subcommand)]
34pub(crate) enum Command {
35    /// Fetch the blobs in a shard as quickly as possible, repeated some number
36    /// of times.
37    S3Fetch(S3FetchArgs),
38    /// Read and write test objects in an S3-compatible store with many
39    /// concurrent readers, comparing HTTP clients.
40    BlobRead(BlobReadArgs),
41}
42
43/// Fetch the blobs in a shard as quickly as possible, repeated some number of
44/// times.
45#[derive(Debug, Clone, clap::Parser)]
46pub struct S3FetchArgs {
47    #[clap(flatten)]
48    shard: StateArgs,
49
50    #[clap(long, default_value_t = 1)]
51    iters: usize,
52
53    #[clap(long)]
54    parse: bool,
55}
56
57/// Runs the given bench command.
58pub async fn run(command: BenchArgs) -> Result<(), anyhow::Error> {
59    match command.command {
60        Command::S3Fetch(args) => bench_s3(&args).await?,
61        Command::BlobRead(args) => blob_read::run(args).await?,
62    }
63
64    Ok(())
65}
66
67async fn bench_s3(args: &S3FetchArgs) -> Result<(), anyhow::Error> {
68    let parse = args.parse;
69    let shard_id = args.shard.shard_id();
70    let state_versions = args.shard.open().await?;
71    let versions = state_versions
72        .fetch_recent_live_diffs::<u64>(&shard_id)
73        .await;
74    let state = state_versions
75        .fetch_current_state::<u64>(&shard_id, versions.0)
76        .await;
77    let state = state.check_ts_codec(&shard_id)?;
78    let snap = state
79        .snapshot(state.since())
80        .expect("since should be available for reads");
81
82    let batch_parts: Vec<_> = stream::iter(&snap)
83        .flat_map(|batch| {
84            batch.part_stream(shard_id, &*state_versions.blob, &*state_versions.metrics)
85        })
86        .try_collect()
87        .await?;
88
89    println!("iter,key,size_bytes,fetch_secs,parse_secs");
90    for iter in 0..args.iters {
91        let start = Instant::now();
92        let mut fetches = Vec::new();
93        for part in &batch_parts {
94            let key = match &**part {
95                BatchPart::Hollow(x) => x.key.complete(&shard_id),
96                BatchPart::Inline { .. } => continue,
97            };
98            let blob = Arc::clone(&state_versions.blob);
99            let metrics = Arc::clone(&state_versions.metrics);
100            let fetch = mz_ore::task::spawn(|| "", async move {
101                let buf = blob.get(&key).await.unwrap().unwrap();
102                let fetch_elapsed = start.elapsed();
103                let buf_len = buf.len();
104                let parse_elapsed = mz_ore::task::spawn_blocking(
105                    || "",
106                    move || {
107                        let start = Instant::now();
108                        if parse {
109                            BlobTraceBatchPart::<u64>::decode(&buf, &metrics.columnar).unwrap();
110                        }
111                        start.elapsed()
112                    },
113                )
114                .await;
115                (
116                    key,
117                    buf_len,
118                    fetch_elapsed.as_secs_f64(),
119                    parse_elapsed.as_secs_f64(),
120                )
121            });
122            fetches.push(fetch);
123        }
124        for fetch in fetches {
125            let (key, size_bytes, fetch_secs, parse_secs) = fetch.await;
126            println!(
127                "{},{},{},{},{}",
128                iter, key, size_bytes, fetch_secs, parse_secs
129            );
130        }
131    }
132
133    Ok(())
134}
mz_persist_client/cli/bench.rs

mz_persist_client/cli/
bench.rs
