mz_persist_client/cli/bench/
blob_read.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//! Concurrent [S3Blob] reads, plus a steady stream of PUTs through the same
11//! client, against an S3-compatible store.
12//!
13//! Written to measure GCS read throughput (PER-136). Over HTTP/2, the SDK's
14//! HTTP client sends every request to a host over one connection, and one
15//! connection to GCS carried only 300 to 400 MiB/s.
16//! `mz_aws_util::pooled_http_client` spreads requests over more connections.
17//! Compare `--client sdk` with `--client pooled` to recheck that, or to tune
18//! the pool's size.
19
20use std::collections::BTreeMap;
21use std::sync::atomic::{AtomicUsize, Ordering};
22use std::sync::{Arc, Mutex};
23use std::time::{Duration, Instant};
24
25use anyhow::bail;
26use aws_smithy_runtime_api::box_error::BoxError;
27use aws_smithy_runtime_api::client::connector_metadata::ConnectorMetadata;
28use aws_smithy_runtime_api::client::http::{
29    HttpClient, HttpConnector, HttpConnectorFuture, HttpConnectorSettings, SharedHttpClient,
30    SharedHttpConnector,
31};
32use aws_smithy_runtime_api::client::orchestrator::HttpRequest;
33use aws_smithy_runtime_api::client::runtime_components::{
34    RuntimeComponents, RuntimeComponentsBuilder,
35};
36use aws_smithy_types::config_bag::ConfigBag;
37use bytes::Bytes;
38use futures::{StreamExt, TryStreamExt};
39use mz_ore::cast::CastLossy;
40use mz_ore::error::ErrorExt;
41use mz_ore::metrics::MetricsRegistry;
42use mz_ore::now::SYSTEM_TIME;
43use mz_ore::secure::SecureString;
44use mz_persist::location::Blob;
45use mz_persist::metrics::S3BlobMetrics;
46use mz_persist::s3::{BLOB_HTTP_POOL_MAX_IN_FLIGHT_PER_CLIENT, S3Blob, S3BlobConfig};
47
48use crate::cfg::PersistConfig;
49
50/// Concurrent persist blob reads against an S3-compatible store.
51#[derive(Debug, clap::Args)]
52pub struct BlobReadArgs {
53    #[clap(subcommand)]
54    command: Command,
55}
56
57#[derive(Debug, clap::Subcommand)]
58enum Command {
59    /// Writes test objects with `S3Blob::set`, so objects over the multipart
60    /// threshold are uploaded in parts, like production blobs.
61    Upload(UploadArgs),
62    /// Reads the test objects with N concurrent readers for each N, while
63    /// writing small objects at a steady rate, and reports latency,
64    /// throughput and errors.
65    Read(ReadArgs),
66}
67
68#[derive(Debug, clap::Args)]
69struct StoreArgs {
70    #[clap(long, default_value = "https://storage.googleapis.com")]
71    endpoint: String,
72    #[clap(long, default_value = "auto")]
73    region: String,
74    #[clap(long)]
75    bucket: String,
76    /// Key prefix of the test objects.
77    #[clap(long, default_value = "blob-read-bench")]
78    prefix: String,
79    #[clap(long, env = "BLOB_READ_BENCH_ACCESS_KEY_ID")]
80    access_key_id: String,
81    #[clap(
82        long,
83        env = "BLOB_READ_BENCH_SECRET_ACCESS_KEY",
84        hide_env_values = true
85    )]
86    secret_access_key: SecureString,
87}
88
89#[derive(Debug, clap::Args)]
90struct UploadArgs {
91    #[clap(flatten)]
92    store: StoreArgs,
93    #[clap(long, value_delimiter = ',', default_value = "8,16,64")]
94    sizes_mib: Vec<usize>,
95    #[clap(long, default_value_t = 70)]
96    count_per_size: usize,
97    #[clap(long, default_value_t = 16)]
98    concurrency: usize,
99}
100
101#[derive(Debug, clap::Args)]
102struct ReadArgs {
103    #[clap(flatten)]
104    store: StoreArgs,
105    #[clap(long, value_enum)]
106    client: ClientKind,
107    /// Runs one round per value, in order.
108    #[clap(long, value_delimiter = ',', default_value = "16,64,256,512")]
109    concurrency: Vec<usize>,
110    /// Readers start no new get after this long. A round lasts until every
111    /// in-flight get finishes.
112    #[clap(long, value_parser = humantime::parse_duration, default_value = "30s")]
113    duration: Duration,
114    /// PUTs started per second during each round. 0 turns them off.
115    #[clap(long, default_value_t = 10)]
116    puts_per_sec: u32,
117    #[clap(long, default_value_t = 1024)]
118    put_size_kib: usize,
119    /// A PUT that takes longer than this counts as stalled.
120    #[clap(long, value_parser = humantime::parse_duration, default_value = "10s")]
121    put_stall: Duration,
122}
123
124/// PUTs cycle through this many keys under [PUT_KEY_PREFIX], overwriting old
125/// ones, so they don't pile up in the bucket.
126const PUT_KEYS: usize = 100;
127/// Keeps PUT objects out of the keys that readers fetch.
128const PUT_KEY_PREFIX: &str = "put/";
129
130#[derive(Debug, Clone, Copy, clap::ValueEnum)]
131enum ClientKind {
132    /// `mz_aws_util::http_client`, the SDK's client.
133    Sdk,
134    /// `mz_aws_util::pooled_http_client`, which spreads requests over a
135    /// growing set of `Sdk` clients.
136    Pooled,
137}
138
139/// Runs the given blob-read bench command.
140pub(crate) async fn run(args: BlobReadArgs) -> Result<(), anyhow::Error> {
141    match args.command {
142        Command::Upload(args) => upload(args).await,
143        Command::Read(args) => read(args).await,
144    }
145}
146
147/// Opens an [S3Blob] with persist's production timeouts and no hedging.
148async fn open(
149    store: &StoreArgs,
150    client: ClientKind,
151    registry: &MetricsRegistry,
152    header_waits: Arc<HeaderWaits>,
153) -> Result<S3Blob, anyhow::Error> {
154    let knobs = PersistConfig::new_default_configs(&crate::BUILD_INFO, SYSTEM_TIME.clone());
155    let metrics = S3BlobMetrics::new(registry);
156    let pool_metrics = metrics.http_pool();
157    let config = S3BlobConfig::new(
158        store.bucket.clone(),
159        store.prefix.clone(),
160        None,
161        Some(store.endpoint.clone()),
162        Some(store.region.clone()),
163        Some((
164            store.access_key_id.clone(),
165            store.secret_access_key.unsecure().to_owned(),
166        )),
167        Box::new(knobs),
168        metrics,
169    )
170    .await?;
171    let http_client = match client {
172        ClientKind::Sdk => SharedHttpClient::new(mz_aws_util::http_client()),
173        ClientKind::Pooled => SharedHttpClient::new(mz_aws_util::pooled_http_client(
174            || *BLOB_HTTP_POOL_MAX_IN_FLIGHT_PER_CLIENT.default(),
175            pool_metrics,
176        )),
177    };
178    let config = config.with_http_client(TimedHttpClient {
179        inner: http_client,
180        waits: header_waits,
181    });
182    Ok(S3Blob::open(config).await?)
183}
184
185async fn upload(args: UploadArgs) -> Result<(), anyhow::Error> {
186    let blob = open(
187        &args.store,
188        ClientKind::Sdk,
189        &MetricsRegistry::new(),
190        Arc::default(),
191    )
192    .await?;
193    let objects: Vec<_> = args
194        .sizes_mib
195        .iter()
196        .flat_map(|size| (0..args.count_per_size).map(move |i| (format!("{size}mib-{i}"), *size)))
197        .collect();
198    let mut uploads = futures::stream::iter(objects)
199        .map(|(key, size_mib)| {
200            let blob = &blob;
201            async move {
202                let mut data = vec![0; size_mib * 1024 * 1024];
203                rand::fill(&mut data[..]);
204                blob.set(&key, Bytes::from(data)).await?;
205                Ok::<_, anyhow::Error>(key)
206            }
207        })
208        .buffer_unordered(args.concurrency);
209    while let Some(key) = uploads.try_next().await? {
210        println!("uploaded {key}");
211    }
212    Ok(())
213}
214
215#[derive(Debug, Default)]
216struct ReaderStats {
217    latencies: Vec<Duration>,
218    failed_latencies: Vec<Duration>,
219    bytes: u64,
220    errors: BTreeMap<String, u64>,
221}
222
223#[derive(Debug, Default)]
224struct PutStats {
225    latencies: Vec<Duration>,
226    stalled: usize,
227    errors: BTreeMap<String, u64>,
228}
229
230async fn read(args: ReadArgs) -> Result<(), anyhow::Error> {
231    let registry = MetricsRegistry::new();
232    let header_waits = Arc::new(HeaderWaits::default());
233    let blob = Arc::new(
234        open(
235            &args.store,
236            args.client,
237            &registry,
238            Arc::clone(&header_waits),
239        )
240        .await?,
241    );
242    let mut keys = Vec::new();
243    blob.list_keys_and_metadata("", &mut |meta| {
244        if !meta.key.starts_with(PUT_KEY_PREFIX) {
245            keys.push(meta.key.to_owned());
246        }
247    })
248    .await?;
249    if keys.is_empty() {
250        bail!("no objects under prefix {}", args.store.prefix);
251    }
252    let keys = Arc::new(keys);
253    let mut put_data = vec![0; args.put_size_kib * 1024];
254    rand::fill(&mut put_data[..]);
255    let put_data = Bytes::from(put_data);
256    let next_put = Arc::new(AtomicUsize::new(0));
257
258    for &concurrency in &args.concurrency {
259        let read_timeouts_before = read_timeouts(&registry);
260        header_waits.take();
261        let next_key = Arc::new(AtomicUsize::new(0));
262        let start = Instant::now();
263        let deadline = start + args.duration;
264        let putter = (args.puts_per_sec > 0).then(|| {
265            let blob = Arc::clone(&blob);
266            let put_data = put_data.clone();
267            let next_put = Arc::clone(&next_put);
268            let period = Duration::from_secs(1) / args.puts_per_sec;
269            let put_stall = args.put_stall;
270            mz_ore::task::spawn(|| "blob_read_bench_putter", async move {
271                // Each PUT runs in its own task, so a stalled one doesn't
272                // delay the next.
273                let mut puts = Vec::new();
274                let mut ticker = tokio::time::interval(period);
275                while Instant::now() < deadline {
276                    ticker.tick().await;
277                    let n = next_put.fetch_add(1, Ordering::Relaxed) % PUT_KEYS;
278                    let key = format!("{PUT_KEY_PREFIX}{n}");
279                    let blob = Arc::clone(&blob);
280                    let put_data = put_data.clone();
281                    puts.push(mz_ore::task::spawn(|| "blob_read_bench_put", async move {
282                        let put_start = Instant::now();
283                        let result = blob.set(&key, put_data).await;
284                        (put_start.elapsed(), result)
285                    }));
286                }
287                let mut stats = PutStats::default();
288                for put in puts {
289                    let (elapsed, result) = put.await;
290                    if elapsed > put_stall {
291                        stats.stalled += 1;
292                    }
293                    match result {
294                        Ok(()) => stats.latencies.push(elapsed),
295                        Err(err) => {
296                            *stats
297                                .errors
298                                .entry(err.display_with_causes().to_string())
299                                .or_default() += 1
300                        }
301                    }
302                }
303                stats
304            })
305        });
306        let readers: Vec<_> = (0..concurrency)
307            .map(|_| {
308                let blob = Arc::clone(&blob);
309                let keys = Arc::clone(&keys);
310                let next_key = Arc::clone(&next_key);
311                mz_ore::task::spawn(|| "blob_read_bench_reader", async move {
312                    let mut stats = ReaderStats::default();
313                    while Instant::now() < deadline {
314                        let key = &keys[next_key.fetch_add(1, Ordering::Relaxed) % keys.len()];
315                        let get_start = Instant::now();
316                        let result = blob.get(key).await;
317                        let elapsed = get_start.elapsed();
318                        let err = match result {
319                            Ok(Some(value)) => {
320                                stats.latencies.push(elapsed);
321                                stats.bytes += u64::try_from(value.len()).expect("fits");
322                                continue;
323                            }
324                            Ok(None) => format!("{key} missing"),
325                            Err(err) => err.display_with_causes().to_string(),
326                        };
327                        stats.failed_latencies.push(elapsed);
328                        *stats.errors.entry(err).or_default() += 1;
329                    }
330                    stats
331                })
332            })
333            .collect();
334
335        let mut total = ReaderStats::default();
336        for reader in readers {
337            let stats = reader.await;
338            total.latencies.extend(stats.latencies);
339            total.failed_latencies.extend(stats.failed_latencies);
340            total.bytes += stats.bytes;
341            for (err, count) in stats.errors {
342                *total.errors.entry(err).or_default() += count;
343            }
344        }
345        let elapsed = start.elapsed();
346        let puts = match putter {
347            Some(putter) => putter.await,
348            None => PutStats::default(),
349        };
350
351        // Read timeouts count PUTs as well as GETs.
352        println!(
353            "client={:?} concurrency={} elapsed={:.1}s read_timeouts={} get_MiB/s={:.1}",
354            args.client,
355            concurrency,
356            elapsed.as_secs_f64(),
357            read_timeouts(&registry) - read_timeouts_before,
358            f64::cast_lossy(total.bytes) / (1024.0 * 1024.0) / elapsed.as_secs_f64(),
359        );
360        println!(
361            "  gets: {} errors={}",
362            summary(total.latencies),
363            total.errors.values().sum::<u64>(),
364        );
365        println!("  failed gets: {}", summary(total.failed_latencies));
366        println!(
367            "  puts: {} errors={} stalled(>{:?})={}",
368            summary(puts.latencies),
369            puts.errors.values().sum::<u64>(),
370            args.put_stall,
371            puts.stalled,
372        );
373        for (method, waits) in header_waits.take() {
374            println!("  {method} wait for headers: {}", summary(waits));
375        }
376        for (err, count) in &total.errors {
377            println!("  get error {count}x: {err}");
378        }
379        for (err, count) in &puts.errors {
380            println!("  put error {count}x: {err}");
381        }
382    }
383    Ok(())
384}
385
386/// Count, p50, p99 and max of `samples`.
387fn summary(mut samples: Vec<Duration>) -> String {
388    samples.sort();
389    let percentile = |p: usize| {
390        samples
391            .get(samples.len().saturating_sub(1) * p / 100)
392            .copied()
393            .unwrap_or_default()
394    };
395    format!(
396        "n={} p50={:?} p99={:?} max={:?}",
397        samples.len(),
398        percentile(50),
399        percentile(99),
400        percentile(100),
401    )
402}
403
404/// Per HTTP method, the time from each request reaching the HTTP client until
405/// its response head arrives or it fails. Includes waiting for a connection,
406/// or for a free stream on a shared HTTP/2 connection. Each SDK retry is a
407/// separate request.
408#[derive(Debug, Default)]
409struct HeaderWaits(Mutex<BTreeMap<String, Vec<Duration>>>);
410
411impl HeaderWaits {
412    fn record(&self, method: String, wait: Duration) {
413        let mut waits = self.0.lock().expect("lock poisoned");
414        waits.entry(method).or_default().push(wait);
415    }
416
417    fn take(&self) -> BTreeMap<String, Vec<Duration>> {
418        std::mem::take(&mut *self.0.lock().expect("lock poisoned"))
419    }
420}
421
422/// Wraps an [HttpClient] to record [HeaderWaits].
423#[derive(Debug)]
424struct TimedHttpClient {
425    inner: SharedHttpClient,
426    waits: Arc<HeaderWaits>,
427}
428
429impl HttpClient for TimedHttpClient {
430    fn http_connector(
431        &self,
432        settings: &HttpConnectorSettings,
433        components: &RuntimeComponents,
434    ) -> SharedHttpConnector {
435        SharedHttpConnector::new(TimedConnector {
436            inner: self.inner.http_connector(settings, components),
437            waits: Arc::clone(&self.waits),
438        })
439    }
440
441    fn validate_base_client_config(
442        &self,
443        runtime_components: &RuntimeComponentsBuilder,
444        cfg: &ConfigBag,
445    ) -> Result<(), BoxError> {
446        self.inner
447            .validate_base_client_config(runtime_components, cfg)
448    }
449
450    fn validate_final_config(
451        &self,
452        runtime_components: &RuntimeComponents,
453        cfg: &ConfigBag,
454    ) -> Result<(), BoxError> {
455        self.inner.validate_final_config(runtime_components, cfg)
456    }
457
458    fn connector_metadata(&self) -> Option<ConnectorMetadata> {
459        self.inner.connector_metadata()
460    }
461}
462
463#[derive(Debug)]
464struct TimedConnector {
465    inner: SharedHttpConnector,
466    waits: Arc<HeaderWaits>,
467}
468
469impl HttpConnector for TimedConnector {
470    fn call(&self, request: HttpRequest) -> HttpConnectorFuture {
471        let method = request.method().to_owned();
472        let waits = Arc::clone(&self.waits);
473        let start = Instant::now();
474        let response = self.inner.call(request);
475        HttpConnectorFuture::new(async move {
476            let response = response.await;
477            waits.record(method, start.elapsed());
478            response
479        })
480    }
481}
482
483/// The SDK's read timeout firings so far: the `mz_persist_s3_read_timeouts`
484/// counter, which can't be read from [S3BlobMetrics] outside its crate.
485fn read_timeouts(registry: &MetricsRegistry) -> u64 {
486    registry
487        .gather()
488        .iter()
489        .find(|family| family.name() == "mz_persist_s3_read_timeouts")
490        .map_or(0, |family| {
491            u64::cast_lossy(family.get_metric()[0].get_counter().value())
492        })
493}
mz_persist_client/cli/bench/blob_read.rs

mz_persist_client/cli/bench/
blob_read.rs
