diff --git a/.github/workflows/bench-image.yml b/.github/workflows/bench-image.yml new file mode 100644 index 0000000..98e164d --- /dev/null +++ b/.github/workflows/bench-image.yml @@ -0,0 +1,54 @@ +name: Bench Image + +on: + workflow_dispatch: {} + +env: + BENCH_REGISTRY_IMAGE: ghcr.io/wallaroolabs/plateau-bench + RUST_VERSION: 1.84.1 + +jobs: + build-and-push: + runs-on: + - ubuntu-latest + steps: + - 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: 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 + uses: docker/build-push-action@v6 + with: + target: batch-load + tags: ${{ steps.meta.outputs.tags }} + labels: ${{ steps.meta.outputs.labels }} + push: true + build-args: | + RUST_VERSION=${{ env.RUST_VERSION }} diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 54d23ff..0acdbde 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -76,6 +76,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 +135,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: | @@ -199,3 +201,4 @@ jobs: - name: Inspect image run: | docker buildx imagetools inspect ${{ env.REGISTRY_IMAGE }}:${{ steps.meta.outputs.version }} + diff --git a/Cargo.lock b/Cargo.lock index 01aea77..9343d61 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -524,8 +524,10 @@ dependencies = [ "arrow-ipc", "arrow-schema", "async-trait", + "clap", "futures", "hdrhistogram", + "humantime", "humantime-serde", "plateau-client", "plateau-server", @@ -537,6 +539,7 @@ dependencies = [ "serde", "serde_json", "tokio", + "toml 0.8.12", "tracing", "tracing-subscriber", ] diff --git a/Dockerfile b/Dockerfile index a1aa43c..8061271 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 -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 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"] diff --git a/bench/Cargo.toml b/bench/Cargo.toml index 597d7ce..ea26749 100644 --- a/bench/Cargo.toml +++ b/bench/Cargo.toml @@ -10,8 +10,11 @@ authors.workspace = true [dependencies] 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" @@ -38,3 +41,10 @@ plateau-client = { workspace = true, features = ["health"] } [[bin]] name = "cv" + +[[bin]] +name = "load" + +[[bin]] +name = "batch-load" +path = "src/bin/batch_load.rs" 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..24faaa8 --- /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 = 200.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..e0d0924 --- /dev/null +++ b/bench/k8s/job.yaml @@ -0,0 +1,94 @@ +# 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 = 2 + rows_min = 1000 + rows_max = 50000 + batch_interval = "15m" +--- +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 + # 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: regcred + 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.wallaroo.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 diff --git a/bench/src/batch.rs b/bench/src/batch.rs new file mode 100644 index 0000000..687256c --- /dev/null +++ b/bench/src/batch.rs @@ -0,0 +1,708 @@ +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 arrow_array::RecordBatch; +use arrow_ipc::writer::FileWriter; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use plateau_client::{Client, Error, MultiChunk}; +use reqwest::StatusCode; +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::Now; + +// ── Config ──────────────────────────────────────────────────────────────────── + +#[derive(Debug, Deserialize)] +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, + /// 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, +} + +#[derive(Debug, Deserialize)] +pub struct BatchConfig { + /// Speed multiplier: 60.0 means 1h intervals fire every 1 minute. + #[serde(default = "default_speed")] + pub speed: f64, + /// Path to the state file. + pub state_file: Option, + /// Directory where generated topic schema files are stored. + pub schemas_dir: Option, + pub topics: TopicsConfig, +} + +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)?) + } +} + +/// 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)] +pub struct BatchState { + /// RFC 3339 timestamp of first-ever run start; schedule anchor. + pub origin: Option, + /// 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, +} + +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); + } +} + +// ── 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 { + 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 `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))]; + 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 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 mut generated = 0; + for idx in 0..t.count { + let path = schema_path(schemas_dir, idx); + if path.exists() { + 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 generated > 0 { + info!("generated {generated} topic schema files in {schemas_dir:?}"); + } + + Ok(()) +} + +// ── Deterministic batch 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 ──────────────────────────────────────────────────────────── +// +// Box is !Send, so the sampler lives in a std::thread and +// communicates with the async worker via channels. + +struct SamplerRequest { + seed: u64, + rows_min: usize, + rows_max: usize, +} + +struct SamplerThread { + sample_path: PathBuf, + req_rx: std::sync::mpsc::Receiver, + result_tx: mpsc::Sender, +} + +impl SamplerThread { + fn run(self) { + 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); + // 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; + } + } + } +} + +// ── Per-partition async worker ──────────────────────────────────────────────── + +struct PartitionWorker { + client: Client, + topic: String, + partition: String, + sample_path: PathBuf, + rows_min: usize, + rows_max: usize, + batch_period: Duration, + stagger_offset: Duration, + state: Arc>, + state_path: PathBuf, +} + +impl PartitionWorker { + 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 { + sample_path: self.sample_path.clone(), + req_rx, + result_tx, + }; + thread::spawn(move || st.run()); + + let resume_from = { + let s = self.state.lock().await; + s.last_batch_for(&self.topic, &self.partition) + }; + + 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() + } + }; + + 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 { + let fire_at = origin + + self.stagger_offset + + self.batch_period * batch_idx as u32; + + 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; + } + } + + let seed = batch_seed(&self.topic, &self.partition, batch_idx); + 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 { + Some(m) => m, + 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) + .await; + + match r { + 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; + continue; + } + Err(e) => { + warn!( + "{}/{} batch {} failed ({} rows): {}", + self.topic, self.partition, batch_idx, batch_rows, e + ); + } + } + + { + 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; + } + } +} + +// ── Window management ───────────────────────────────────────────────────────── + +#[allow(clippy::too_many_arguments)] +fn spawn_window( + client: &Client, + topics: &TopicsConfig, + schemas_dir: &Path, + state: &Arc>, + state_path: &Path, + window_start: usize, + batch_period: Duration, + speed: f64, + schemas_seed: u64, +) -> Vec> { + let window_end = (window_start + topics.active).min(topics.count); + + // 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 { + Duration::ZERO + }; + + let mut handles = vec![]; + let mut slot = 0usize; + + 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..tp.partitions { + let worker = PartitionWorker { + client: client.clone(), + topic: name.clone(), + partition: format!("partition-{p}"), + sample_path: sample_path.clone(), + rows_min: tp.rows_min, + rows_max: tp.rows_max, + 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<()> { + // 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?; + + let schemas_dir = config + .schemas_dir + .clone() + .unwrap_or_else(|| PathBuf::from("batch-schemas")); + + let mut state = BatchState::load(state_path); + + // 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()); + } + if state.window_since.is_none() { + state.window_since = Some(humantime::format_rfc3339(SystemTime::now()).to_string()); + } + state.save(state_path)?; + + let schemas_seed = state.schemas_seed.unwrap(); + + // Generate any missing schema files. + ensure_schemas(&schemas_dir, &config.topics, schemas_seed)?; + + let state = Arc::new(Mutex::new(state)); + + let batch_period = config.topics.batch_interval.div_f64(config.speed); + let rotation_period = config.topics.rotation_interval.div_f64(config.speed); + + // 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 { + (elapsed.as_nanos() / rotation_period.as_nanos()) as usize + }; + 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)?; + } + } + + 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, + schemas_seed, + ); + + // Progress reporter. + 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")); + } + } + }); + + // Rotation loop. + loop { + tokio::time::sleep(rotation_period).await; + + // 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, + 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 new file mode 100644 index 0000000..d8071b9 --- /dev/null +++ b/bench/src/bin/batch_load.rs @@ -0,0 +1,70 @@ +use std::path::PathBuf; + +use clap::Parser; +use tracing_subscriber::{fmt, EnvFilter}; + +use bench::batch::{run_batch, BatchConfig}; + +/// Batch-job load generator for an existing Plateau server. +/// +/// 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 # compress time: 1h intervals fire every 1 minute +/// schemas_dir = "batch-schemas" +/// +/// [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) +/// 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 +/// +/// All values drawn from a range use a clamped normal distribution. +#[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. + #[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/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..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; @@ -80,19 +81,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 +92,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 +111,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 +206,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(); 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.