From a59ace150a6f00ee30874948a8c9b0bd461fc8f5 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:48:45 -0500 Subject: [PATCH 1/2] bench: add an object store that simulates remote storage over local files Adds `SimulatedObjectStore`, which presents a local filesystem the way S3 or GCS present themselves, so benchmarks can exercise realistic IO without a network. The point is fidelity of the request *pattern*, not just of the wall clock. Neither `AmazonS3` nor `GoogleCloudStorage` implements `ObjectStore::get_ranges`; they inherit the trait default, which merges ranges less than 1MiB apart and issues the merged chunks as up to 10 concurrent GETs. So this store implements only the HTTP-shaped primitives a real remote store implements and leaves `get_ranges`, `list_with_offset` and `rename_opts` to the trait defaults, letting arrow-rs's coalescing and fan-out run against it unmodified. Each simulated request pays for a connection from a bounded pool, a time to first byte drawn from an S3-shaped distribution, and transfer time at a fixed per-connection bandwidth. LIST is paginated at 1000 keys with pages charged serially, since each needs the previous continuation token. Not committed yet: the benchmark runner still uses the old store. Co-Authored-By: Claude Opus 5 --- .../util/simulated_object_store/latency.rs | 137 ++++ .../src/util/simulated_object_store/mod.rs | 698 ++++++++++++++++++ 2 files changed, 835 insertions(+) create mode 100644 benchmarks/src/util/simulated_object_store/latency.rs create mode 100644 benchmarks/src/util/simulated_object_store/mod.rs diff --git a/benchmarks/src/util/simulated_object_store/latency.rs b/benchmarks/src/util/simulated_object_store/latency.rs new file mode 100644 index 0000000000000..010d2ead1f6fa --- /dev/null +++ b/benchmarks/src/util/simulated_object_store/latency.rs @@ -0,0 +1,137 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Time-to-first-byte distributions and how they are sampled. + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +/// GET time-to-first-byte distribution, inspired by real S3 latencies. +/// +/// 20 values: 11x P50 (~25-35ms), 5x P75-P90 (~70-110ms), 2x P95 (~120-150ms), +/// 2x P99 (~180-200ms). +/// Sorted: 25,25,28,28,30,30,30,30,32,32,35, 70,85,100,100,110, 130,150, 180,200 +/// P50≈32ms, P90≈110ms, P99≈200ms +pub const GET_TTFB_MS: &[u64] = &[ + 30, 100, 25, 85, 32, 200, 28, 130, 35, 70, 30, 150, 30, 110, 28, 180, 32, 25, 100, 30, +]; + +/// LIST time-to-first-byte distribution, generally higher than GET. +/// +/// This is the cost of a *single page* of results, not of a whole listing. +/// +/// 20 values: 11x P50 (~40-70ms), 5x P75-P90 (~120-180ms), 2x P95 (~200-250ms), +/// 2x P99 (~300-400ms). +/// Sorted: 40,40,50,50,55,55,60,60,65,65,70, 120,140,160,160,180, 210,250, 300,400 +/// P50≈65ms, P90≈180ms, P99≈400ms +pub const LIST_TTFB_MS: &[u64] = &[ + 55, 160, 40, 140, 65, 400, 50, 210, 70, 120, 60, 250, 55, 180, 50, 300, 65, 40, 160, + 60, +]; + +/// Draws request latencies from a fixed distribution. +/// +/// Draws are deterministic: the Nth draw of a run is always the same value, so +/// benchmark runs stay reproducible. +/// +/// The draw counter is hashed rather than used to index the table directly. A +/// plain round robin resonates with the fixed fan-out of +/// [`coalesce_ranges`](object_store::coalesce_ranges), which issues requests in +/// waves of 10: with a 20-entry table every vectored read would see the same +/// two fixed sets of latencies, and a read that happened to issue exactly 10 or +/// 20 requests would always land on exactly the table mean. +#[derive(Debug)] +pub struct LatencySampler { + table: &'static [u64], + draws: AtomicU64, +} + +impl LatencySampler { + pub const fn new(table: &'static [u64]) -> Self { + assert!(!table.is_empty()); + Self { + table, + draws: AtomicU64::new(0), + } + } + + /// Draw the next latency. + pub fn sample(&self) -> Duration { + let draw = self.draws.fetch_add(1, Ordering::Relaxed); + let idx = (splitmix64(draw) % self.table.len() as u64) as usize; + Duration::from_millis(self.table[idx]) + } +} + +/// The SplitMix64 finalizer: stateless, cheap, and mixes well enough that +/// consecutive draw indices land on unrelated table entries. +const fn splitmix64(draw: u64) -> u64 { + let mut z = draw.wrapping_add(0x9E37_79B9_7F4A_7C15); + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sampler_is_deterministic() { + let a = LatencySampler::new(GET_TTFB_MS); + let b = LatencySampler::new(GET_TTFB_MS); + for _ in 0..50 { + assert_eq!(a.sample(), b.sample()); + } + } + + #[test] + fn sampler_reproduces_the_table_distribution() { + let sampler = LatencySampler::new(GET_TTFB_MS); + // Over many draws every table entry should come up, and the mean should + // land near the table mean. A round robin would guarantee this trivially; + // the point here is that hashing does not skew it. + let draws = 20_000; + let total: u64 = (0..draws) + .map(|_| sampler.sample().as_millis() as u64) + .sum(); + let table_mean = + GET_TTFB_MS.iter().sum::() as f64 / GET_TTFB_MS.len() as f64; + let sampled_mean = total as f64 / draws as f64; + assert!( + (sampled_mean - table_mean).abs() < table_mean * 0.05, + "sampled mean {sampled_mean} too far from table mean {table_mean}" + ); + } + + #[test] + fn consecutive_draws_do_not_cycle_with_the_coalesce_fan_out() { + // `coalesce_ranges` issues requests in waves of 10. Two consecutive + // waves must not see the same multiset of latencies, which is exactly + // what a round robin over a 20 entry table would produce. + let sampler = LatencySampler::new(GET_TTFB_MS); + let wave = |sampler: &LatencySampler| -> Vec { + let mut v: Vec<_> = (0..10).map(|_| sampler.sample()).collect(); + v.sort(); + v + }; + let first = wave(&sampler); + let second = wave(&sampler); + let third = wave(&sampler); + assert!(first != second || second != third); + } +} diff --git a/benchmarks/src/util/simulated_object_store/mod.rs b/benchmarks/src/util/simulated_object_store/mod.rs new file mode 100644 index 0000000000000..97cbdf7549218 --- /dev/null +++ b/benchmarks/src/util/simulated_object_store/mod.rs @@ -0,0 +1,698 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! An [`ObjectStore`] that makes a local filesystem behave like remote object +//! storage (S3, GCS), so that benchmarks can be run against realistic IO +//! without a network. +//! +//! # Why not just sleep before each call +//! +//! What matters is the *request pattern*, not only the wall clock. A store that +//! sleeps once per [`ObjectStore`] method call charges one round trip where the +//! real store pays several, and it is blind to exactly the thing an IO +//! optimisation changes. +//! +//! The concrete case is [`ObjectStore::get_ranges`]. Neither `AmazonS3` nor +//! `GoogleCloudStorage` implements it; they inherit the trait default, which +//! merges ranges less than +//! [`OBJECT_STORE_COALESCE_DEFAULT`](object_store::OBJECT_STORE_COALESCE_DEFAULT) +//! bytes apart and then issues the merged chunks as up to 10 *concurrent* GETs. +//! A vectored read of a row group whose column chunks sit more than 1MiB apart +//! is therefore many round trips against S3, and how many is data dependent. +//! +//! So this store deliberately implements **only** the HTTP-shaped primitives a +//! real remote store implements, and leaves `get_ranges` — along with +//! [`ObjectStore::list_with_offset`] and [`ObjectStore::rename_opts`] — to the +//! trait default. Everything arrow-rs does above the wire (coalescing, fan-out, +//! pagination) then runs against the simulator unmodified. +//! +//! Note this is the opposite of the advice in [`ObjectStore`]'s "Wrappers" +//! section, which tells wrappers to implement every method so they do not lose +//! the wrapped store's overrides. That advice is for observability wrappers. +//! Here the wrapped store's overrides are the problem: `LocalFileSystem` +//! implements `get_ranges` as a sequence of positional reads with no coalescing +//! at all, which is precisely the behaviour we need to shed. Do not add +//! `#[deny(clippy::missing_trait_methods)]` to the impl below. +//! +//! # What a simulated request costs +//! +//! 1. waiting for a free connection, bounded by +//! [`SimulatedStoreConfig::max_concurrent_requests`], +//! 2. time to first byte, drawn from a latency distribution (see [`latency`]), +//! 3. transfer time at +//! [`SimulatedStoreConfig::connection_bytes_per_second`], charged as the +//! response body is consumed. +//! +//! Step 3 is what makes coalescing decisions meaningful: merging across a gap +//! trades bytes for round trips, and without a per-byte cost a simulator says +//! that trade is always free. +//! +//! # What is deliberately not modelled +//! +//! arrow-rs never splits one large range into several smaller concurrent +//! requests, so neither does this store: a 200MB coalesced range is one GET on +//! one connection, and is bandwidth bound. Request failures, retries and +//! per-prefix throttling are also not modelled. + +mod latency; + +use std::fmt; +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use futures::StreamExt; +use futures::stream::BoxStream; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, + ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, + Result, +}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; + +pub use latency::{GET_TTFB_MS, LIST_TTFB_MS, LatencySampler}; + +/// The connection pool is owned by the store and never closed. +const POOL_OPEN: &str = "connection pool is never closed"; + +/// How the simulated remote store behaves. +#[derive(Debug, Clone)] +pub struct SimulatedStoreConfig { + /// Time-to-first-byte distribution for GET, in milliseconds. + pub get_ttfb_ms: &'static [u64], + /// Time-to-first-byte distribution for a single LIST page, in milliseconds. + pub list_ttfb_ms: &'static [u64], + /// Sustained throughput of a *single* connection. + /// + /// S3 and GCS deliver roughly 50-100MB/s on one connection no matter how + /// large the object is, which is why one huge coalesced range is slow even + /// though it is only one round trip. + pub connection_bytes_per_second: u64, + /// Maximum number of requests in flight, i.e. the client connection pool. + pub max_concurrent_requests: usize, + /// Keys returned per LIST page. Both S3 and GCS default to 1000. + pub list_page_size: usize, + /// Transfer time is accumulated and only slept once it exceeds this, so + /// that small body chunks do not become a storm of sub-millisecond timer + /// waits that the runtime cannot honour accurately anyway. + pub min_transfer_sleep: Duration, + /// Reads no larger than this are served by one blocking read of the local + /// file rather than a chunked stream. Above it the body is streamed, so a + /// whole-object GET of a multi-gigabyte CSV is not buffered in memory. + pub max_eager_read_bytes: u64, +} + +impl Default for SimulatedStoreConfig { + fn default() -> Self { + Self { + get_ttfb_ms: GET_TTFB_MS, + list_ttfb_ms: LIST_TTFB_MS, + connection_bytes_per_second: 100 * 1024 * 1024, + max_concurrent_requests: 128, + list_page_size: 1000, + min_transfer_sleep: Duration::from_micros(500), + max_eager_read_bytes: 64 * 1024 * 1024, + } + } +} + +/// Presents `T` — in practice a `LocalFileSystem` — as if it were remote object +/// storage. See the [module docs](self). +#[derive(Debug)] +pub struct SimulatedObjectStore { + inner: T, + config: SimulatedStoreConfig, + connections: Arc, + get_ttfb: Arc, + list_ttfb: Arc, +} + +impl SimulatedObjectStore { + pub fn new(inner: T) -> Self { + Self::with_config(inner, SimulatedStoreConfig::default()) + } + + pub fn with_config(inner: T, config: SimulatedStoreConfig) -> Self { + Self { + inner, + connections: Arc::new(Semaphore::new(config.max_concurrent_requests)), + get_ttfb: Arc::new(LatencySampler::new(config.get_ttfb_ms)), + list_ttfb: Arc::new(LatencySampler::new(config.list_ttfb_ms)), + config, + } + } + + pub fn config(&self) -> &SimulatedStoreConfig { + &self.config + } + + /// Take a connection from the pool, waiting if all of them are busy. + async fn connection(&self) -> OwnedSemaphorePermit { + Arc::clone(&self.connections) + .acquire_owned() + .await + .expect(POOL_OPEN) + } + + /// Wrap a response so that its body costs transfer time, holding the + /// connection until the body has been consumed. + async fn simulate_body( + &self, + result: GetResult, + permit: OwnedSemaphorePermit, + ) -> Result { + let rate = self.config.connection_bytes_per_second; + let len = result.range.end - result.range.start; + + if len <= self.config.max_eager_read_bytes { + // Every read the Parquet reader makes is bounded, and its caller + // collects the body in full, so charging the whole transfer up + // front is exact. Doing it this way also keeps `LocalFileSystem`'s + // single-blocking-read fast path instead of paying a + // `spawn_blocking` per 8KiB chunk. + let meta = result.meta.clone(); + let range = result.range.clone(); + let attributes = result.attributes.clone(); + let bytes = result.bytes().await?; + tokio::time::sleep(transfer_time(bytes.len() as u64, rate)).await; + drop(permit); + return Ok(GetResult { + payload: GetResultPayload::Stream( + futures::stream::once(async move { Ok(bytes) }).boxed(), + ), + meta, + range, + attributes, + }); + } + + Ok(self.throttle_body(result, permit)) + } + + /// Charge transfer time incrementally as the body is consumed, so that a + /// large or unbounded read is neither buffered in memory nor able to hide + /// its time-to-first-byte behind its throughput. + fn throttle_body( + &self, + result: GetResult, + permit: OwnedSemaphorePermit, + ) -> GetResult { + let rate = self.config.connection_bytes_per_second; + let min_sleep = self.config.min_transfer_sleep; + let meta = result.meta.clone(); + let range = result.range.clone(); + let attributes = result.attributes.clone(); + + let payload = futures::stream::unfold( + (result.into_stream(), Duration::ZERO, permit), + move |(mut body, mut owed, permit)| async move { + let item = body.next().await?; + if let Ok(bytes) = &item { + owed += transfer_time(bytes.len() as u64, rate); + if owed >= min_sleep { + tokio::time::sleep(owed).await; + owed = Duration::ZERO; + } + } + // `permit` rides along so the connection stays checked out + // until the body is fully read or the stream is dropped. + Some((item, (body, owed, permit))) + }, + ) + .boxed(); + + GetResult { + payload: GetResultPayload::Stream(payload), + meta, + range, + attributes, + } + } +} + +fn transfer_time(bytes: u64, bytes_per_second: u64) -> Duration { + Duration::from_secs_f64(bytes as f64 / bytes_per_second as f64) +} + +impl fmt::Display for SimulatedObjectStore { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "SimulatedObjectStore({})", self.inner) + } +} + +// NOTE: `get_ranges`, `list_with_offset` and `rename_opts` are intentionally +// left to the `ObjectStore` defaults, exactly as the real remote stores leave +// them. See the module docs before adding any of them here. +#[async_trait] +impl ObjectStore for SimulatedObjectStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> Result { + let _permit = self.connection().await; + let bytes = payload.content_length() as u64; + let cost = self.get_ttfb.sample() + + transfer_time(bytes, self.config.connection_bytes_per_second); + tokio::time::sleep(cost).await; + self.inner.put_opts(location, payload, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> Result> { + // Only the "create upload" request is charged; the parts go through the + // returned `MultipartUpload`, which is not wrapped. Benchmarks read far + // more than they write, so this has not been worth the extra machinery. + let _permit = self.connection().await; + tokio::time::sleep(self.get_ttfb.sample()).await; + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> Result { + // A HEAD moves no bytes, but the `GetResult` it returns still reports a + // range spanning the whole object. Charging that range as a transfer + // would make every metadata probe cost as much as downloading the file. + let head = options.head; + let permit = self.connection().await; + tokio::time::sleep(self.get_ttfb.sample()).await; + let result = self.inner.get_opts(location, options).await?; + if head { + return Ok(result); + } + self.simulate_body(result, permit).await + } + + fn delete_stream( + &self, + locations: BoxStream<'static, Result>, + ) -> BoxStream<'static, Result> { + let ttfb = Arc::clone(&self.get_ttfb); + let connections = Arc::clone(&self.connections); + let locations = locations + .then(move |location| { + let ttfb = Arc::clone(&ttfb); + let connections = Arc::clone(&connections); + async move { + let _permit = connections.acquire_owned().await.expect(POOL_OPEN); + tokio::time::sleep(ttfb.sample()).await; + location + } + }) + .boxed(); + self.inner.delete_stream(locations) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result> { + let inner = self.inner.list(prefix); + let page_size = self.config.list_page_size; + let ttfb = Arc::clone(&self.list_ttfb); + let connections = Arc::clone(&self.connections); + + futures::stream::unfold((inner, 0usize), move |(mut inner, seen)| { + let ttfb = Arc::clone(&ttfb); + let connections = Arc::clone(&connections); + async move { + // S3 and GCS cap a listing response at `page_size` keys, and + // each page needs the continuation token from the page before + // it. Pages are therefore serial round trips, not a fan-out: + // listing 5000 files costs five times the latency of listing + // 500, which is why partition discovery hurts on remote stores. + if seen % page_size == 0 { + let _permit = connections.acquire_owned().await.expect(POOL_OPEN); + tokio::time::sleep(ttfb.sample()).await; + } + let item = inner.next().await?; + Some((item, (inner, seen + 1))) + } + }) + .boxed() + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { + // The real stores paginate this internally and hand back the assembled + // result, so the caller makes one call but pays for every page. The + // page count is only known after listing, so the cost is charged + // afterwards; the total is what matters. + let result = self.inner.list_with_delimiter(prefix).await?; + let entries = result.objects.len() + result.common_prefixes.len(); + let pages = entries.div_ceil(self.config.list_page_size).max(1); + for _ in 0..pages { + let _permit = self.connection().await; + tokio::time::sleep(self.list_ttfb.sample()).await; + } + Ok(result) + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + options: CopyOptions, + ) -> Result<()> { + // Server side copy: one round trip, no bytes over the wire. + let _permit = self.connection().await; + tokio::time::sleep(self.get_ttfb.sample()).await; + self.inner.copy_opts(from, to, options).await + } +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use bytes::Bytes; + use datafusion_common::instant::Instant; + use object_store::memory::InMemory; + use object_store::{ObjectStoreExt, PutPayload}; + + use super::*; + + /// Records the range of every `get_opts` reaching the inner store, which is + /// what a real store would turn into one HTTP request each. + #[derive(Debug, Default)] + struct RequestLog { + gets: AtomicUsize, + ranges: Mutex>, + } + + #[derive(Debug)] + struct Recording { + inner: InMemory, + log: Arc, + } + + impl fmt::Display for Recording { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "Recording") + } + } + + #[async_trait] + impl ObjectStore for Recording { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> Result { + self.inner.put_opts(location, payload, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> Result> { + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts( + &self, + location: &Path, + options: GetOptions, + ) -> Result { + self.log.gets.fetch_add(1, Ordering::Relaxed); + let result = self.inner.get_opts(location, options).await?; + self.log + .ranges + .lock() + .unwrap() + .push((result.range.start, result.range.end)); + Ok(result) + } + + fn delete_stream( + &self, + locations: BoxStream<'static, Result>, + ) -> BoxStream<'static, Result> { + self.inner.delete_stream(locations) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result> { + self.inner.list(prefix) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + options: CopyOptions, + ) -> Result<()> { + self.inner.copy_opts(from, to, options).await + } + } + + /// A config with no latency and effectively infinite bandwidth, so tests + /// assert on request *shape* without waiting. + fn instant() -> SimulatedStoreConfig { + SimulatedStoreConfig { + get_ttfb_ms: &[0], + list_ttfb_ms: &[0], + connection_bytes_per_second: u64::MAX, + ..Default::default() + } + } + + async fn store_with( + config: SimulatedStoreConfig, + len: usize, + ) -> (SimulatedObjectStore, Arc, Path) { + let log = Arc::new(RequestLog::default()); + let inner = Recording { + inner: InMemory::new(), + log: Arc::clone(&log), + }; + let path = Path::from("data.parquet"); + inner + .put(&path, PutPayload::from(vec![7u8; len])) + .await + .unwrap(); + (SimulatedObjectStore::with_config(inner, config), log, path) + } + + #[tokio::test] + async fn distant_ranges_become_separate_requests() { + // Ranges more than OBJECT_STORE_COALESCE_DEFAULT (1MiB) apart are + // separate GETs against S3. A store that overrode `get_ranges` would + // report one. + let (store, log, path) = store_with(instant(), 8 * 1024 * 1024).await; + let ranges = [0..1024, 3 * 1024 * 1024..3 * 1024 * 1024 + 1024]; + + let data = store.get_ranges(&path, &ranges).await.unwrap(); + + assert_eq!(data.len(), 2); + assert_eq!(log.gets.load(Ordering::Relaxed), 2); + } + + #[tokio::test] + async fn nearby_ranges_are_coalesced_into_one_request() { + // Under 1MiB apart, so arrow-rs merges them and reads the gap too. + let (store, log, path) = store_with(instant(), 8 * 1024 * 1024).await; + let ranges = [0..1024, 2048..4096]; + + store.get_ranges(&path, &ranges).await.unwrap(); + + assert_eq!(log.gets.load(Ordering::Relaxed), 1); + assert_eq!(log.ranges.lock().unwrap().as_slice(), &[(0, 4096)]); + } + + /// A Parquet file whose column chunks are far enough apart that arrow-rs + /// will not coalesce them: random `i64`s do not compress, so each chunk is + /// well over the 1MiB coalescing threshold. + fn wide_parquet(rows: usize, cols: usize) -> Bytes { + use arrow::array::{ArrayRef, Int64Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use parquet::arrow::ArrowWriter; + use rand::{Rng, SeedableRng, rngs::StdRng}; + + let mut rng = StdRng::seed_from_u64(42); + let schema = Arc::new(Schema::new( + (0..cols) + .map(|i| Field::new(format!("c{i}"), DataType::Int64, false)) + .collect::>(), + )); + let arrays: Vec = (0..cols) + .map(|_| { + let values: Vec = (0..rows).map(|_| rng.random()).collect(); + Arc::new(Int64Array::from(values)) as ArrayRef + }) + .collect(); + let batch = RecordBatch::try_new(Arc::clone(&schema), arrays).unwrap(); + + let mut buf = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + Bytes::from(buf) + } + + /// Run `sql` against a Parquet file served by the simulated store, and + /// return how many requests reached the wire during execution. + async fn requests_for_query(sql: &str) -> usize { + use datafusion::execution::object_store::ObjectStoreUrl; + use datafusion::prelude::{ParquetReadOptions, SessionContext}; + + let log = Arc::new(RequestLog::default()); + let recording = Recording { + inner: InMemory::new(), + log: Arc::clone(&log), + }; + recording + .put( + &Path::from("t.parquet"), + PutPayload::from(wide_parquet(200_000, 8)), + ) + .await + .unwrap(); + + let ctx = SessionContext::new(); + let url = ObjectStoreUrl::parse("mem://").unwrap(); + ctx.register_object_store( + url.as_ref(), + Arc::new(SimulatedObjectStore::with_config(recording, instant())), + ); + ctx.register_parquet("t", "mem:///t.parquet", ParquetReadOptions::default()) + .await + .unwrap(); + + // Registration infers the schema, which is IO of its own. Only count + // what the scan itself costs. + log.gets.store(0, Ordering::Relaxed); + ctx.sql(sql).await.unwrap().collect().await.unwrap(); + log.gets.load(Ordering::Relaxed) + } + + #[tokio::test] + async fn projecting_more_columns_costs_more_requests() { + // The whole point of the module: DataFusion asks for both column chunks + // in a single `get_byte_ranges` call, and the store must turn that into + // as many requests as S3 would. A store that overrode `get_ranges` and + // slept once would report the same count for both queries. + let one = requests_for_query("SELECT sum(c0) FROM t").await; + let two = requests_for_query("SELECT sum(c0), sum(c7) FROM t").await; + + assert!( + two > one, + "reading two distant column chunks ({two} requests) should cost more \ + than reading one ({one} requests)" + ); + } + + #[tokio::test] + async fn returned_ranges_are_correct() { + let (store, _log, path) = store_with(instant(), 4096).await; + let ranges = [10..20, 100..108, 3000..3001]; + + let data = store.get_ranges(&path, &ranges).await.unwrap(); + + assert_eq!( + data, + vec![ + Bytes::from(vec![7u8; 10]), + Bytes::from(vec![7u8; 8]), + Bytes::from(vec![7u8; 1]) + ] + ); + } + + #[tokio::test] + async fn head_requests_do_not_pay_for_a_body() { + let config = SimulatedStoreConfig { + get_ttfb_ms: &[0], + list_ttfb_ms: &[0], + // 1KiB/s, so charging the 1MiB object as a transfer would take + // roughly a quarter of an hour. + connection_bytes_per_second: 1024, + ..Default::default() + }; + let (store, _log, path) = store_with(config, 1024 * 1024).await; + + let start = Instant::now(); + store.head(&path).await.unwrap(); + + assert!( + start.elapsed() < Duration::from_secs(1), + "a HEAD moves no bytes and must not be charged as a transfer" + ); + } + + #[tokio::test] + async fn transfer_time_scales_with_bytes() { + let config = SimulatedStoreConfig { + get_ttfb_ms: &[0], + list_ttfb_ms: &[0], + // 1MiB/s, so a 256KiB read should take about a quarter second. + connection_bytes_per_second: 1024 * 1024, + ..Default::default() + }; + let (store, _log, path) = store_with(config, 1024 * 1024).await; + + let start = Instant::now(); + store.get_range(&path, 0..256 * 1024).await.unwrap(); + let elapsed = start.elapsed(); + + assert!( + elapsed >= Duration::from_millis(200), + "expected a bandwidth cost, took {elapsed:?}" + ); + } + + #[tokio::test] + async fn list_is_charged_per_page() { + let config = SimulatedStoreConfig { + get_ttfb_ms: &[0], + list_ttfb_ms: &[20], + list_page_size: 3, + connection_bytes_per_second: u64::MAX, + ..Default::default() + }; + let log = Arc::new(RequestLog::default()); + let inner = Recording { + inner: InMemory::new(), + log, + }; + for i in 0..7 { + inner + .put(&Path::from(format!("part-{i}")), PutPayload::from("x")) + .await + .unwrap(); + } + let store = SimulatedObjectStore::with_config(inner, config); + + let start = Instant::now(); + let listed: Vec<_> = store.list(None).collect().await; + let elapsed = start.elapsed(); + + assert_eq!(listed.len(), 7); + // 7 keys at 3 per page is 3 serial pages, so at least 60ms. + assert!( + elapsed >= Duration::from_millis(60), + "expected three serial LIST pages, took {elapsed:?}" + ); + } +} From af1b50b17a6abe4376963f7a7cf24429b97bbb4c Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:48:45 -0500 Subject: [PATCH 2/2] bench: use the simulated remote store for --simulate-latency Replaces `LatencyObjectStore`, which overrode `get_ranges` and slept once per call. That charged one round trip where S3 pays one per coalesced chunk, and it modelled no per-byte cost at all, so a 4 byte footer read and a 100MB column chunk read cost the same. It also charged a single latency for an entire LIST stream, where the real APIs paginate at 1000 keys. Numbers from previous --simulate-latency runs are not comparable across this change. Co-Authored-By: Claude Opus 5 --- benchmarks/bench.sh | 3 +- benchmarks/src/util/latency_object_store.rs | 157 -------------------- benchmarks/src/util/mod.rs | 2 +- benchmarks/src/util/options.rs | 25 ++-- 4 files changed, 19 insertions(+), 168 deletions(-) delete mode 100644 benchmarks/src/util/latency_object_store.rs diff --git a/benchmarks/bench.sh b/benchmarks/bench.sh index 52b78c844a73a..d723d8809eb87 100755 --- a/benchmarks/bench.sh +++ b/benchmarks/bench.sh @@ -165,7 +165,8 @@ CARGO_COMMAND command that runs the benchmark binary DATAFUSION_DIR directory to use (default $DATAFUSION_DIR) RESULTS_NAME folder where the benchmark files are stored PREFER_HASH_JOIN Prefer hash join algorithm (default true) -SIMULATE_LATENCY Simulate object store latency to mimic S3 (default false) +SIMULATE_LATENCY Serve local files as if they were remote object storage: per-request + latency, per-connection bandwidth and paginated LIST (default false) DATAFUSION_* Set the given datafusion configuration " exit 1 diff --git a/benchmarks/src/util/latency_object_store.rs b/benchmarks/src/util/latency_object_store.rs deleted file mode 100644 index 9ef8d1b78b751..0000000000000 --- a/benchmarks/src/util/latency_object_store.rs +++ /dev/null @@ -1,157 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! An ObjectStore wrapper that adds simulated S3-like latency to get and list operations. -//! -//! Cycles through a fixed latency distribution inspired by real S3 performance: -//! - P50: ~30ms -//! - P75-P90: ~100-120ms -//! - P99: ~150-200ms - -use std::fmt; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::Duration; - -use async_trait::async_trait; -use futures::StreamExt; -use futures::stream::BoxStream; -use object_store::path::Path; -use object_store::{ - CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, - ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result, -}; - -/// GET latency distribution, inspired by S3 latencies. -/// Deterministic but shuffled to avoid artificial patterns. -/// 20 values: 11x P50 (~25-35ms), 5x P75-P90 (~70-110ms), 2x P95 (~120-150ms), 2x P99 (~180-200ms) -/// Sorted: 25,25,28,28,30,30,30,30,32,32,35, 70,85,100,100,110, 130,150, 180,200 -/// P50≈32ms, P90≈110ms, P99≈200ms -const GET_LATENCIES_MS: &[u64] = &[ - 30, 100, 25, 85, 32, 200, 28, 130, 35, 70, 30, 150, 30, 110, 28, 180, 32, 25, 100, 30, -]; - -/// LIST latency distribution, generally higher than GET. -/// 20 values: 11x P50 (~40-70ms), 5x P75-P90 (~120-180ms), 2x P95 (~200-250ms), 2x P99 (~300-400ms) -/// Sorted: 40,40,50,50,55,55,60,60,65,65,70, 120,140,160,160,180, 210,250, 300,400 -/// P50≈65ms, P90≈180ms, P99≈400ms -const LIST_LATENCIES_MS: &[u64] = &[ - 55, 160, 40, 140, 65, 400, 50, 210, 70, 120, 60, 250, 55, 180, 50, 300, 65, 40, 160, - 60, -]; - -/// An ObjectStore wrapper that injects simulated latency on get and list calls. -#[derive(Debug)] -pub struct LatencyObjectStore { - inner: T, - get_counter: AtomicUsize, - list_counter: AtomicUsize, -} - -impl LatencyObjectStore { - pub fn new(inner: T) -> Self { - Self { - inner, - get_counter: AtomicUsize::new(0), - list_counter: AtomicUsize::new(0), - } - } - - fn next_get_latency(&self) -> Duration { - let idx = - self.get_counter.fetch_add(1, Ordering::Relaxed) % GET_LATENCIES_MS.len(); - Duration::from_millis(GET_LATENCIES_MS[idx]) - } - - fn next_list_latency(&self) -> Duration { - let idx = - self.list_counter.fetch_add(1, Ordering::Relaxed) % LIST_LATENCIES_MS.len(); - Duration::from_millis(LIST_LATENCIES_MS[idx]) - } -} - -impl fmt::Display for LatencyObjectStore { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "LatencyObjectStore({})", self.inner) - } -} - -#[async_trait] -impl ObjectStore for LatencyObjectStore { - async fn put_opts( - &self, - location: &Path, - payload: PutPayload, - opts: PutOptions, - ) -> Result { - self.inner.put_opts(location, payload, opts).await - } - - async fn put_multipart_opts( - &self, - location: &Path, - opts: PutMultipartOptions, - ) -> Result> { - self.inner.put_multipart_opts(location, opts).await - } - - async fn get_opts(&self, location: &Path, options: GetOptions) -> Result { - tokio::time::sleep(self.next_get_latency()).await; - self.inner.get_opts(location, options).await - } - - async fn get_ranges( - &self, - location: &Path, - ranges: &[std::ops::Range], - ) -> Result> { - tokio::time::sleep(self.next_get_latency()).await; - self.inner.get_ranges(location, ranges).await - } - - fn delete_stream( - &self, - locations: BoxStream<'static, Result>, - ) -> BoxStream<'static, Result> { - self.inner.delete_stream(locations) - } - - fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result> { - let latency = self.next_list_latency(); - let stream = self.inner.list(prefix); - futures::stream::once(async move { - tokio::time::sleep(latency).await; - futures::stream::empty() - }) - .flatten() - .chain(stream) - .boxed() - } - - async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { - tokio::time::sleep(self.next_list_latency()).await; - self.inner.list_with_delimiter(prefix).await - } - - async fn copy_opts( - &self, - from: &Path, - to: &Path, - options: CopyOptions, - ) -> Result<()> { - self.inner.copy_opts(from, to, options).await - } -} diff --git a/benchmarks/src/util/mod.rs b/benchmarks/src/util/mod.rs index 43855ea468ef5..12d03e45a1a77 100644 --- a/benchmarks/src/util/mod.rs +++ b/benchmarks/src/util/mod.rs @@ -16,11 +16,11 @@ // under the License. //! Shared benchmark utilities -pub mod latency_object_store; mod memory; mod memory_pool; mod options; mod run; +pub mod simulated_object_store; pub use memory::print_memory_stats; pub use memory_pool::PeakRecordingPool; diff --git a/benchmarks/src/util/options.rs b/benchmarks/src/util/options.rs index c744d0bf31c7f..a0911654b5f89 100644 --- a/benchmarks/src/util/options.rs +++ b/benchmarks/src/util/options.rs @@ -30,7 +30,9 @@ use datafusion::{ use datafusion_common::{DataFusionError, Result}; use object_store::local::LocalFileSystem; -use super::{latency_object_store::LatencyObjectStore, memory_pool::PeakRecordingPool}; +use super::{ + memory_pool::PeakRecordingPool, simulated_object_store::SimulatedObjectStore, +}; // Common benchmark options (don't use doc comments otherwise this doc // shows up in help files) @@ -66,8 +68,10 @@ pub struct CommonOpt { #[arg(short, long, env)] pub debug: bool, - /// Simulate object store latency to mimic remote storage (e.g. S3). - /// Adds random latency in the range 20-200ms to each object store operation. + /// Serve the local filesystem as if it were remote object storage (e.g. S3 + /// or GCS): per-request latency, per-connection bandwidth, a bounded + /// connection pool, and paginated LIST. See + /// [`crate::util::simulated_object_store`]. #[arg(long = "simulate-latency", env)] pub simulate_latency: bool, } @@ -135,17 +139,20 @@ impl CommonOpt { Ok(rt_builder) } - /// Build the runtime environment, optionally wrapping the local filesystem - /// with a throttled object store to simulate remote storage latency. + /// Build the runtime environment, optionally presenting the local + /// filesystem as remote object storage. pub fn build_runtime(&self) -> Result> { let rt = self.runtime_env_builder()?.build_arc()?; if self.simulate_latency { - let store: Arc = - Arc::new(LatencyObjectStore::new(LocalFileSystem::new())); + let store = SimulatedObjectStore::new(LocalFileSystem::new()); + let config = store.config().clone(); let url = ObjectStoreUrl::parse("file:///")?; - rt.register_object_store(url.as_ref(), store); + rt.register_object_store(url.as_ref(), Arc::new(store)); println!( - "Simulating S3-like object store latency (get: 25-200ms, list: 40-400ms)" + "Simulating remote object storage (get TTFB: 25-200ms, list TTFB: \ + 40-400ms/page, {} MB/s per connection, {} connections)", + config.connection_bytes_per_second / (1024 * 1024), + config.max_concurrent_requests, ); } Ok(rt)