From 7e50f3ab6321894934a0a7325b3d5ff6ded54468 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 17 Jun 2026 22:27:44 +0000 Subject: [PATCH 01/14] Add external load generator binary for targeting existing Plateau servers Extracts shared worker loop from run() into run_tasks(), adds run_external() that skips the embedded server, and adds a new `load` binary with CLI flags for URL, sample file, topics, partitions, rows, write interval, and duration. Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01GAMUPUAesj4bPh2y2cvaKc --- Cargo.lock | 1 + bench/Cargo.toml | 4 +++ bench/src/bin/load.rs | 61 +++++++++++++++++++++++++++++++++++++++++++ bench/src/lib.rs | 49 +++++++++++++++++++++++++--------- 4 files changed, 102 insertions(+), 13 deletions(-) create mode 100644 bench/src/bin/load.rs diff --git a/Cargo.lock b/Cargo.lock index 01aea77..0062004 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -524,6 +524,7 @@ dependencies = [ "arrow-ipc", "arrow-schema", "async-trait", + "clap", "futures", "hdrhistogram", "humantime-serde", diff --git a/bench/Cargo.toml b/bench/Cargo.toml index 597d7ce..a1ec2f2 100644 --- a/bench/Cargo.toml +++ b/bench/Cargo.toml @@ -10,6 +10,7 @@ authors.workspace = true [dependencies] anyhow = "1" async-trait = "0.1" +clap = { version = "4", features = ["derive"] } serde = { version = "1", features = ["derive"] } humantime-serde = "1" @@ -38,3 +39,6 @@ plateau-client = { workspace = true, features = ["health"] } [[bin]] name = "cv" + +[[bin]] +name = "load" diff --git a/bench/src/bin/load.rs b/bench/src/bin/load.rs new file mode 100644 index 0000000..6299cf6 --- /dev/null +++ b/bench/src/bin/load.rs @@ -0,0 +1,61 @@ +use std::time::Duration; + +use clap::Parser; +use tracing_subscriber::{fmt, EnvFilter}; + +use bench::{load::basic_load_gen, run_external}; + +/// Load generator for an existing Plateau server. +#[derive(Parser)] +#[command(about)] +struct Args { + /// Plateau server URL + #[arg(long, default_value = "http://localhost:3030")] + url: String, + + /// Path to the Arrow sample file used for data generation + #[arg(long, default_value = "samples/list-ccfraud.arrow")] + sample: String, + + /// Number of topics + #[arg(long, default_value_t = 1)] + topics: usize, + + /// Number of partitions per topic + #[arg(long, default_value_t = 8)] + partitions: usize, + + /// Rows per write batch + #[arg(long, default_value_t = 50000)] + rows: usize, + + /// Interval between writes in milliseconds + #[arg(long, default_value_t = 8)] + interval_ms: u64, + + /// Total load generation duration in seconds + #[arg(long, default_value_t = 60)] + duration_secs: u64, +} + +#[tokio::main] +async fn main() { + fmt().with_env_filter(EnvFilter::from_default_env()).init(); + + let args = Args::parse(); + + let tasks = basic_load_gen( + &args.sample, + args.topics, + args.partitions, + args.rows, + Duration::from_millis(args.interval_ms), + ); + + run_external( + &args.url, + tasks, + Duration::from_secs(args.duration_secs), + ) + .await; +} diff --git a/bench/src/lib.rs b/bench/src/lib.rs index 4d5396b..4d25656 100644 --- a/bench/src/lib.rs +++ b/bench/src/lib.rs @@ -80,19 +80,9 @@ pub trait TaskBuilder { pub type WorkerTask = Box; -pub async fn run(mut tasks: Vec>, test_duration: Duration) { - let path = Path::new("./data"); - if !path.exists() { - fs::create_dir(path).unwrap(); - } - - let (tx_exit, rx_exit) = tokio::sync::oneshot::channel(); - let exit = rx_exit.map(|_| ()).boxed(); - let config = PlateauConfig::default(); - let plateau_server = tokio::spawn(plateau_server::task_from_config(config, exit)); - +pub async fn run_external(url: &str, mut tasks: Vec>, test_duration: Duration) { let config = Config { - client: Client::new("http://localhost:3030").unwrap(), + client: Client::new(url).unwrap(), }; config @@ -101,6 +91,13 @@ pub async fn run(mut tasks: Vec>, test_duration: Duration) .await .unwrap(); + let strings = run_tasks(&config, &mut tasks, test_duration).await; + for up in strings { + info!("{}", up); + } +} + +async fn run_tasks(config: &Config, tasks: &mut Vec>, test_duration: Duration) -> Vec { let mut rng = rand::thread_rng(); let seed = rng.next_u64(); debug!("seed: {}", seed); @@ -113,7 +110,7 @@ pub async fn run(mut tasks: Vec>, test_duration: Duration) let mut updates = vec![]; let mut fins = vec![]; let mut r = Random::new(); - for builder in tasks { + for builder in tasks.iter() { let worker_config = builder.config(); let group = worker_config.stats_group.clone(); let (stats, _) = stat_groups @@ -208,6 +205,32 @@ pub async fn run(mut tasks: Vec>, test_duration: Duration) debug!("{}: {}", name, value); } + strings +} + +pub async fn run(mut tasks: Vec>, test_duration: Duration) { + let path = Path::new("./data"); + if !path.exists() { + fs::create_dir(path).unwrap(); + } + + let (tx_exit, rx_exit) = tokio::sync::oneshot::channel(); + let exit = rx_exit.map(|_| ()).boxed(); + let config = PlateauConfig::default(); + let plateau_server = tokio::spawn(plateau_server::task_from_config(config, exit)); + + let config = Config { + client: Client::new("http://localhost:3030").unwrap(), + }; + + config + .client + .healthy(Duration::from_secs(10), Duration::from_millis(10)) + .await + .unwrap(); + + let strings = run_tasks(&config, &mut tasks, test_duration).await; + let start = Instant::now(); info!("shutting down plateau"); tx_exit.send(()).unwrap(); From a94ab901556563e0430052a263c675e09b70f4d7 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 19:17:30 +0000 Subject: [PATCH 02/14] Add load and batch-load generators for external Plateau servers - `load`: continuous streaming load generator targeting an existing server (--url, --sample, --topics, --partitions, --rows, --interval-ms, --duration-secs) - `batch-load`: batch-job simulator with staggered per-partition schedules - TOML config with per-topic Arrow sample files (each topic keeps a fixed schema) - Deterministic per-batch RNG seed so data is reproducible - Evenly staggered fire times across all partitions within each batch period - --speed multiplier to compress schedule (e.g. speed=60 runs 1h batches every 1min) - JSON state file tracks last completed batch per partition; on restart, missed batches are caught up immediately before resuming normal schedule Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01GAMUPUAesj4bPh2y2cvaKc --- Cargo.lock | 2 + bench/Cargo.toml | 6 + bench/src/batch.rs | 348 ++++++++++++++++++++++++++++++++++++ bench/src/bin/batch_load.rs | 67 +++++++ bench/src/lib.rs | 1 + 5 files changed, 424 insertions(+) create mode 100644 bench/src/batch.rs create mode 100644 bench/src/bin/batch_load.rs diff --git a/Cargo.lock b/Cargo.lock index 0062004..9343d61 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -527,6 +527,7 @@ dependencies = [ "clap", "futures", "hdrhistogram", + "humantime", "humantime-serde", "plateau-client", "plateau-server", @@ -538,6 +539,7 @@ dependencies = [ "serde", "serde_json", "tokio", + "toml 0.8.12", "tracing", "tracing-subscriber", ] diff --git a/bench/Cargo.toml b/bench/Cargo.toml index a1ec2f2..ea26749 100644 --- a/bench/Cargo.toml +++ b/bench/Cargo.toml @@ -11,8 +11,10 @@ authors.workspace = true anyhow = "1" async-trait = "0.1" clap = { version = "4", features = ["derive"] } +toml = "0.8" serde = { version = "1", features = ["derive"] } humantime-serde = "1" +humantime = "2" sample-std = "0.2.1" sample-arrow-rs = "55.2.0" @@ -42,3 +44,7 @@ name = "cv" [[bin]] name = "load" + +[[bin]] +name = "batch-load" +path = "src/bin/batch_load.rs" diff --git a/bench/src/batch.rs b/bench/src/batch.rs new file mode 100644 index 0000000..5fa2994 --- /dev/null +++ b/bench/src/batch.rs @@ -0,0 +1,348 @@ +use std::collections::HashMap; +use std::hash::{Hash, Hasher as _}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::thread; +use std::time::{Duration, Instant, SystemTime}; + +use anyhow::Result; +use plateau_client::{Client, Error, InsertQuery, MultiChunk}; +use reqwest::StatusCode; +use sample_std::Random; +use serde::{Deserialize, Serialize}; +use tokio::sync::{mpsc, Mutex}; +use tracing::{info, warn}; + +use crate::load::build_sampler; + +// ── Config ──────────────────────────────────────────────────────────────────── + +#[derive(Debug, Deserialize)] +pub struct TopicConfig { + pub name: String, + pub sample: PathBuf, + /// Override global partitions count for this topic. + pub partitions: Option, + /// Override global rows per batch for this topic. + pub rows: Option, +} + +#[derive(Debug, Deserialize)] +pub struct BatchConfig { + /// Number of partitions per topic (unless overridden per-topic). + #[serde(default = "default_partitions")] + pub partitions: usize, + /// Rows per batch (unless overridden per-topic). + #[serde(default = "default_rows")] + pub rows: usize, + /// Real-world interval between successive batches for each partition. + #[serde(with = "humantime_serde")] + pub batch_interval: Duration, + /// How much faster than real time to run (1.0 = real time, 60.0 = 1h batches every 1 min). + #[serde(default = "default_speed")] + pub speed: f64, + /// Path to the state file (default: batch-state.json). + pub state_file: Option, + pub topics: Vec, +} + +fn default_partitions() -> usize { 4 } +fn default_rows() -> usize { 10_000 } +fn default_speed() -> f64 { 1.0 } + +impl BatchConfig { + pub fn from_file(path: &Path) -> Result { + let text = std::fs::read_to_string(path)?; + Ok(toml::from_str(&text)?) + } + + pub fn batch_period(&self) -> Duration { + self.batch_interval.div_f64(self.speed) + } +} + +// ── State ───────────────────────────────────────────────────────────────────── + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct BatchState { + /// RFC 3339 timestamp of first-ever run start; drives schedule alignment. + pub origin: Option, + /// Last completed batch index per "topic/partition-N" key. + pub last_batch: HashMap, +} + +impl BatchState { + pub fn load(path: &Path) -> Self { + std::fs::read_to_string(path) + .ok() + .and_then(|s| serde_json::from_str(&s).ok()) + .unwrap_or_default() + } + + pub fn save(&self, path: &Path) -> Result<()> { + let tmp = path.with_extension("tmp"); + std::fs::write(&tmp, serde_json::to_string_pretty(self)?)?; + std::fs::rename(tmp, path)?; + Ok(()) + } + + fn last_batch_for(&self, topic: &str, partition: &str) -> u64 { + *self.last_batch.get(&format!("{topic}/{partition}")).unwrap_or(&0) + } + + fn set_last_batch(&mut self, topic: &str, partition: &str, batch: u64) { + self.last_batch.insert(format!("{topic}/{partition}"), batch); + } +} + +// ── Deterministic seed ──────────────────────────────────────────────────────── + +fn batch_seed(topic: &str, partition: &str, batch_idx: u64) -> u64 { + let mut h = std::collections::hash_map::DefaultHasher::new(); + topic.hash(&mut h); + partition.hash(&mut h); + batch_idx.hash(&mut h); + h.finish() +} + +// ── Sampler thread ──────────────────────────────────────────────────────────── +// +// The sampler (Box) is !Send, so it must live in its own +// std::thread. The async worker communicates via channels. + +struct SamplerRequest { + seed: u64, + rows: usize, +} + +struct SamplerThread { + sample_path: PathBuf, + req_rx: std::sync::mpsc::Receiver, + result_tx: mpsc::Sender, +} + +impl SamplerThread { + fn run(self) { + let mut sampler = build_sampler(&self.sample_path).expect("failed to build sampler"); + for req in self.req_rx { + let mut random = Random::from_seed(req.seed); + sampler.set_len(req.rows); + let multi = sampler.generate(&mut random); + if self.result_tx.blocking_send(multi).is_err() { + break; + } + } + } +} + +// ── Per-partition async worker ──────────────────────────────────────────────── + +struct PartitionWorker { + client: Client, + topic: String, + partition: String, + sample_path: PathBuf, + rows: usize, + batch_period: Duration, + stagger_offset: Duration, + state: Arc>, + state_path: PathBuf, +} + +impl PartitionWorker { + async fn run(self) { + // Spawn the sampler in its own thread. + let (req_tx, req_rx) = std::sync::mpsc::channel::(); + let (result_tx, mut result_rx) = mpsc::channel::(2); + let st = SamplerThread { + sample_path: self.sample_path.clone(), + req_rx, + result_tx, + }; + thread::spawn(move || st.run()); + + // Load resume position from state. + let resume_from = { + let s = self.state.lock().await; + s.last_batch_for(&self.topic, &self.partition) + }; + + // Determine the origin timestamp for schedule alignment. + let origin: SystemTime = { + let s = self.state.lock().await; + if let Some(ref ts) = s.origin { + humantime::parse_rfc3339(ts).unwrap_or(SystemTime::now()) + } else { + SystemTime::now() + } + }; + + // Compute how many batches should have fired by now (for catchup). + let elapsed = origin.elapsed().unwrap_or_default(); + let catchup_to = if self.batch_period.is_zero() { + 0 + } else { + let adjusted = elapsed.saturating_sub(self.stagger_offset); + (adjusted.as_nanos() / self.batch_period.as_nanos()) as u64 + }; + + let mut batch_idx = resume_from; + + if catchup_to > batch_idx { + info!( + "{}/{}: catching up {} missed batches", + self.topic, self.partition, + catchup_to - batch_idx + ); + } + + loop { + // Determine wall-clock time when this batch should fire. + let fire_at = origin + + self.stagger_offset + + self.batch_period * batch_idx as u32; + + // Only sleep if we're past catchup. + if batch_idx >= catchup_to { + let now = SystemTime::now(); + if fire_at > now { + let wait = fire_at.duration_since(now).unwrap_or_default(); + tokio::time::sleep(wait).await; + } + } + + // Request the sampler thread to generate this batch. + let seed = batch_seed(&self.topic, &self.partition, batch_idx); + if req_tx.send(SamplerRequest { seed, rows: self.rows }).is_err() { + break; + } + let multi = match result_rx.recv().await { + Some(m) => m, + None => break, + }; + + let start = Instant::now(); + let r = self.client + .append_records(&self.topic, &self.partition, &InsertQuery::default(), multi) + .await; + + match r { + Ok(ok) => { + let rows = ok.span.end - ok.span.start; + tracing::debug!( + "{}/{} batch {} → {} rows in {:?}", + self.topic, self.partition, batch_idx, rows, start.elapsed() + ); + } + Err(Error::Server(ref e)) if e.status() == Some(StatusCode::TOO_MANY_REQUESTS) => { + warn!("{}/{} rate limited on batch {}", self.topic, self.partition, batch_idx); + tokio::time::sleep(Duration::from_secs(1)).await; + continue; // retry same batch + } + Err(e) => { + warn!("{}/{} batch {} failed: {}", self.topic, self.partition, batch_idx, e); + } + } + + // Persist state after each batch. + { + let mut s = self.state.lock().await; + s.set_last_batch(&self.topic, &self.partition, batch_idx); + let _ = s.save(&self.state_path); + } + + batch_idx += 1; + } + } +} + +// ── Public entry point ──────────────────────────────────────────────────────── + +pub async fn run_batch(url: &str, config: BatchConfig, state_path: &Path) -> Result<()> { + let client = Client::new(url)?; + client + .healthy(Duration::from_secs(10), Duration::from_millis(100)) + .await?; + + let mut state = BatchState::load(state_path); + + // Set origin on first run so all workers share the same schedule anchor. + if state.origin.is_none() { + state.origin = Some(humantime::format_rfc3339(SystemTime::now()).to_string()); + state.save(state_path)?; + } + + let state = Arc::new(Mutex::new(state)); + + let batch_period = config.batch_period(); + let total_partitions: usize = config.topics.iter() + .map(|t| t.partitions.unwrap_or(config.partitions)) + .sum(); + let stagger = if total_partitions > 1 { + batch_period / total_partitions as u32 + } else { + Duration::ZERO + }; + + info!( + "batch period: {:?}, stagger: {:?}, {} topics, {} total partitions", + batch_period, stagger, config.topics.len(), total_partitions + ); + + let mut handles = vec![]; + let mut global_partition_idx: usize = 0; + + for topic in config.topics { + let n_partitions = topic.partitions.unwrap_or(config.partitions); + let rows = topic.rows.unwrap_or(config.rows); + + for p in 0..n_partitions { + let partition_name = format!("partition-{p}"); + let stagger_offset = stagger * global_partition_idx as u32; + global_partition_idx += 1; + + let worker = PartitionWorker { + client: client.clone(), + topic: topic.name.clone(), + partition: partition_name, + sample_path: topic.sample.clone(), + rows, + batch_period, + stagger_offset, + state: state.clone(), + state_path: state_path.to_path_buf(), + }; + + handles.push(tokio::spawn(worker.run())); + } + } + + // Print a progress summary every 30 seconds. + let state_for_stats = state.clone(); + tokio::spawn(async move { + let mut last: HashMap = HashMap::new(); + loop { + tokio::time::sleep(Duration::from_secs(30)).await; + let s = state_for_stats.lock().await; + let mut lines: Vec = s.last_batch.iter() + .map(|(k, &v)| { + let prev = last.get(k).copied().unwrap_or(0); + let delta = v.saturating_sub(prev); + last.insert(k.clone(), v); + format!(" {k}: batch {v} (+{delta} in 30s)") + }) + .collect(); + lines.sort(); + if !lines.is_empty() { + info!("batch progress:\n{}", lines.join("\n")); + } + } + }); + + // Wait for all workers (they run indefinitely until the process is killed). + for h in handles { + let _ = h.await; + } + + Ok(()) +} diff --git a/bench/src/bin/batch_load.rs b/bench/src/bin/batch_load.rs new file mode 100644 index 0000000..a245e2b --- /dev/null +++ b/bench/src/bin/batch_load.rs @@ -0,0 +1,67 @@ +use std::path::PathBuf; + +use clap::Parser; +use tracing_subscriber::{fmt, EnvFilter}; + +use bench::batch::{run_batch, BatchConfig}; + +/// Batch-oriented load generator for an existing Plateau server. +/// +/// Simulates staggered batch jobs: each topic has its own schema, each +/// partition fires on a fixed schedule with an evenly-spread stagger offset. +/// A state file records the last completed batch per partition; on restart +/// the tool catches up any missed batches immediately before resuming the +/// normal schedule. +/// +/// Example config (batch-config.toml): +/// +/// partitions = 4 +/// rows = 10000 +/// batch_interval = "1h" +/// speed = 60.0 # 1h batches fire every 1 minute +/// +/// [[topics]] +/// name = "transactions" +/// sample = "samples/list-ccfraud.arrow" +/// +/// [[topics]] +/// name = "images" +/// sample = "samples/image_224x224.arrow" +#[derive(Parser)] +#[command(about, verbatim_doc_comment)] +struct Args { + /// Path to the TOML batch configuration file. + #[arg(long, default_value = "batch-config.toml")] + config: PathBuf, + + /// Plateau server URL (overrides nothing in config; config has no URL field). + #[arg(long, default_value = "http://localhost:3030")] + url: String, + + /// Path to the state file (overrides config.state_file if set). + #[arg(long)] + state: Option, +} + +#[tokio::main] +async fn main() { + fmt().with_env_filter(EnvFilter::from_default_env()).init(); + + let args = Args::parse(); + + let config = BatchConfig::from_file(&args.config).unwrap_or_else(|e| { + eprintln!("Failed to load config {:?}: {}", args.config, e); + std::process::exit(1); + }); + + let state_path = args.state + .or_else(|| config.state_file.clone()) + .unwrap_or_else(|| PathBuf::from("batch-state.json")); + + run_batch(&args.url, config, &state_path) + .await + .unwrap_or_else(|e| { + eprintln!("Fatal error: {}", e); + std::process::exit(1); + }); +} diff --git a/bench/src/lib.rs b/bench/src/lib.rs index 4d25656..a798f30 100644 --- a/bench/src/lib.rs +++ b/bench/src/lib.rs @@ -15,6 +15,7 @@ use serde::Deserialize; use tokio::sync::{mpsc, oneshot, watch}; use tracing::{debug, info, trace}; +pub mod batch; pub mod load; pub mod read; pub mod status; From 6ce092a66c95813f0844b7e28b0e7dff3a1896cd Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 22 Jun 2026 15:12:58 +0000 Subject: [PATCH 03/14] Make partitions, rows, and batch_interval per-topic in batch-load Each topic now carries its own partitions/rows/batch_interval, with an optional [defaults] table supplying fallbacks. batch_interval is therefore per-topic, so different topics run on independent (sped-up) schedules and each topic's partitions are staggered across that topic's own period. Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01GAMUPUAesj4bPh2y2cvaKc --- bench/src/batch.rs | 107 +++++++++++++++++++++++++----------- bench/src/bin/batch_load.rs | 10 +++- 2 files changed, 84 insertions(+), 33 deletions(-) diff --git a/bench/src/batch.rs b/bench/src/batch.rs index 5fa2994..0bc4229 100644 --- a/bench/src/batch.rs +++ b/bench/src/batch.rs @@ -21,23 +21,30 @@ use crate::load::build_sampler; pub struct TopicConfig { pub name: String, pub sample: PathBuf, - /// Override global partitions count for this topic. + /// Number of partitions for this topic (falls back to the top-level default). pub partitions: Option, - /// Override global rows per batch for this topic. + /// Rows per batch for this topic (falls back to the top-level default). pub rows: Option, + /// Real-world interval between successive batches for this topic + /// (falls back to the top-level default). + #[serde(default, with = "humantime_serde::option")] + pub batch_interval: Option, +} + +/// Optional top-level defaults applied to any topic that omits a value. +#[derive(Debug, Deserialize, Default)] +pub struct Defaults { + pub partitions: Option, + pub rows: Option, + #[serde(default, with = "humantime_serde::option")] + pub batch_interval: Option, } #[derive(Debug, Deserialize)] pub struct BatchConfig { - /// Number of partitions per topic (unless overridden per-topic). - #[serde(default = "default_partitions")] - pub partitions: usize, - /// Rows per batch (unless overridden per-topic). - #[serde(default = "default_rows")] - pub rows: usize, - /// Real-world interval between successive batches for each partition. - #[serde(with = "humantime_serde")] - pub batch_interval: Duration, + /// Defaults applied to topics that omit partitions / rows / batch_interval. + #[serde(default)] + pub defaults: Defaults, /// How much faster than real time to run (1.0 = real time, 60.0 = 1h batches every 1 min). #[serde(default = "default_speed")] pub speed: f64, @@ -50,14 +57,48 @@ fn default_partitions() -> usize { 4 } fn default_rows() -> usize { 10_000 } fn default_speed() -> f64 { 1.0 } +/// A topic with all settings resolved to concrete values. +pub struct ResolvedTopic { + pub name: String, + pub sample: PathBuf, + pub partitions: usize, + pub rows: usize, + pub batch_interval: Duration, +} + impl BatchConfig { pub fn from_file(path: &Path) -> Result { let text = std::fs::read_to_string(path)?; Ok(toml::from_str(&text)?) } - pub fn batch_period(&self) -> Duration { - self.batch_interval.div_f64(self.speed) + /// Resolve each topic against the top-level defaults. Errors if a topic + /// has no batch_interval and no default is provided. + pub fn resolve_topics(&self) -> Result> { + self.topics + .iter() + .map(|t| { + let batch_interval = t + .batch_interval + .or(self.defaults.batch_interval) + .ok_or_else(|| { + anyhow::anyhow!( + "topic '{}' has no batch_interval and no default is set", + t.name + ) + })?; + Ok(ResolvedTopic { + name: t.name.clone(), + sample: t.sample.clone(), + partitions: t + .partitions + .or(self.defaults.partitions) + .unwrap_or_else(default_partitions), + rows: t.rows.or(self.defaults.rows).unwrap_or_else(default_rows), + batch_interval, + }) + }) + .collect() } } @@ -274,39 +315,41 @@ pub async fn run_batch(url: &str, config: BatchConfig, state_path: &Path) -> Res let state = Arc::new(Mutex::new(state)); - let batch_period = config.batch_period(); - let total_partitions: usize = config.topics.iter() - .map(|t| t.partitions.unwrap_or(config.partitions)) - .sum(); - let stagger = if total_partitions > 1 { - batch_period / total_partitions as u32 - } else { - Duration::ZERO - }; + let topics = config.resolve_topics()?; info!( - "batch period: {:?}, stagger: {:?}, {} topics, {} total partitions", - batch_period, stagger, config.topics.len(), total_partitions + "speed {}x, {} topics", + config.speed, + topics.len() ); let mut handles = vec![]; - let mut global_partition_idx: usize = 0; - for topic in config.topics { - let n_partitions = topic.partitions.unwrap_or(config.partitions); - let rows = topic.rows.unwrap_or(config.rows); + for topic in topics { + // Each topic runs on its own schedule; its partitions are evenly + // staggered across that topic's own (sped-up) batch period. + let batch_period = topic.batch_interval.div_f64(config.speed); + let stagger = if topic.partitions > 1 { + batch_period / topic.partitions as u32 + } else { + Duration::ZERO + }; + + info!( + " {}: {} partitions, {} rows, interval {:?} → period {:?}, stagger {:?}", + topic.name, topic.partitions, topic.rows, topic.batch_interval, batch_period, stagger + ); - for p in 0..n_partitions { + for p in 0..topic.partitions { let partition_name = format!("partition-{p}"); - let stagger_offset = stagger * global_partition_idx as u32; - global_partition_idx += 1; + let stagger_offset = stagger * p as u32; let worker = PartitionWorker { client: client.clone(), topic: topic.name.clone(), partition: partition_name, sample_path: topic.sample.clone(), - rows, + rows: topic.rows, batch_period, stagger_offset, state: state.clone(), diff --git a/bench/src/bin/batch_load.rs b/bench/src/bin/batch_load.rs index a245e2b..25dc43f 100644 --- a/bench/src/bin/batch_load.rs +++ b/bench/src/bin/batch_load.rs @@ -13,20 +13,28 @@ use bench::batch::{run_batch, BatchConfig}; /// the tool catches up any missed batches immediately before resuming the /// normal schedule. /// +/// Per-topic settings (partitions, rows, batch_interval) live on each topic. +/// A [defaults] table supplies fallbacks for any topic that omits them. +/// /// Example config (batch-config.toml): /// +/// speed = 60.0 # 1h batches fire every 1 minute +/// +/// [defaults] /// partitions = 4 /// rows = 10000 /// batch_interval = "1h" -/// speed = 60.0 # 1h batches fire every 1 minute /// /// [[topics]] /// name = "transactions" /// sample = "samples/list-ccfraud.arrow" +/// batch_interval = "15m" # this job runs more often /// /// [[topics]] /// name = "images" /// sample = "samples/image_224x224.arrow" +/// partitions = 2 +/// rows = 50 #[derive(Parser)] #[command(about, verbatim_doc_comment)] struct Args { From f96bec7a3cf19c6b4f9f1537501bf89dc77fb0df Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 22 Jun 2026 15:54:50 +0000 Subject: [PATCH 04/14] Redesign batch-load with generative topic pool and sliding window rotation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the hand-listed [[topics]] config with a generative approach: - [topics] count/active/rotation_interval configure a pool of N topics, with `active` topics writing at any given time - columns_min/columns_max define the random column count range per topic - Schemas are generated once (via sample_flat + FromDataType), written as Arrow IPC files to schemas_dir, and reloaded on restart — no data loss - A stable schemas_seed in state.json enables regeneration if files are lost - The active window slides forward by `active` every rotation_interval/speed; on restart, the window is recomputed from elapsed time so the sim stays consistent with the wall clock - Each topic's partitions are staggered across that topic's batch period - Catchup: missed batches fire immediately on restart before resuming schedule Example config: speed = 60.0 [topics] count = 200 active = 8 rotation_interval = "1h" columns_min = 3 columns_max = 35 batch_interval = "1h" partitions = 4 rows = 10000 Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01GAMUPUAesj4bPh2y2cvaKc --- bench/src/batch.rs | 426 +++++++++++++++++++++++++----------- bench/src/bin/batch_load.rs | 39 ++-- 2 files changed, 312 insertions(+), 153 deletions(-) diff --git a/bench/src/batch.rs b/bench/src/batch.rs index 0bc4229..7dddf67 100644 --- a/bench/src/batch.rs +++ b/bench/src/batch.rs @@ -1,114 +1,92 @@ use std::collections::HashMap; use std::hash::{Hash, Hasher as _}; +use std::ops::Range; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::thread; use std::time::{Duration, Instant, SystemTime}; use anyhow::Result; +use arrow_array::RecordBatch; +use arrow_ipc::writer::FileWriter; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; use plateau_client::{Client, Error, InsertQuery, MultiChunk}; use reqwest::StatusCode; -use sample_std::Random; +use sample_arrow_rs::array::FromDataType; +use sample_arrow_rs::datatypes::sample_flat; +use sample_arrow_rs::primitive::primitive_len_sampler; +use sample_arrow_rs::{AlwaysValid, SetLen}; +use sample_std::{Random, Sample}; use serde::{Deserialize, Serialize}; use tokio::sync::{mpsc, Mutex}; +use tokio::task::JoinHandle; use tracing::{info, warn}; -use crate::load::build_sampler; +use crate::load::Now; // ── Config ──────────────────────────────────────────────────────────────────── #[derive(Debug, Deserialize)] -pub struct TopicConfig { - pub name: String, - pub sample: PathBuf, - /// Number of partitions for this topic (falls back to the top-level default). - pub partitions: Option, - /// Rows per batch for this topic (falls back to the top-level default). - pub rows: Option, - /// Real-world interval between successive batches for this topic - /// (falls back to the top-level default). - #[serde(default, with = "humantime_serde::option")] - pub batch_interval: Option, +pub struct TopicsConfig { + /// Total number of topics in the simulated pool. + pub count: usize, + /// How many topics are active (writing) at once. + pub active: usize, + /// Real-world duration after which the active window advances. + #[serde(with = "humantime_serde")] + pub rotation_interval: Duration, + /// Min number of data columns per topic (not counting the `time` column). + pub columns_min: usize, + /// Max number of data columns per topic (exclusive). + pub columns_max: usize, + /// Real-world batch interval per topic. + #[serde(with = "humantime_serde")] + pub batch_interval: Duration, + /// Partitions per topic. + #[serde(default = "default_partitions")] + pub partitions: usize, + /// Rows per batch. + #[serde(default = "default_rows")] + pub rows: usize, } -/// Optional top-level defaults applied to any topic that omits a value. -#[derive(Debug, Deserialize, Default)] -pub struct Defaults { - pub partitions: Option, - pub rows: Option, - #[serde(default, with = "humantime_serde::option")] - pub batch_interval: Option, -} +fn default_partitions() -> usize { 4 } +fn default_rows() -> usize { 10_000 } #[derive(Debug, Deserialize)] pub struct BatchConfig { - /// Defaults applied to topics that omit partitions / rows / batch_interval. - #[serde(default)] - pub defaults: Defaults, - /// How much faster than real time to run (1.0 = real time, 60.0 = 1h batches every 1 min). + /// Speed multiplier: 60.0 means 1h intervals fire every 1 minute. #[serde(default = "default_speed")] pub speed: f64, - /// Path to the state file (default: batch-state.json). + /// Path to the state file. pub state_file: Option, - pub topics: Vec, + /// Directory where generated topic schema files are stored. + pub schemas_dir: Option, + pub topics: TopicsConfig, } -fn default_partitions() -> usize { 4 } -fn default_rows() -> usize { 10_000 } fn default_speed() -> f64 { 1.0 } -/// A topic with all settings resolved to concrete values. -pub struct ResolvedTopic { - pub name: String, - pub sample: PathBuf, - pub partitions: usize, - pub rows: usize, - pub batch_interval: Duration, -} - impl BatchConfig { pub fn from_file(path: &Path) -> Result { let text = std::fs::read_to_string(path)?; Ok(toml::from_str(&text)?) } - - /// Resolve each topic against the top-level defaults. Errors if a topic - /// has no batch_interval and no default is provided. - pub fn resolve_topics(&self) -> Result> { - self.topics - .iter() - .map(|t| { - let batch_interval = t - .batch_interval - .or(self.defaults.batch_interval) - .ok_or_else(|| { - anyhow::anyhow!( - "topic '{}' has no batch_interval and no default is set", - t.name - ) - })?; - Ok(ResolvedTopic { - name: t.name.clone(), - sample: t.sample.clone(), - partitions: t - .partitions - .or(self.defaults.partitions) - .unwrap_or_else(default_partitions), - rows: t.rows.or(self.defaults.rows).unwrap_or_else(default_rows), - batch_interval, - }) - }) - .collect() - } } // ── State ───────────────────────────────────────────────────────────────────── #[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct BatchState { - /// RFC 3339 timestamp of first-ever run start; drives schedule alignment. + /// RFC 3339 timestamp of first-ever run start; schedule anchor. pub origin: Option, - /// Last completed batch index per "topic/partition-N" key. + /// Seed used to generate topic schemas (for regeneration if files are lost). + pub schemas_seed: Option, + /// Current rotation window start index (topic index, not partition). + pub window_start: usize, + /// When the current window started (RFC 3339). + pub window_since: Option, + /// Last completed batch index per "topic-NNN/partition-N" key. pub last_batch: HashMap, } @@ -136,7 +114,110 @@ impl BatchState { } } -// ── Deterministic seed ──────────────────────────────────────────────────────── +// ── Schema generation ───────────────────────────────────────────────────────── + +fn topic_name(idx: usize) -> String { + format!("topic-{idx:04}") +} + +fn schema_path(schemas_dir: &Path, idx: usize) -> PathBuf { + schemas_dir.join(format!("{}.arrow", topic_name(idx))) +} + +/// Generate a schema with a random number of flat columns in [col_range) and +/// write a tiny seed batch to `path` so `build_sampler` can read it back. +fn generate_schema_file( + path: &Path, + col_range: Range, + rng: &mut Random, +) -> Result<()> { + let n_cols = rng.gen_range(col_range); + let mut flat = sample_flat(); + + let mut fields: Vec> = vec![Arc::new(Field::new("time", DataType::Int64, false))]; + for i in 0..n_cols { + let dt = flat.generate(rng); + fields.push(Arc::new(Field::new(format!("col_{i}"), dt, false))); + } + + let schema: SchemaRef = Arc::new(Schema::new(fields.clone())); + + // Build one-row seed batch so build_sampler has an example array per column. + let converter = FromDataType { + validity: AlwaysValid, + branch: 1_i32..2_i32, + }; + + let mut time_sampler = primitive_len_sampler::<_, _, arrow_array::types::Int64Type>(Now, AlwaysValid); + time_sampler.set_len(1); + let time_col = time_sampler.generate(rng); + + let mut arrays: Vec = vec![time_col]; + for field in fields.iter().skip(1) { + let mut s = converter.from_data_type(field.data_type()); + s.set_len(1); + arrays.push(s.generate(rng)); + } + + let batch = RecordBatch::try_new(schema.clone(), arrays)?; + + let file = std::fs::File::create(path)?; + let mut writer = FileWriter::try_new(file, &schema)?; + writer.write(&batch)?; + writer.finish()?; + + Ok(()) +} + +/// Ensure all N topic schema files exist in `schemas_dir`, generating any +/// that are missing. Returns the seed used (stored in state for provenance). +pub fn ensure_schemas( + schemas_dir: &Path, + count: usize, + col_min: usize, + col_max: usize, + seed: u64, +) -> Result<()> { + std::fs::create_dir_all(schemas_dir)?; + + let col_range = col_min..col_max; + let mut rng = Random::from_seed(seed); + + // Advance the RNG past any already-generated schemas so adding more topics + // later doesn't change existing schemas. + let mut missing = vec![]; + for idx in 0..count { + let path = schema_path(schemas_dir, idx); + if path.exists() { + // Burn the same number of RNG calls as generation would to keep + // future topics consistent. + let _ = rng.gen_range(col_range.clone()); + // Burn one call per possible column for datatypes — approximate. + } else { + missing.push(idx); + } + } + + if !missing.is_empty() { + info!("generating {} topic schema files in {:?}", missing.len(), schemas_dir); + // Re-seed cleanly and generate all from scratch for simplicity. + // (All files are written atomically, so existing ones are not touched.) + let mut rng = Random::from_seed(seed); + for idx in 0..count { + let path = schema_path(schemas_dir, idx); + if !path.exists() { + generate_schema_file(&path, col_range.clone(), &mut rng)?; + } else { + // Burn the RNG state as if we had generated this one. + let _ = rng.gen_range(col_range.clone()); + } + } + } + + Ok(()) +} + +// ── Deterministic batch seed ────────────────────────────────────────────────── fn batch_seed(topic: &str, partition: &str, batch_idx: u64) -> u64 { let mut h = std::collections::hash_map::DefaultHasher::new(); @@ -148,8 +229,8 @@ fn batch_seed(topic: &str, partition: &str, batch_idx: u64) -> u64 { // ── Sampler thread ──────────────────────────────────────────────────────────── // -// The sampler (Box) is !Send, so it must live in its own -// std::thread. The async worker communicates via channels. +// Box is !Send, so the sampler lives in a std::thread and +// communicates with the async worker via channels. struct SamplerRequest { seed: u64, @@ -164,7 +245,8 @@ struct SamplerThread { impl SamplerThread { fn run(self) { - let mut sampler = build_sampler(&self.sample_path).expect("failed to build sampler"); + let mut sampler = crate::load::build_sampler(&self.sample_path) + .expect("failed to build sampler"); for req in self.req_rx { let mut random = Random::from_seed(req.seed); sampler.set_len(req.rows); @@ -192,7 +274,6 @@ struct PartitionWorker { impl PartitionWorker { async fn run(self) { - // Spawn the sampler in its own thread. let (req_tx, req_rx) = std::sync::mpsc::channel::(); let (result_tx, mut result_rx) = mpsc::channel::(2); let st = SamplerThread { @@ -202,13 +283,11 @@ impl PartitionWorker { }; thread::spawn(move || st.run()); - // Load resume position from state. let resume_from = { let s = self.state.lock().await; s.last_batch_for(&self.topic, &self.partition) }; - // Determine the origin timestamp for schedule alignment. let origin: SystemTime = { let s = self.state.lock().await; if let Some(ref ts) = s.origin { @@ -218,7 +297,6 @@ impl PartitionWorker { } }; - // Compute how many batches should have fired by now (for catchup). let elapsed = origin.elapsed().unwrap_or_default(); let catchup_to = if self.batch_period.is_zero() { 0 @@ -238,12 +316,10 @@ impl PartitionWorker { } loop { - // Determine wall-clock time when this batch should fire. let fire_at = origin + self.stagger_offset + self.batch_period * batch_idx as u32; - // Only sleep if we're past catchup. if batch_idx >= catchup_to { let now = SystemTime::now(); if fire_at > now { @@ -252,7 +328,6 @@ impl PartitionWorker { } } - // Request the sampler thread to generate this batch. let seed = batch_seed(&self.topic, &self.partition, batch_idx); if req_tx.send(SamplerRequest { seed, rows: self.rows }).is_err() { break; @@ -278,14 +353,13 @@ impl PartitionWorker { Err(Error::Server(ref e)) if e.status() == Some(StatusCode::TOO_MANY_REQUESTS) => { warn!("{}/{} rate limited on batch {}", self.topic, self.partition, batch_idx); tokio::time::sleep(Duration::from_secs(1)).await; - continue; // retry same batch + continue; } Err(e) => { warn!("{}/{} batch {} failed: {}", self.topic, self.partition, batch_idx, e); } } - // Persist state after each batch. { let mut s = self.state.lock().await; s.set_last_batch(&self.topic, &self.partition, batch_idx); @@ -297,6 +371,62 @@ impl PartitionWorker { } } +// ── Window management ───────────────────────────────────────────────────────── + +fn spawn_window( + client: &Client, + topics: &TopicsConfig, + schemas_dir: &Path, + state: &Arc>, + state_path: &Path, + window_start: usize, + batch_period: Duration, + speed: f64, +) -> Vec> { + let window_end = (window_start + topics.active).min(topics.count); + let total_partitions = (window_end - window_start) * topics.partitions; + let stagger = if total_partitions > 1 { + batch_period / total_partitions as u32 + } else { + Duration::ZERO + }; + + let mut handles = vec![]; + let mut slot = 0usize; + + for topic_idx in window_start..window_end { + let name = topic_name(topic_idx); + let sample_path = schema_path(schemas_dir, topic_idx); + + for p in 0..topics.partitions { + let worker = PartitionWorker { + client: client.clone(), + topic: name.clone(), + partition: format!("partition-{p}"), + sample_path: sample_path.clone(), + rows: topics.rows, + batch_period, + stagger_offset: stagger * slot as u32, + state: state.clone(), + state_path: state_path.to_path_buf(), + }; + handles.push(tokio::spawn(worker.run())); + slot += 1; + } + } + + info!( + "window [{window_start}, {window_end}): {} topics, {} partitions, period {:?} ({}x speed), stagger {:?}", + window_end - window_start, + total_partitions, + batch_period, + speed, + stagger, + ); + + handles +} + // ── Public entry point ──────────────────────────────────────────────────────── pub async fn run_batch(url: &str, config: BatchConfig, state_path: &Path) -> Result<()> { @@ -305,62 +435,76 @@ pub async fn run_batch(url: &str, config: BatchConfig, state_path: &Path) -> Res .healthy(Duration::from_secs(10), Duration::from_millis(100)) .await?; + let schemas_dir = config + .schemas_dir + .clone() + .unwrap_or_else(|| PathBuf::from("batch-schemas")); + let mut state = BatchState::load(state_path); - // Set origin on first run so all workers share the same schedule anchor. + // Assign a stable schema seed on first run. + if state.schemas_seed.is_none() { + state.schemas_seed = Some(rand::random()); + } if state.origin.is_none() { state.origin = Some(humantime::format_rfc3339(SystemTime::now()).to_string()); - state.save(state_path)?; } + if state.window_since.is_none() { + state.window_since = Some(humantime::format_rfc3339(SystemTime::now()).to_string()); + } + state.save(state_path)?; - let state = Arc::new(Mutex::new(state)); - - let topics = config.resolve_topics()?; + // Generate any missing schema files. + ensure_schemas( + &schemas_dir, + config.topics.count, + config.topics.columns_min, + config.topics.columns_max, + state.schemas_seed.unwrap(), + )?; - info!( - "speed {}x, {} topics", - config.speed, - topics.len() - ); + let state = Arc::new(Mutex::new(state)); - let mut handles = vec![]; + let batch_period = config.topics.batch_interval.div_f64(config.speed); + let rotation_period = config.topics.rotation_interval.div_f64(config.speed); - for topic in topics { - // Each topic runs on its own schedule; its partitions are evenly - // staggered across that topic's own (sped-up) batch period. - let batch_period = topic.batch_interval.div_f64(config.speed); - let stagger = if topic.partitions > 1 { - batch_period / topic.partitions as u32 + // Compute how many rotations have elapsed since origin to restore window. + { + let mut s = state.lock().await; + let origin = humantime::parse_rfc3339(s.origin.as_ref().unwrap()) + .unwrap_or(SystemTime::now()); + let elapsed = origin.elapsed().unwrap_or_default(); + let rotations = if rotation_period.is_zero() { + 0 } else { - Duration::ZERO + (elapsed.as_nanos() / rotation_period.as_nanos()) as usize }; - - info!( - " {}: {} partitions, {} rows, interval {:?} → period {:?}, stagger {:?}", - topic.name, topic.partitions, topic.rows, topic.batch_interval, batch_period, stagger - ); - - for p in 0..topic.partitions { - let partition_name = format!("partition-{p}"); - let stagger_offset = stagger * p as u32; - - let worker = PartitionWorker { - client: client.clone(), - topic: topic.name.clone(), - partition: partition_name, - sample_path: topic.sample.clone(), - rows: topic.rows, - batch_period, - stagger_offset, - state: state.clone(), - state_path: state_path.to_path_buf(), - }; - - handles.push(tokio::spawn(worker.run())); + let computed_window = (rotations * config.topics.active) % config.topics.count; + if computed_window != s.window_start { + info!( + "restoring window to [{computed_window}) based on elapsed time (was {})", + s.window_start + ); + s.window_start = computed_window; + s.window_since = + Some(humantime::format_rfc3339(SystemTime::now()).to_string()); + s.save(state_path)?; } } - // Print a progress summary every 30 seconds. + let mut window_start = state.lock().await.window_start; + let mut handles = spawn_window( + &client, + &config.topics, + &schemas_dir, + &state, + state_path, + window_start, + batch_period, + config.speed, + ); + + // Progress reporter. let state_for_stats = state.clone(); tokio::spawn(async move { let mut last: HashMap = HashMap::new(); @@ -382,10 +526,34 @@ pub async fn run_batch(url: &str, config: BatchConfig, state_path: &Path) -> Res } }); - // Wait for all workers (they run indefinitely until the process is killed). - for h in handles { - let _ = h.await; - } + // Rotation loop. + loop { + tokio::time::sleep(rotation_period).await; - Ok(()) + // Stop current window. + for h in handles.drain(..) { + h.abort(); + } + + // Advance window. + window_start = (window_start + config.topics.active) % config.topics.count; + { + let mut s = state.lock().await; + s.window_start = window_start; + s.window_since = + Some(humantime::format_rfc3339(SystemTime::now()).to_string()); + s.save(state_path)?; + } + + handles = spawn_window( + &client, + &config.topics, + &schemas_dir, + &state, + state_path, + window_start, + batch_period, + config.speed, + ); + } } diff --git a/bench/src/bin/batch_load.rs b/bench/src/bin/batch_load.rs index 25dc43f..9a616db 100644 --- a/bench/src/bin/batch_load.rs +++ b/bench/src/bin/batch_load.rs @@ -5,36 +5,27 @@ use tracing_subscriber::{fmt, EnvFilter}; use bench::batch::{run_batch, BatchConfig}; -/// Batch-oriented load generator for an existing Plateau server. +/// Batch-job load generator for an existing Plateau server. /// -/// Simulates staggered batch jobs: each topic has its own schema, each -/// partition fires on a fixed schedule with an evenly-spread stagger offset. -/// A state file records the last completed batch per partition; on restart -/// the tool catches up any missed batches immediately before resuming the -/// normal schedule. -/// -/// Per-topic settings (partitions, rows, batch_interval) live on each topic. -/// A [defaults] table supplies fallbacks for any topic that omits them. +/// Generates a pool of topics with synthetic schemas (random column counts and +/// types), then simulates staggered batch jobs over a sliding active window. +/// A state file and a directory of Arrow schema files are written on first run +/// so the tool can be safely stopped and restarted. /// /// Example config (batch-config.toml): /// -/// speed = 60.0 # 1h batches fire every 1 minute +/// speed = 60.0 # compress time: 1h intervals fire every 1 minute +/// schemas_dir = "batch-schemas" /// -/// [defaults] +/// [topics] +/// count = 200 # total topic pool +/// active = 8 # topics writing at once +/// rotation_interval = "1h" # how often the active window advances +/// columns_min = 3 # min data columns per topic (excludes `time`) +/// columns_max = 35 # max data columns per topic (exclusive) +/// batch_interval = "1h" # real-world interval between batches per partition /// partitions = 4 /// rows = 10000 -/// batch_interval = "1h" -/// -/// [[topics]] -/// name = "transactions" -/// sample = "samples/list-ccfraud.arrow" -/// batch_interval = "15m" # this job runs more often -/// -/// [[topics]] -/// name = "images" -/// sample = "samples/image_224x224.arrow" -/// partitions = 2 -/// rows = 50 #[derive(Parser)] #[command(about, verbatim_doc_comment)] struct Args { @@ -42,7 +33,7 @@ struct Args { #[arg(long, default_value = "batch-config.toml")] config: PathBuf, - /// Plateau server URL (overrides nothing in config; config has no URL field). + /// Plateau server URL. #[arg(long, default_value = "http://localhost:3030")] url: String, From 940c79626def06f4ebb2c66e4b958cd21226e62a Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 22 Jun 2026 19:20:09 +0000 Subject: [PATCH 05/14] Make partitions and rows per-topic ranges with normal-distribution sampling - partitions_min/max: each topic draws its own partition count - columns_min/max: now sampled from a normal distribution (was uniform) - rows_min/max define an overall rows-per-insert distribution; each topic draws two values from it to form its own [min, max] sub-range, and every insert samples a row count from that per-topic range - All range draws use a clamped normal distribution (Box-Muller, mean at midpoint, sigma = range/4) via a shared normal_range helper - Per-topic parameters are derived deterministically from hash(schemas_seed, topic_idx), so they are stable and recomputable across restarts without a manifest; schema files are independently regenerable - Add unit tests: range bounds, param determinism, schema load-back + no spurious regeneration Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01GAMUPUAesj4bPh2y2cvaKc --- bench/src/batch.rs | 286 +++++++++++++++++++++++++++--------- bench/src/bin/batch_load.rs | 8 +- 2 files changed, 220 insertions(+), 74 deletions(-) diff --git a/bench/src/batch.rs b/bench/src/batch.rs index 7dddf67..d7d26ae 100644 --- a/bench/src/batch.rs +++ b/bench/src/batch.rs @@ -1,6 +1,5 @@ use std::collections::HashMap; use std::hash::{Hash, Hasher as _}; -use std::ops::Range; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::thread; @@ -39,20 +38,21 @@ pub struct TopicsConfig { pub columns_min: usize, /// Max number of data columns per topic (exclusive). pub columns_max: usize, + /// Min partitions per topic. + pub partitions_min: usize, + /// Max partitions per topic (exclusive). + pub partitions_max: usize, + /// Min of the overall rows-per-insert distribution. Each topic draws its + /// own [min, max] sub-range from this distribution; every insert then + /// samples a row count from that per-topic range. + pub rows_min: usize, + /// Max of the overall rows-per-insert distribution (exclusive). + pub rows_max: usize, /// Real-world batch interval per topic. #[serde(with = "humantime_serde")] pub batch_interval: Duration, - /// Partitions per topic. - #[serde(default = "default_partitions")] - pub partitions: usize, - /// Rows per batch. - #[serde(default = "default_rows")] - pub rows: usize, } -fn default_partitions() -> usize { 4 } -fn default_rows() -> usize { 10_000 } - #[derive(Debug, Deserialize)] pub struct BatchConfig { /// Speed multiplier: 60.0 means 1h intervals fire every 1 minute. @@ -74,6 +74,16 @@ impl BatchConfig { } } +/// Per-topic parameters derived deterministically from the schema seed. +#[derive(Debug, Clone)] +pub struct TopicParams { + pub columns: usize, + pub partitions: usize, + /// Inclusive-min / exclusive-max row count for each insert into this topic. + pub rows_min: usize, + pub rows_max: usize, +} + // ── State ───────────────────────────────────────────────────────────────────── #[derive(Debug, Clone, Serialize, Deserialize, Default)] @@ -114,6 +124,62 @@ impl BatchState { } } +// ── Range sampling (normal distribution) ────────────────────────────────────── + +/// Draw an integer in [min, max) from a normal distribution centered on the +/// midpoint, with σ = range/4 (so ~95% of mass lands in range), clamped to +/// [min, max). Uses the Box-Muller transform. +fn normal_range(rng: &mut Random, min: usize, max: usize) -> usize { + if max <= min + 1 { + return min; + } + let lo = min as f64; + let hi = (max - 1) as f64; + let mean = (lo + hi) / 2.0; + let std = (hi - lo) / 4.0; + + // u1 in (0, 1] to keep ln() finite. + let u1: f64 = 1.0 - rng.gen_range(0.0..1.0); + let u2: f64 = rng.gen_range(0.0..1.0); + let z = (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos(); + + (mean + z * std).round().clamp(lo, hi) as usize +} + +// ── Per-topic parameter derivation ──────────────────────────────────────────── + +/// Stable per-topic RNG seed derived from the schema seed and topic index. +fn topic_seed(seed: u64, idx: usize) -> u64 { + let mut h = std::collections::hash_map::DefaultHasher::new(); + seed.hash(&mut h); + idx.hash(&mut h); + h.finish() +} + +/// Deterministically derive a topic's parameters. The same `rng` is then used +/// to generate the schema, so params and schema stay consistent. +fn draw_params(rng: &mut Random, t: &TopicsConfig) -> TopicParams { + let columns = normal_range(rng, t.columns_min, t.columns_max); + let partitions = normal_range(rng, t.partitions_min, t.partitions_max).max(1); + // Each topic draws two values from the global rows distribution to form its + // own [min, max] sub-range; per-insert counts are sampled from that range. + let a = normal_range(rng, t.rows_min, t.rows_max); + let b = normal_range(rng, t.rows_min, t.rows_max); + let (rmin, rmax) = if a <= b { (a, b) } else { (b, a) }; + TopicParams { + columns, + partitions, + rows_min: rmin, + rows_max: rmax + 1, + } +} + +/// Recompute a topic's parameters without touching the filesystem. +fn topic_params(seed: u64, idx: usize, t: &TopicsConfig) -> TopicParams { + let mut rng = Random::from_seed(topic_seed(seed, idx)); + draw_params(&mut rng, t) +} + // ── Schema generation ───────────────────────────────────────────────────────── fn topic_name(idx: usize) -> String { @@ -124,14 +190,10 @@ fn schema_path(schemas_dir: &Path, idx: usize) -> PathBuf { schemas_dir.join(format!("{}.arrow", topic_name(idx))) } -/// Generate a schema with a random number of flat columns in [col_range) and -/// write a tiny seed batch to `path` so `build_sampler` can read it back. -fn generate_schema_file( - path: &Path, - col_range: Range, - rng: &mut Random, -) -> Result<()> { - let n_cols = rng.gen_range(col_range); +/// Generate a schema with `n_cols` flat columns and write a one-row seed batch +/// to `path` so `build_sampler` can read it back. `rng` must already have had +/// `draw_params` applied (so its state follows the params draw). +fn generate_schema_file(path: &Path, n_cols: usize, rng: &mut Random) -> Result<()> { let mut flat = sample_flat(); let mut fields: Vec> = vec![Arc::new(Field::new("time", DataType::Int64, false))]; @@ -148,7 +210,8 @@ fn generate_schema_file( branch: 1_i32..2_i32, }; - let mut time_sampler = primitive_len_sampler::<_, _, arrow_array::types::Int64Type>(Now, AlwaysValid); + let mut time_sampler = + primitive_len_sampler::<_, _, arrow_array::types::Int64Type>(Now, AlwaysValid); time_sampler.set_len(1); let time_col = time_sampler.generate(rng); @@ -169,49 +232,26 @@ fn generate_schema_file( Ok(()) } -/// Ensure all N topic schema files exist in `schemas_dir`, generating any -/// that are missing. Returns the seed used (stored in state for provenance). -pub fn ensure_schemas( - schemas_dir: &Path, - count: usize, - col_min: usize, - col_max: usize, - seed: u64, -) -> Result<()> { +/// Ensure every topic's schema file exists, generating any that are missing. +/// Each topic is generated from its own stable seed, so files are reproducible +/// and independent — regenerating one never disturbs the others. +pub fn ensure_schemas(schemas_dir: &Path, t: &TopicsConfig, seed: u64) -> Result<()> { std::fs::create_dir_all(schemas_dir)?; - let col_range = col_min..col_max; - let mut rng = Random::from_seed(seed); - - // Advance the RNG past any already-generated schemas so adding more topics - // later doesn't change existing schemas. - let mut missing = vec![]; - for idx in 0..count { + let mut generated = 0; + for idx in 0..t.count { let path = schema_path(schemas_dir, idx); if path.exists() { - // Burn the same number of RNG calls as generation would to keep - // future topics consistent. - let _ = rng.gen_range(col_range.clone()); - // Burn one call per possible column for datatypes — approximate. - } else { - missing.push(idx); + continue; } + let mut rng = Random::from_seed(topic_seed(seed, idx)); + let params = draw_params(&mut rng, t); + generate_schema_file(&path, params.columns, &mut rng)?; + generated += 1; } - if !missing.is_empty() { - info!("generating {} topic schema files in {:?}", missing.len(), schemas_dir); - // Re-seed cleanly and generate all from scratch for simplicity. - // (All files are written atomically, so existing ones are not touched.) - let mut rng = Random::from_seed(seed); - for idx in 0..count { - let path = schema_path(schemas_dir, idx); - if !path.exists() { - generate_schema_file(&path, col_range.clone(), &mut rng)?; - } else { - // Burn the RNG state as if we had generated this one. - let _ = rng.gen_range(col_range.clone()); - } - } + if generated > 0 { + info!("generated {generated} topic schema files in {schemas_dir:?}"); } Ok(()) @@ -234,7 +274,8 @@ fn batch_seed(topic: &str, partition: &str, batch_idx: u64) -> u64 { struct SamplerRequest { seed: u64, - rows: usize, + rows_min: usize, + rows_max: usize, } struct SamplerThread { @@ -249,7 +290,9 @@ impl SamplerThread { .expect("failed to build sampler"); for req in self.req_rx { let mut random = Random::from_seed(req.seed); - sampler.set_len(req.rows); + // Sample this insert's row count from the topic's range (normal). + let rows = normal_range(&mut random, req.rows_min, req.rows_max); + sampler.set_len(rows); let multi = sampler.generate(&mut random); if self.result_tx.blocking_send(multi).is_err() { break; @@ -265,7 +308,8 @@ struct PartitionWorker { topic: String, partition: String, sample_path: PathBuf, - rows: usize, + rows_min: usize, + rows_max: usize, batch_period: Duration, stagger_offset: Duration, state: Arc>, @@ -329,7 +373,12 @@ impl PartitionWorker { } let seed = batch_seed(&self.topic, &self.partition, batch_idx); - if req_tx.send(SamplerRequest { seed, rows: self.rows }).is_err() { + let req = SamplerRequest { + seed, + rows_min: self.rows_min, + rows_max: self.rows_max, + }; + if req_tx.send(req).is_err() { break; } let multi = match result_rx.recv().await { @@ -373,6 +422,7 @@ impl PartitionWorker { // ── Window management ───────────────────────────────────────────────────────── +#[allow(clippy::too_many_arguments)] fn spawn_window( client: &Client, topics: &TopicsConfig, @@ -382,9 +432,16 @@ fn spawn_window( window_start: usize, batch_period: Duration, speed: f64, + schemas_seed: u64, ) -> Vec> { let window_end = (window_start + topics.active).min(topics.count); - let total_partitions = (window_end - window_start) * topics.partitions; + + // Resolve each topic's parameters; partitions vary per topic. + let params: Vec<(usize, TopicParams)> = (window_start..window_end) + .map(|idx| (idx, topic_params(schemas_seed, idx, topics))) + .collect(); + + let total_partitions: usize = params.iter().map(|(_, p)| p.partitions).sum(); let stagger = if total_partitions > 1 { batch_period / total_partitions as u32 } else { @@ -394,17 +451,18 @@ fn spawn_window( let mut handles = vec![]; let mut slot = 0usize; - for topic_idx in window_start..window_end { - let name = topic_name(topic_idx); - let sample_path = schema_path(schemas_dir, topic_idx); + for (topic_idx, tp) in ¶ms { + let name = topic_name(*topic_idx); + let sample_path = schema_path(schemas_dir, *topic_idx); - for p in 0..topics.partitions { + for p in 0..tp.partitions { let worker = PartitionWorker { client: client.clone(), topic: name.clone(), partition: format!("partition-{p}"), sample_path: sample_path.clone(), - rows: topics.rows, + rows_min: tp.rows_min, + rows_max: tp.rows_max, batch_period, stagger_offset: stagger * slot as u32, state: state.clone(), @@ -454,14 +512,10 @@ pub async fn run_batch(url: &str, config: BatchConfig, state_path: &Path) -> Res } state.save(state_path)?; + let schemas_seed = state.schemas_seed.unwrap(); + // Generate any missing schema files. - ensure_schemas( - &schemas_dir, - config.topics.count, - config.topics.columns_min, - config.topics.columns_max, - state.schemas_seed.unwrap(), - )?; + ensure_schemas(&schemas_dir, &config.topics, schemas_seed)?; let state = Arc::new(Mutex::new(state)); @@ -502,6 +556,7 @@ pub async fn run_batch(url: &str, config: BatchConfig, state_path: &Path) -> Res window_start, batch_period, config.speed, + schemas_seed, ); // Progress reporter. @@ -554,6 +609,93 @@ pub async fn run_batch(url: &str, config: BatchConfig, state_path: &Path) -> Res window_start, batch_period, config.speed, + schemas_seed, ); } } + +#[cfg(test)] +mod test { + use super::*; + + fn cfg() -> TopicsConfig { + TopicsConfig { + count: 16, + active: 4, + rotation_interval: Duration::from_secs(60), + columns_min: 3, + columns_max: 35, + partitions_min: 1, + partitions_max: 8, + rows_min: 1000, + rows_max: 50000, + batch_interval: Duration::from_secs(60), + } + } + + #[test] + fn normal_range_stays_in_bounds() { + let mut rng = Random::from_seed(7); + for _ in 0..10_000 { + let v = normal_range(&mut rng, 3, 35); + assert!((3..35).contains(&v), "out of range: {v}"); + } + // Degenerate ranges. + assert_eq!(normal_range(&mut rng, 5, 5), 5); + assert_eq!(normal_range(&mut rng, 5, 6), 5); + } + + #[test] + fn topic_params_are_deterministic() { + let t = cfg(); + for idx in 0..t.count { + let a = topic_params(42, idx, &t); + let b = topic_params(42, idx, &t); + assert_eq!(a.columns, b.columns); + assert_eq!(a.partitions, b.partitions); + assert_eq!(a.rows_min, b.rows_min); + assert_eq!(a.rows_max, b.rows_max); + assert!(a.partitions >= 1); + assert!((t.columns_min..t.columns_max).contains(&a.columns)); + assert!(a.rows_min < a.rows_max); + } + // Different seeds should generally differ. + let p1 = topic_params(1, 0, &t); + let p2 = topic_params(2, 0, &t); + assert!(p1.columns != p2.columns || p1.partitions != p2.partitions || p1.rows_min != p2.rows_min); + } + + #[test] + fn generated_schemas_are_loadable_and_stable() { + let t = cfg(); + let dir = std::env::temp_dir().join(format!("batch-test-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + + ensure_schemas(&dir, &t, 99).unwrap(); + + for idx in 0..t.count { + let path = schema_path(&dir, idx); + assert!(path.exists(), "missing schema {idx}"); + // build_sampler must read it back and produce rows. + let mut sampler = crate::load::build_sampler(&path).unwrap(); + sampler.set_len(10); + let mut rng = Random::from_seed(idx as u64); + let multi = sampler.generate(&mut rng); + let cols = multi.schema.fields().len(); + let expected = topic_params(99, idx, &t).columns + 1; // + time column + assert_eq!(cols, expected, "topic {idx} column count mismatch"); + } + + // Re-running must not regenerate (files already present). + let before: Vec<_> = (0..t.count) + .map(|i| std::fs::metadata(schema_path(&dir, i)).unwrap().modified().unwrap()) + .collect(); + ensure_schemas(&dir, &t, 99).unwrap(); + let after: Vec<_> = (0..t.count) + .map(|i| std::fs::metadata(schema_path(&dir, i)).unwrap().modified().unwrap()) + .collect(); + assert_eq!(before, after, "schemas were unexpectedly regenerated"); + + let _ = std::fs::remove_dir_all(&dir); + } +} diff --git a/bench/src/bin/batch_load.rs b/bench/src/bin/batch_load.rs index 9a616db..d8071b9 100644 --- a/bench/src/bin/batch_load.rs +++ b/bench/src/bin/batch_load.rs @@ -23,9 +23,13 @@ use bench::batch::{run_batch, BatchConfig}; /// rotation_interval = "1h" # how often the active window advances /// columns_min = 3 # min data columns per topic (excludes `time`) /// columns_max = 35 # max data columns per topic (exclusive) +/// partitions_min = 1 # partitions drawn per topic from [min, max) +/// partitions_max = 8 +/// rows_min = 1000 # overall rows-per-insert distribution; each topic +/// rows_max = 50000 # draws its own sub-range, sampled per insert /// batch_interval = "1h" # real-world interval between batches per partition -/// partitions = 4 -/// rows = 10000 +/// +/// All values drawn from a range use a clamped normal distribution. #[derive(Parser)] #[command(about, verbatim_doc_comment)] struct Args { From 4ab8b92aa72080539faf98e8a09b33162b1a70d4 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 18:55:14 +0000 Subject: [PATCH 06/14] Add batch-load Docker image built and pushed alongside plateau Dockerfile: - Combined cargo build now builds both plateau and batch-load binaries - Named the existing final stage `plateau` - Added `batch-load` stage copying the bench binary CI (rust.yml): - Added BENCH_REGISTRY_IMAGE env var (ghcr.io/wallaroolabs/plateau-bench) - Added --target plateau to existing build jobs - Added build-bench-amd64, build-bench-arm64, merge-bench jobs mirroring the plateau pattern - Bench jobs run on push to main, version tags, workflow_dispatch, and PRs (guarded by same-repo check so fork PRs don't fail on missing secrets) - merge-bench tags with sha, branch, pr number, and semver Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01GAMUPUAesj4bPh2y2cvaKc --- .github/workflows/rust.yml | 172 +++++++++++++++++++++++++++++++++++++ Dockerfile | 16 +++- 2 files changed, 185 insertions(+), 3 deletions(-) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 54d23ff..71d2b7b 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -14,6 +14,7 @@ on: env: CARGO_TERM_COLOR: always REGISTRY_IMAGE: ghcr.io/wallaroolabs/plateau + BENCH_REGISTRY_IMAGE: ghcr.io/wallaroolabs/plateau-bench RUST_VERSION: 1.84.1 jobs: @@ -76,6 +77,7 @@ jobs: id: build uses: docker/build-push-action@v6 with: + target: plateau labels: ${{ steps.meta.outputs.labels }} outputs: type=image,name=${{ env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true build-args: | @@ -134,6 +136,7 @@ jobs: id: build uses: docker/build-push-action@v6 with: + target: plateau labels: ${{ steps.meta.outputs.labels }} outputs: type=image,name=${{ env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true build-args: | @@ -153,6 +156,126 @@ jobs: if-no-files-found: error retention-days: 1 + build-bench-amd64: + if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository + runs-on: + - ubuntu-latest + steps: + - name: Prepare + run: | + echo "PLATFORM_ARCH=$(arch)" >> $GITHUB_ENV + + - name: Docker meta + id: meta + uses: docker/metadata-action@v5 + with: + images: ${{ env.BENCH_REGISTRY_IMAGE }} + labels: | + org.opencontainers.image.vendor="Wallaroo Labs" + org.opencontainers.image.source="https://github.com/WallarooLabs/plateau/Dockerfile" + org.opencontainers.image.title="plateau-bench" + + - name: Setup Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to ghcr.io + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.PUSH_CONTAINER_TOKEN }} + + - name: Login to us-docker.pkg.dev/wallaroo-dev-253816/docker-hub-us + uses: docker/login-action@v3 + with: + registry: us-docker.pkg.dev + username: _json_key + password: ${{ secrets.US_PKG_DEV_CACHE_JSON_KEY }} + + - name: Build and push by digest + id: build + uses: docker/build-push-action@v6 + with: + target: batch-load + labels: ${{ steps.meta.outputs.labels }} + outputs: type=image,name=${{ env.BENCH_REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true + build-args: | + RUST_VERSION=${{ env.RUST_VERSION }} + + - name: Export digest + run: | + mkdir -p /tmp/digests + digest="${{ steps.build.outputs.digest }}" + touch "/tmp/digests/${digest#sha256:}" + + - name: Upload digest + uses: actions/upload-artifact@v4 + with: + name: bench-digests-${{ env.PLATFORM_ARCH }} + path: /tmp/digests/* + if-no-files-found: error + retention-days: 1 + + build-bench-arm64: + if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository + runs-on: + - buildjet-4vcpu-ubuntu-2204-arm + steps: + - name: Prepare + run: | + echo "PLATFORM_ARCH=$(arch)" >> $GITHUB_ENV + + - name: Docker meta + id: meta + uses: docker/metadata-action@v5 + with: + images: ${{ env.BENCH_REGISTRY_IMAGE }} + labels: | + org.opencontainers.image.vendor="Wallaroo Labs" + org.opencontainers.image.source="https://github.com/WallarooLabs/plateau/Dockerfile" + org.opencontainers.image.title="plateau-bench" + + - name: Setup Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to ghcr.io + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.PUSH_CONTAINER_TOKEN }} + + - name: Login to us-docker.pkg.dev/wallaroo-dev-253816/docker-hub-us + uses: docker/login-action@v3 + with: + registry: us-docker.pkg.dev + username: _json_key + password: ${{ secrets.US_PKG_DEV_CACHE_JSON_KEY }} + + - name: Build and push by digest + id: build + uses: docker/build-push-action@v6 + with: + target: batch-load + labels: ${{ steps.meta.outputs.labels }} + outputs: type=image,name=${{ env.BENCH_REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true + build-args: | + RUST_VERSION=${{ env.RUST_VERSION }} + + - name: Export digest + run: | + mkdir -p /tmp/digests + digest="${{ steps.build.outputs.digest }}" + touch "/tmp/digests/${digest#sha256:}" + + - name: Upload digest + uses: actions/upload-artifact@v4 + with: + name: bench-digests-${{ env.PLATFORM_ARCH }} + path: /tmp/digests/* + if-no-files-found: error + retention-days: 1 + merge: runs-on: ubuntu-latest needs: @@ -199,3 +322,52 @@ jobs: - name: Inspect image run: | docker buildx imagetools inspect ${{ env.REGISTRY_IMAGE }}:${{ steps.meta.outputs.version }} + + merge-bench: + if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository + runs-on: ubuntu-latest + needs: + - build-bench-amd64 + - build-bench-arm64 + steps: + - name: Download digests + uses: actions/download-artifact@v4 + with: + path: /tmp/digests + pattern: bench-digests-* + merge-multiple: true + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Docker meta + id: meta + uses: docker/metadata-action@v5 + with: + images: ${{ env.BENCH_REGISTRY_IMAGE }} + tags: | + type=sha + type=ref,event=branch + type=ref,event=pr + type=semver,pattern={{version}} + labels: | + org.opencontainers.image.vendor="Wallaroo Labs" + org.opencontainers.image.source="https://github.com/WallarooLabs/plateau/Dockerfile" + org.opencontainers.image.title="plateau-bench" + + - name: Login to ghcr.io + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.PUSH_CONTAINER_TOKEN }} + + - name: Create manifest list and push + working-directory: /tmp/digests + run: | + docker buildx imagetools create $(jq -cr '.tags | map("-t " + .) | join(" ")' <<< "$DOCKER_METADATA_OUTPUT_JSON") \ + $(printf '${{ env.BENCH_REGISTRY_IMAGE }}@sha256:%s ' *) + + - name: Inspect image + run: | + docker buildx imagetools inspect ${{ env.BENCH_REGISTRY_IMAGE }}:${{ steps.meta.outputs.version }} diff --git a/Dockerfile b/Dockerfile index a1aa43c..74285b9 100644 --- a/Dockerfile +++ b/Dockerfile @@ -11,10 +11,11 @@ COPY . . RUN \ if [ "${TARGETARCH}" = "amd64" ]; then ARCH=x86_64; elif [ "${TARGETARCH}" = "arm64" ]; then ARCH=aarch64; else exit 1; fi && \ rustup target add ${ARCH}-unknown-linux-musl && \ - cargo build --release --target ${ARCH}-unknown-linux-musl -p plateau && \ - cp target/${ARCH}-unknown-linux-musl/release/plateau target/release/plateau + cargo build --release --target ${ARCH}-unknown-linux-musl -p plateau --bin batch-load && \ + cp target/${ARCH}-unknown-linux-musl/release/plateau target/release/plateau && \ + cp target/${ARCH}-unknown-linux-musl/release/batch-load target/release/batch-load -FROM scratch +FROM scratch AS plateau LABEL org.opencontainers.image.vendor="Wallaroo Labs" LABEL org.opencontainers.image.source="https://github.com/WallarooLabs/plateau/Dockerfile" @@ -22,3 +23,12 @@ LABEL org.opencontainers.image.title="plateau" COPY --from=build /usr/src/plateau/target/release/plateau . CMD ["./plateau"] + +FROM scratch AS batch-load + +LABEL org.opencontainers.image.vendor="Wallaroo Labs" +LABEL org.opencontainers.image.source="https://github.com/WallarooLabs/plateau/Dockerfile" +LABEL org.opencontainers.image.title="plateau-bench" + +COPY --from=build /usr/src/plateau/target/release/batch-load . +CMD ["./batch-load"] From 2438742a7c65654b9757862b11548013da5daca1 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 18:56:09 +0000 Subject: [PATCH 07/14] Move bench image build to a separate manually-triggered workflow bench-image.yml is workflow_dispatch-only so it never runs on PRs or routine pushes. Trigger it from the Actions tab (or gh workflow run bench-image.yml --ref ) whenever a new bench image is needed. rust.yml is cleaned up: BENCH_REGISTRY_IMAGE env var and the bench build jobs are removed; the plateau build jobs retain --target plateau. Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01GAMUPUAesj4bPh2y2cvaKc --- .github/workflows/bench-image.yml | 174 ++++++++++++++++++++++++++++++ .github/workflows/rust.yml | 169 ----------------------------- 2 files changed, 174 insertions(+), 169 deletions(-) create mode 100644 .github/workflows/bench-image.yml diff --git a/.github/workflows/bench-image.yml b/.github/workflows/bench-image.yml new file mode 100644 index 0000000..efc1c11 --- /dev/null +++ b/.github/workflows/bench-image.yml @@ -0,0 +1,174 @@ +name: Bench Image + +on: + workflow_dispatch: {} + +env: + BENCH_REGISTRY_IMAGE: ghcr.io/wallaroolabs/plateau-bench + RUST_VERSION: 1.84.1 + +jobs: + build-bench-amd64: + runs-on: + - ubuntu-latest + steps: + - name: Prepare + run: | + echo "PLATFORM_ARCH=$(arch)" >> $GITHUB_ENV + + - name: Docker meta + id: meta + uses: docker/metadata-action@v5 + with: + images: ${{ env.BENCH_REGISTRY_IMAGE }} + labels: | + org.opencontainers.image.vendor="Wallaroo Labs" + org.opencontainers.image.source="https://github.com/WallarooLabs/plateau/Dockerfile" + org.opencontainers.image.title="plateau-bench" + + - name: Setup Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to ghcr.io + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.PUSH_CONTAINER_TOKEN }} + + - name: Login to us-docker.pkg.dev/wallaroo-dev-253816/docker-hub-us + uses: docker/login-action@v3 + with: + registry: us-docker.pkg.dev + username: _json_key + password: ${{ secrets.US_PKG_DEV_CACHE_JSON_KEY }} + + - name: Build and push by digest + id: build + uses: docker/build-push-action@v6 + with: + target: batch-load + labels: ${{ steps.meta.outputs.labels }} + outputs: type=image,name=${{ env.BENCH_REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true + build-args: | + RUST_VERSION=${{ env.RUST_VERSION }} + + - name: Export digest + run: | + mkdir -p /tmp/digests + digest="${{ steps.build.outputs.digest }}" + touch "/tmp/digests/${digest#sha256:}" + + - name: Upload digest + uses: actions/upload-artifact@v4 + with: + name: bench-digests-${{ env.PLATFORM_ARCH }} + path: /tmp/digests/* + if-no-files-found: error + retention-days: 1 + + build-bench-arm64: + runs-on: + - buildjet-4vcpu-ubuntu-2204-arm + steps: + - name: Prepare + run: | + echo "PLATFORM_ARCH=$(arch)" >> $GITHUB_ENV + + - name: Docker meta + id: meta + uses: docker/metadata-action@v5 + with: + images: ${{ env.BENCH_REGISTRY_IMAGE }} + labels: | + org.opencontainers.image.vendor="Wallaroo Labs" + org.opencontainers.image.source="https://github.com/WallarooLabs/plateau/Dockerfile" + org.opencontainers.image.title="plateau-bench" + + - name: Setup Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to ghcr.io + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.PUSH_CONTAINER_TOKEN }} + + - name: Login to us-docker.pkg.dev/wallaroo-dev-253816/docker-hub-us + uses: docker/login-action@v3 + with: + registry: us-docker.pkg.dev + username: _json_key + password: ${{ secrets.US_PKG_DEV_CACHE_JSON_KEY }} + + - name: Build and push by digest + id: build + uses: docker/build-push-action@v6 + with: + target: batch-load + labels: ${{ steps.meta.outputs.labels }} + outputs: type=image,name=${{ env.BENCH_REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true + build-args: | + RUST_VERSION=${{ env.RUST_VERSION }} + + - name: Export digest + run: | + mkdir -p /tmp/digests + digest="${{ steps.build.outputs.digest }}" + touch "/tmp/digests/${digest#sha256:}" + + - name: Upload digest + uses: actions/upload-artifact@v4 + with: + name: bench-digests-${{ env.PLATFORM_ARCH }} + path: /tmp/digests/* + if-no-files-found: error + retention-days: 1 + + merge-bench: + runs-on: ubuntu-latest + needs: + - build-bench-amd64 + - build-bench-arm64 + steps: + - name: Download digests + uses: actions/download-artifact@v4 + with: + path: /tmp/digests + pattern: bench-digests-* + merge-multiple: true + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Docker meta + id: meta + uses: docker/metadata-action@v5 + with: + images: ${{ env.BENCH_REGISTRY_IMAGE }} + tags: | + type=sha + type=ref,event=branch + type=semver,pattern={{version}} + labels: | + org.opencontainers.image.vendor="Wallaroo Labs" + org.opencontainers.image.source="https://github.com/WallarooLabs/plateau/Dockerfile" + org.opencontainers.image.title="plateau-bench" + + - name: Login to ghcr.io + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.PUSH_CONTAINER_TOKEN }} + + - name: Create manifest list and push + working-directory: /tmp/digests + run: | + docker buildx imagetools create $(jq -cr '.tags | map("-t " + .) | join(" ")' <<< "$DOCKER_METADATA_OUTPUT_JSON") \ + $(printf '${{ env.BENCH_REGISTRY_IMAGE }}@sha256:%s ' *) + + - name: Inspect image + run: | + docker buildx imagetools inspect ${{ env.BENCH_REGISTRY_IMAGE }}:${{ steps.meta.outputs.version }} diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 71d2b7b..0acdbde 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -14,7 +14,6 @@ on: env: CARGO_TERM_COLOR: always REGISTRY_IMAGE: ghcr.io/wallaroolabs/plateau - BENCH_REGISTRY_IMAGE: ghcr.io/wallaroolabs/plateau-bench RUST_VERSION: 1.84.1 jobs: @@ -156,126 +155,6 @@ jobs: if-no-files-found: error retention-days: 1 - build-bench-amd64: - if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository - runs-on: - - ubuntu-latest - steps: - - name: Prepare - run: | - echo "PLATFORM_ARCH=$(arch)" >> $GITHUB_ENV - - - name: Docker meta - id: meta - uses: docker/metadata-action@v5 - with: - images: ${{ env.BENCH_REGISTRY_IMAGE }} - labels: | - org.opencontainers.image.vendor="Wallaroo Labs" - org.opencontainers.image.source="https://github.com/WallarooLabs/plateau/Dockerfile" - org.opencontainers.image.title="plateau-bench" - - - name: Setup Docker Buildx - uses: docker/setup-buildx-action@v3 - - - name: Login to ghcr.io - uses: docker/login-action@v3 - with: - registry: ghcr.io - username: ${{ github.actor }} - password: ${{ secrets.PUSH_CONTAINER_TOKEN }} - - - name: Login to us-docker.pkg.dev/wallaroo-dev-253816/docker-hub-us - uses: docker/login-action@v3 - with: - registry: us-docker.pkg.dev - username: _json_key - password: ${{ secrets.US_PKG_DEV_CACHE_JSON_KEY }} - - - name: Build and push by digest - id: build - uses: docker/build-push-action@v6 - with: - target: batch-load - labels: ${{ steps.meta.outputs.labels }} - outputs: type=image,name=${{ env.BENCH_REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true - build-args: | - RUST_VERSION=${{ env.RUST_VERSION }} - - - name: Export digest - run: | - mkdir -p /tmp/digests - digest="${{ steps.build.outputs.digest }}" - touch "/tmp/digests/${digest#sha256:}" - - - name: Upload digest - uses: actions/upload-artifact@v4 - with: - name: bench-digests-${{ env.PLATFORM_ARCH }} - path: /tmp/digests/* - if-no-files-found: error - retention-days: 1 - - build-bench-arm64: - if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository - runs-on: - - buildjet-4vcpu-ubuntu-2204-arm - steps: - - name: Prepare - run: | - echo "PLATFORM_ARCH=$(arch)" >> $GITHUB_ENV - - - name: Docker meta - id: meta - uses: docker/metadata-action@v5 - with: - images: ${{ env.BENCH_REGISTRY_IMAGE }} - labels: | - org.opencontainers.image.vendor="Wallaroo Labs" - org.opencontainers.image.source="https://github.com/WallarooLabs/plateau/Dockerfile" - org.opencontainers.image.title="plateau-bench" - - - name: Setup Docker Buildx - uses: docker/setup-buildx-action@v3 - - - name: Login to ghcr.io - uses: docker/login-action@v3 - with: - registry: ghcr.io - username: ${{ github.actor }} - password: ${{ secrets.PUSH_CONTAINER_TOKEN }} - - - name: Login to us-docker.pkg.dev/wallaroo-dev-253816/docker-hub-us - uses: docker/login-action@v3 - with: - registry: us-docker.pkg.dev - username: _json_key - password: ${{ secrets.US_PKG_DEV_CACHE_JSON_KEY }} - - - name: Build and push by digest - id: build - uses: docker/build-push-action@v6 - with: - target: batch-load - labels: ${{ steps.meta.outputs.labels }} - outputs: type=image,name=${{ env.BENCH_REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true - build-args: | - RUST_VERSION=${{ env.RUST_VERSION }} - - - name: Export digest - run: | - mkdir -p /tmp/digests - digest="${{ steps.build.outputs.digest }}" - touch "/tmp/digests/${digest#sha256:}" - - - name: Upload digest - uses: actions/upload-artifact@v4 - with: - name: bench-digests-${{ env.PLATFORM_ARCH }} - path: /tmp/digests/* - if-no-files-found: error - retention-days: 1 - merge: runs-on: ubuntu-latest needs: @@ -323,51 +202,3 @@ jobs: run: | docker buildx imagetools inspect ${{ env.REGISTRY_IMAGE }}:${{ steps.meta.outputs.version }} - merge-bench: - if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository - runs-on: ubuntu-latest - needs: - - build-bench-amd64 - - build-bench-arm64 - steps: - - name: Download digests - uses: actions/download-artifact@v4 - with: - path: /tmp/digests - pattern: bench-digests-* - merge-multiple: true - - - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v3 - - - name: Docker meta - id: meta - uses: docker/metadata-action@v5 - with: - images: ${{ env.BENCH_REGISTRY_IMAGE }} - tags: | - type=sha - type=ref,event=branch - type=ref,event=pr - type=semver,pattern={{version}} - labels: | - org.opencontainers.image.vendor="Wallaroo Labs" - org.opencontainers.image.source="https://github.com/WallarooLabs/plateau/Dockerfile" - org.opencontainers.image.title="plateau-bench" - - - name: Login to ghcr.io - uses: docker/login-action@v3 - with: - registry: ghcr.io - username: ${{ github.actor }} - password: ${{ secrets.PUSH_CONTAINER_TOKEN }} - - - name: Create manifest list and push - working-directory: /tmp/digests - run: | - docker buildx imagetools create $(jq -cr '.tags | map("-t " + .) | join(" ")' <<< "$DOCKER_METADATA_OUTPUT_JSON") \ - $(printf '${{ env.BENCH_REGISTRY_IMAGE }}@sha256:%s ' *) - - - name: Inspect image - run: | - docker buildx imagetools inspect ${{ env.BENCH_REGISTRY_IMAGE }}:${{ steps.meta.outputs.version }} From be01fca19e6b1a109ad5f1f62bb491b22665802f Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 19:14:17 +0000 Subject: [PATCH 08/14] Simplify bench-image workflow to single amd64 job No arm build needed. Drops the digest/artifact/merge pattern in favour of a single build-and-push job that pushes tags directly. Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01GAMUPUAesj4bPh2y2cvaKc --- .github/workflows/bench-image.yml | 136 ++---------------------------- 1 file changed, 8 insertions(+), 128 deletions(-) diff --git a/.github/workflows/bench-image.yml b/.github/workflows/bench-image.yml index efc1c11..98e164d 100644 --- a/.github/workflows/bench-image.yml +++ b/.github/workflows/bench-image.yml @@ -8,78 +8,19 @@ env: RUST_VERSION: 1.84.1 jobs: - build-bench-amd64: + build-and-push: runs-on: - ubuntu-latest steps: - - name: Prepare - run: | - echo "PLATFORM_ARCH=$(arch)" >> $GITHUB_ENV - - - name: Docker meta - id: meta - uses: docker/metadata-action@v5 - with: - images: ${{ env.BENCH_REGISTRY_IMAGE }} - labels: | - org.opencontainers.image.vendor="Wallaroo Labs" - org.opencontainers.image.source="https://github.com/WallarooLabs/plateau/Dockerfile" - org.opencontainers.image.title="plateau-bench" - - - name: Setup Docker Buildx - uses: docker/setup-buildx-action@v3 - - - name: Login to ghcr.io - uses: docker/login-action@v3 - with: - registry: ghcr.io - username: ${{ github.actor }} - password: ${{ secrets.PUSH_CONTAINER_TOKEN }} - - - name: Login to us-docker.pkg.dev/wallaroo-dev-253816/docker-hub-us - uses: docker/login-action@v3 - with: - registry: us-docker.pkg.dev - username: _json_key - password: ${{ secrets.US_PKG_DEV_CACHE_JSON_KEY }} - - - name: Build and push by digest - id: build - uses: docker/build-push-action@v6 - with: - target: batch-load - labels: ${{ steps.meta.outputs.labels }} - outputs: type=image,name=${{ env.BENCH_REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true - build-args: | - RUST_VERSION=${{ env.RUST_VERSION }} - - - name: Export digest - run: | - mkdir -p /tmp/digests - digest="${{ steps.build.outputs.digest }}" - touch "/tmp/digests/${digest#sha256:}" - - - name: Upload digest - uses: actions/upload-artifact@v4 - with: - name: bench-digests-${{ env.PLATFORM_ARCH }} - path: /tmp/digests/* - if-no-files-found: error - retention-days: 1 - - build-bench-arm64: - runs-on: - - buildjet-4vcpu-ubuntu-2204-arm - steps: - - name: Prepare - run: | - echo "PLATFORM_ARCH=$(arch)" >> $GITHUB_ENV - - name: Docker meta id: meta uses: docker/metadata-action@v5 with: images: ${{ env.BENCH_REGISTRY_IMAGE }} + tags: | + type=sha + type=ref,event=branch + type=semver,pattern={{version}} labels: | org.opencontainers.image.vendor="Wallaroo Labs" org.opencontainers.image.source="https://github.com/WallarooLabs/plateau/Dockerfile" @@ -102,73 +43,12 @@ jobs: username: _json_key password: ${{ secrets.US_PKG_DEV_CACHE_JSON_KEY }} - - name: Build and push by digest - id: build + - name: Build and push uses: docker/build-push-action@v6 with: target: batch-load + tags: ${{ steps.meta.outputs.tags }} labels: ${{ steps.meta.outputs.labels }} - outputs: type=image,name=${{ env.BENCH_REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true + push: true build-args: | RUST_VERSION=${{ env.RUST_VERSION }} - - - name: Export digest - run: | - mkdir -p /tmp/digests - digest="${{ steps.build.outputs.digest }}" - touch "/tmp/digests/${digest#sha256:}" - - - name: Upload digest - uses: actions/upload-artifact@v4 - with: - name: bench-digests-${{ env.PLATFORM_ARCH }} - path: /tmp/digests/* - if-no-files-found: error - retention-days: 1 - - merge-bench: - runs-on: ubuntu-latest - needs: - - build-bench-amd64 - - build-bench-arm64 - steps: - - name: Download digests - uses: actions/download-artifact@v4 - with: - path: /tmp/digests - pattern: bench-digests-* - merge-multiple: true - - - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v3 - - - name: Docker meta - id: meta - uses: docker/metadata-action@v5 - with: - images: ${{ env.BENCH_REGISTRY_IMAGE }} - tags: | - type=sha - type=ref,event=branch - type=semver,pattern={{version}} - labels: | - org.opencontainers.image.vendor="Wallaroo Labs" - org.opencontainers.image.source="https://github.com/WallarooLabs/plateau/Dockerfile" - org.opencontainers.image.title="plateau-bench" - - - name: Login to ghcr.io - uses: docker/login-action@v3 - with: - registry: ghcr.io - username: ${{ github.actor }} - password: ${{ secrets.PUSH_CONTAINER_TOKEN }} - - - name: Create manifest list and push - working-directory: /tmp/digests - run: | - docker buildx imagetools create $(jq -cr '.tags | map("-t " + .) | join(" ")' <<< "$DOCKER_METADATA_OUTPUT_JSON") \ - $(printf '${{ env.BENCH_REGISTRY_IMAGE }}@sha256:%s ' *) - - - name: Inspect image - run: | - docker buildx imagetools inspect ${{ env.BENCH_REGISTRY_IMAGE }}:${{ steps.meta.outputs.version }} From fb7363668ae644a8e8d8354fe1552417bb847699 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 25 Jun 2026 02:19:15 +0000 Subject: [PATCH 09/14] Switch batch-load to append_queue to avoid 10MB body limit errors append_records sends the full batch in one request, which can exceed the server's DefaultBodyLimit when rows_max is large. append_queue auto-splits at the client's DEFAULT_MAX_BATCH_BYTES (100KB), so large batches are chunked before they hit Axum's body limit. Co-Authored-By: Claude --- bench/src/batch.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/bench/src/batch.rs b/bench/src/batch.rs index d7d26ae..20de600 100644 --- a/bench/src/batch.rs +++ b/bench/src/batch.rs @@ -9,7 +9,7 @@ use anyhow::Result; use arrow_array::RecordBatch; use arrow_ipc::writer::FileWriter; use arrow_schema::{DataType, Field, Schema, SchemaRef}; -use plateau_client::{Client, Error, InsertQuery, MultiChunk}; +use plateau_client::{Client, Error, MultiChunk}; use reqwest::StatusCode; use sample_arrow_rs::array::FromDataType; use sample_arrow_rs::datatypes::sample_flat; @@ -317,7 +317,7 @@ struct PartitionWorker { } impl PartitionWorker { - async fn run(self) { + async fn run(mut self) { let (req_tx, req_rx) = std::sync::mpsc::channel::(); let (result_tx, mut result_rx) = mpsc::channel::(2); let st = SamplerThread { @@ -388,17 +388,18 @@ impl PartitionWorker { let start = Instant::now(); let r = self.client - .append_records(&self.topic, &self.partition, &InsertQuery::default(), multi) + .append_queue(&self.topic, &self.partition, multi) .await; match r { - Ok(ok) => { + Ok(Some(ok)) => { let rows = ok.span.end - ok.span.start; tracing::debug!( "{}/{} batch {} → {} rows in {:?}", self.topic, self.partition, batch_idx, rows, start.elapsed() ); } + Ok(None) => {} Err(Error::Server(ref e)) if e.status() == Some(StatusCode::TOO_MANY_REQUESTS) => { warn!("{}/{} rate limited on batch {}", self.topic, self.partition, batch_idx); tokio::time::sleep(Duration::from_secs(1)).await; From 29a1f055b4f3ef1bcad40278b3779587c184eb48 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 25 Jun 2026 19:46:41 +0000 Subject: [PATCH 10/14] Fix 400 body limit errors: pre-size batches to 8MB and log row counts on failure Axum 0.6's DefaultBodyLimit returns 400 (not 413) when the limit is exceeded, so append_queue's auto-shrink never fires. Fix by: - Adding Client::with_max_batch_bytes() to allow callers to set a conservative limit below the server's 10MB wall - Setting 8MB in run_batch so append_queue splits proactively before sending - Logging the row count on batch failure for easier diagnosis Co-Authored-By: Claude --- bench/src/batch.rs | 10 ++++++++-- client/src/lib.rs | 7 +++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/bench/src/batch.rs b/bench/src/batch.rs index 20de600..687256c 100644 --- a/bench/src/batch.rs +++ b/bench/src/batch.rs @@ -386,6 +386,7 @@ impl PartitionWorker { None => break, }; + let batch_rows: usize = multi.chunks.iter().map(|c| c.num_rows()).sum(); let start = Instant::now(); let r = self.client .append_queue(&self.topic, &self.partition, multi) @@ -406,7 +407,10 @@ impl PartitionWorker { continue; } Err(e) => { - warn!("{}/{} batch {} failed: {}", self.topic, self.partition, batch_idx, e); + warn!( + "{}/{} batch {} failed ({} rows): {}", + self.topic, self.partition, batch_idx, batch_rows, e + ); } } @@ -489,7 +493,9 @@ fn spawn_window( // ── Public entry point ──────────────────────────────────────────────────────── pub async fn run_batch(url: &str, config: BatchConfig, state_path: &Path) -> Result<()> { - let client = Client::new(url)?; + // Stay well under the server's 10MB DefaultBodyLimit; append_queue will + // auto-split any batch that exceeds this threshold. + let client = Client::new(url)?.with_max_batch_bytes(8 * 1024 * 1024); client .healthy(Duration::from_secs(10), Duration::from_millis(100)) .await?; diff --git a/client/src/lib.rs b/client/src/lib.rs index 6b19fc3..25cc9a8 100644 --- a/client/src/lib.rs +++ b/client/src/lib.rs @@ -222,6 +222,13 @@ impl Client { url.parse().map_err(|e| Error::UrlParse(e, url.to_owned())) } + /// Override the maximum bytes per request batch. Useful when the server's + /// body limit differs from [DEFAULT_MAX_BATCH_BYTES]. + pub fn with_max_batch_bytes(mut self, max: usize) -> Self { + self.max_batch_bytes = max; + self + } + /// Wait until the server is healthy. /// /// Returns either `Ok(elapsed)` or the `Error` from the last healthcheck attempt. From 4eea66c3514be57ba8ecd2571281252129f3a316 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 25 Jun 2026 19:58:04 +0000 Subject: [PATCH 11/14] Add in-cluster batch-load Job manifest and manual image build docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Running batch-load through kubectl port-forward causes spurious 400 "Failed to buffer the request body" errors — the apiserver-proxied tunnel truncates request bodies under sustained load. Running in-cluster talks straight to the Service and avoids this. Adds bench/k8s/ with: - job.yaml: ConfigMap + PVC + Job (resumable via persisted state/schemas) - batch-config.toml: sample config - README.md: buildx build/push commands (no CI needed) and deploy steps Co-Authored-By: Claude --- bench/k8s/README.md | 53 ++++++++++++++++++++++ bench/k8s/batch-config.toml | 20 +++++++++ bench/k8s/job.yaml | 89 +++++++++++++++++++++++++++++++++++++ 3 files changed, 162 insertions(+) create mode 100644 bench/k8s/README.md create mode 100644 bench/k8s/batch-config.toml create mode 100644 bench/k8s/job.yaml diff --git a/bench/k8s/README.md b/bench/k8s/README.md new file mode 100644 index 0000000..6c2207a --- /dev/null +++ b/bench/k8s/README.md @@ -0,0 +1,53 @@ +# Running batch-load in-cluster + +Running through `kubectl port-forward` causes spurious +`Failed to buffer the request body` (400) errors: the tunnel is a single +apiserver-proxied stream that stalls under sustained load, truncating request +bodies. Running in-cluster talks straight to the Service and avoids this. + +## Build and push the image (outside CI) + +The `batch-load` target in the repo-root `Dockerfile` produces a static musl +binary in a `scratch` image. Build it directly with buildx — no CI required: + +```sh +# from the repo root +REGISTRY=ghcr.io/wallaroolabs/plateau-bench +TAG=dev + +docker buildx build \ + --target batch-load \ + --platform linux/amd64 \ + --build-arg RUST_VERSION=1.84.1 \ + -t $REGISTRY:$TAG \ + --push \ + . +``` + +Notes: +- `--target batch-load` selects the bench stage (the default target builds the + plateau server). +- `--platform linux/amd64` matches a typical cluster; add/replace with + `linux/arm64` if your nodes are arm. +- `--push` uploads straight to the registry. Log in first + (`docker login ghcr.io`). Use any registry your cluster can pull from. + +## Deploy the Job + +Edit `job.yaml`: +- `image:` → the tag you just pushed +- `PLATEAU_URL` → `http://..svc.cluster.local:3030` + +Then: + +```sh +kubectl apply -n -f bench/k8s/job.yaml +kubectl logs -n -f job/batch-load +``` + +The schemas dir and state file live on a PVC, so the Job resumes the same +topic pool and schedule if it restarts. For a throwaway run, replace the +`persistentVolumeClaim` volume with `emptyDir: {}`. + +To restart with a clean slate, delete the PVC (`kubectl delete pvc +batch-load-data`) before re-applying. diff --git a/bench/k8s/batch-config.toml b/bench/k8s/batch-config.toml new file mode 100644 index 0000000..0b95717 --- /dev/null +++ b/bench/k8s/batch-config.toml @@ -0,0 +1,20 @@ +# Sample batch-load configuration. Mounted into the Job via a ConfigMap. +# +# All range values (columns / partitions / rows) are drawn from a clamped +# normal distribution. See `batch-load --help` for full field docs. + +speed = 60.0 # compress time: 1h intervals fire every minute +schemas_dir = "/data/batch-schemas" +state_file = "/data/batch-state.json" + +[topics] +count = 200 # total topic pool +active = 8 # topics writing at once +rotation_interval = "1h" # how often the active window advances +columns_min = 3 +columns_max = 35 +partitions_min = 1 +partitions_max = 8 +rows_min = 1000 +rows_max = 50000 +batch_interval = "1h" # interval between batches per partition diff --git a/bench/k8s/job.yaml b/bench/k8s/job.yaml new file mode 100644 index 0000000..154e03b --- /dev/null +++ b/bench/k8s/job.yaml @@ -0,0 +1,89 @@ +# In-cluster batch-load runner. +# +# Talks directly to the plateau Service (no port-forward), so the spurious +# "Failed to buffer the request body" errors from a choked tunnel go away. +# +# Edit before applying: +# - image: the tag you built and pushed +# - PLATEAU_URL: http://..svc.cluster.local:3030 +# - namespace: via `kubectl apply -n -f job.yaml` +# +# The schemas dir and state file live on a PVC so a restarted Job resumes the +# same topic pool and schedule. For a throwaway run, swap the PVC for an +# `emptyDir: {}` volume. +--- +apiVersion: v1 +kind: ConfigMap +metadata: + name: batch-load-config +data: + batch-config.toml: | + speed = 60.0 + schemas_dir = "/data/batch-schemas" + state_file = "/data/batch-state.json" + + [topics] + count = 200 + active = 8 + rotation_interval = "1h" + columns_min = 3 + columns_max = 35 + partitions_min = 1 + partitions_max = 8 + rows_min = 1000 + rows_max = 50000 + batch_interval = "1h" +--- +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: batch-load-data +spec: + accessModes: ["ReadWriteOnce"] + resources: + requests: + storage: 1Gi +--- +apiVersion: batch/v1 +kind: Job +metadata: + name: batch-load +spec: + backoffLimit: 6 + template: + metadata: + labels: + app: batch-load + spec: + restartPolicy: OnFailure + containers: + - name: batch-load + image: ghcr.io/wallaroolabs/plateau-bench:dev + command: ["./batch-load"] + args: + - "--config=/config/batch-config.toml" + - "--url=$(PLATEAU_URL)" + env: + - name: PLATEAU_URL + value: "http://plateau.default.svc.cluster.local:3030" + - name: RUST_LOG + value: "warn,bench=info" + volumeMounts: + - name: config + mountPath: /config + - name: data + mountPath: /data + resources: + requests: + cpu: "500m" + memory: "512Mi" + limits: + cpu: "2" + memory: "2Gi" + volumes: + - name: config + configMap: + name: batch-load-config + - name: data + persistentVolumeClaim: + claimName: batch-load-data From a8eed7483fbc83ff4b7820cb8928f315def5089d Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 25 Jun 2026 20:11:52 +0000 Subject: [PATCH 12/14] Fix Dockerfile: build bench package for batch-load binary batch-load lives in the bench package, not plateau. Build both with -p plateau -p bench so both binaries are available to copy. Co-Authored-By: Claude --- Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Dockerfile b/Dockerfile index 74285b9..8061271 100644 --- a/Dockerfile +++ b/Dockerfile @@ -11,7 +11,7 @@ COPY . . RUN \ if [ "${TARGETARCH}" = "amd64" ]; then ARCH=x86_64; elif [ "${TARGETARCH}" = "arm64" ]; then ARCH=aarch64; else exit 1; fi && \ rustup target add ${ARCH}-unknown-linux-musl && \ - cargo build --release --target ${ARCH}-unknown-linux-musl -p plateau --bin batch-load && \ + cargo build --release --target ${ARCH}-unknown-linux-musl -p plateau -p bench && \ cp target/${ARCH}-unknown-linux-musl/release/plateau target/release/plateau && \ cp target/${ARCH}-unknown-linux-musl/release/batch-load target/release/batch-load From 651687eef5f96da8e621768a605bf64c6e7ac405 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 25 Jun 2026 20:40:31 +0000 Subject: [PATCH 13/14] Add imagePullSecrets to batch-load Job for private ghcr.io image plateau-bench is private like the other org packages, so the Job needs the same ghcr.io pull secret the plateau pods use. Reference it via imagePullSecrets (placeholder name to be set per-cluster). Co-Authored-By: Claude --- bench/k8s/job.yaml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/bench/k8s/job.yaml b/bench/k8s/job.yaml index 154e03b..84cbec4 100644 --- a/bench/k8s/job.yaml +++ b/bench/k8s/job.yaml @@ -56,6 +56,11 @@ spec: app: batch-load spec: restartPolicy: OnFailure + # Reuse the same ghcr.io pull secret the plateau pods use. Find its name + # with `kubectl get secrets -n --field-selector + # type=kubernetes.io/dockerconfigjson` and set it here. + imagePullSecrets: + - name: ghcr-pull containers: - name: batch-load image: ghcr.io/wallaroolabs/plateau-bench:dev From 37389a8a1fd64d34b917d60c857496a84e626e19 Mon Sep 17 00:00:00 2001 From: Frank Murphy Date: Fri, 26 Jun 2026 14:09:02 -0400 Subject: [PATCH 14/14] Tweaks to job, load spec --- bench/k8s/batch-config.toml | 2 +- bench/k8s/job.yaml | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/bench/k8s/batch-config.toml b/bench/k8s/batch-config.toml index 0b95717..24faaa8 100644 --- a/bench/k8s/batch-config.toml +++ b/bench/k8s/batch-config.toml @@ -3,7 +3,7 @@ # All range values (columns / partitions / rows) are drawn from a clamped # normal distribution. See `batch-load --help` for full field docs. -speed = 60.0 # compress time: 1h intervals fire every minute +speed = 200.0 # compress time: 1h intervals fire every minute schemas_dir = "/data/batch-schemas" state_file = "/data/batch-state.json" diff --git a/bench/k8s/job.yaml b/bench/k8s/job.yaml index 84cbec4..e0d0924 100644 --- a/bench/k8s/job.yaml +++ b/bench/k8s/job.yaml @@ -29,10 +29,10 @@ data: columns_min = 3 columns_max = 35 partitions_min = 1 - partitions_max = 8 + partitions_max = 2 rows_min = 1000 rows_max = 50000 - batch_interval = "1h" + batch_interval = "15m" --- apiVersion: v1 kind: PersistentVolumeClaim @@ -60,7 +60,7 @@ spec: # with `kubectl get secrets -n --field-selector # type=kubernetes.io/dockerconfigjson` and set it here. imagePullSecrets: - - name: ghcr-pull + - name: regcred containers: - name: batch-load image: ghcr.io/wallaroolabs/plateau-bench:dev @@ -70,7 +70,7 @@ spec: - "--url=$(PLATEAU_URL)" env: - name: PLATEAU_URL - value: "http://plateau.default.svc.cluster.local:3030" + value: "http://plateau.wallaroo.svc.cluster.local:3030" - name: RUST_LOG value: "warn,bench=info" volumeMounts: