krabka_bench_driver/soak/
drill.rs

1//! The disaster-recovery drill's side of a soak: the declared load run against
2//! the source cluster until it is lost, and the read-back of the cluster
3//! restored from the capture.
4//!
5//! `krabka-soak-drill` is the command-line front end. Everything it does
6//! besides parsing arguments and listening for `SIGTERM` is here, so the load
7//! and the read-back are exercised against a live broker by the
8//! `soak_drill` suite of `krabka-broker`.
9
10pub mod cli;
11#[cfg(test)]
12mod tests;
13
14use {
15    super::{
16        client::{
17            Endpoint,
18            load::{CommitTracker, LoadContext, prepare, producer_name, start},
19            now_ms,
20            readback::{commit_checks, read_back, round_trip},
21            windows::Windows,
22        },
23        contract::{Contract, Profile},
24        evidence::{Evidence, FINAL_COMMIT_CHECK},
25        recovery::{
26            Boundary, CommitHistory, LoadRecord, ProducerAcks, RestoredReadings, bound_evidence,
27            committed_by,
28        },
29    },
30    anyhow::{Context, Result, bail},
31    krabka_units::prelude::*,
32    serde::Serialize,
33    std::{fs, future::Future, path::Path, sync::Arc},
34    tokio::time,
35};
36
37/// A disaster-recovery contract's objectives and paths, in the units the
38/// drill's shell side reads: whole milliseconds and topic shapes.
39#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
40pub struct Objectives {
41    pub post_capture_load_ms: i64,
42    pub rpo_ms: i64,
43    pub rto_ms: i64,
44    pub paths: Vec<PathShape>,
45}
46
47/// The topic a path writes to, as the drill must create it.
48#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
49pub struct PathShape {
50    pub name: String,
51    pub storage: String,
52    pub partitions: i32,
53    pub min_insync_replicas: i16,
54}
55
56/// Read and validate the contract at `path`, which must be a
57/// `disaster_recovery` contract.
58///
59/// # Errors
60///
61/// The file cannot be read, does not parse, or is another profile.
62pub fn load_contract(path: &Path) -> Result<Contract> {
63    let contract = Contract::from_file(path)?;
64    if contract.profile != Profile::DisasterRecovery {
65        bail!(
66            "{} is a {:?} contract; the drill runs only disaster_recovery contracts",
67            path.display(),
68            contract.profile
69        );
70    }
71    Ok(contract)
72}
73
74/// The objectives and the paths of a disaster-recovery contract.
75///
76/// # Errors
77///
78/// The contract declares no `disaster_recovery` objectives.
79pub fn objectives(contract: &Contract) -> Result<Objectives> {
80    let declared = contract
81        .disaster_recovery
82        .as_ref()
83        .context("a disaster-recovery contract declares its objectives")?;
84    Ok(Objectives {
85        post_capture_load_ms: declared.post_capture_load.millis_i64(),
86        rpo_ms: declared.rpo.millis_i64(),
87        rto_ms: declared.rto.millis_i64(),
88        paths: contract
89            .paths
90            .iter()
91            .map(|path| PathShape {
92                name: path.name.clone(),
93                storage: serde_json::to_value(path.storage)
94                    .ok()
95                    .and_then(|value| value.as_str().map(str::to_owned))
96                    .unwrap_or_default(),
97                partitions: path.partitions,
98                min_insync_replicas: contract.cluster.min_insync_replicas,
99            })
100            .collect(),
101    })
102}
103
104/// Write `value` as JSON through a rename, so a reader never sees half a file.
105///
106/// # Errors
107///
108/// The file cannot be written or renamed into place.
109pub fn write_json(path: &Path, value: &impl Serialize) -> Result<()> {
110    let partial = path.with_extension("partial");
111    fs::write(&partial, serde_json::to_vec_pretty(value)?)
112        .with_context(|| format!("write {}", partial.display()))?;
113    fs::rename(&partial, path).with_context(|| format!("rename to {}", path.display()))
114}
115
116/// Read a load record that [`run_load`] wrote.
117///
118/// # Errors
119///
120/// The file cannot be read or is not a load record.
121pub fn read_load(path: &Path) -> Result<LoadRecord> {
122    let text = fs::read_to_string(path).with_context(|| format!("read {}", path.display()))?;
123    serde_json::from_str(&text).with_context(|| format!("parse load record {}", path.display()))
124}
125
126/// Run the contract's load until `lost` completes or the contract's duration
127/// elapses, and record every acknowledgement and commit with its time.
128///
129/// When `lost` completes first, every send in flight is abandoned, because no
130/// broker is left to answer it.
131///
132/// # Errors
133///
134/// Returns a consumer startup error, or source loss during startup, before
135/// starting the measured load.
136pub async fn run_load(
137    contract: &Contract,
138    endpoint: Endpoint,
139    broker_image: String,
140    lost: impl Future<Output = ()>,
141) -> Result<LoadRecord> {
142    tokio::pin!(lost);
143    let prepared = tokio::select! {
144        biased;
145        () = &mut lost => bail!("source cluster lost before the measured load started"),
146        prepared = prepare(contract, &endpoint) => prepared.context("prepare the drill consumers")?,
147    };
148    let ctx = Arc::new(
149        LoadContext::builder()
150            .contract(contract.clone())
151            .endpoint(endpoint)
152            .windows(Arc::new(Windows::new(Vec::new())))
153            .record_ack_times(true)
154            .build(),
155    );
156    let started = now_ms();
157    let running = start(&ctx, prepared).context("build the drill's latency histograms")?;
158    tracing::info!(contract = %contract.name, "drill load started");
159    tokio::select! {
160        () = &mut lost => {
161            ctx.abandon();
162            tracing::info!("the source cluster is lost; abandoning the sends in flight");
163        }
164        () = time::sleep(contract.duration.to_std()) => {
165            ctx.stop();
166            tracing::info!("the load ran its whole duration without a loss");
167        }
168    }
169    let ended = now_ms();
170    let trackers = running.trackers.clone();
171    let outcome = running.finish().await.context("a drill load task failed")?;
172    let paths = outcome
173        .measurements(contract, Time::from_millis(ended - started))
174        .context("build the drill's latency histograms")?;
175    Ok(LoadRecord {
176        acks: outcome
177            .producers
178            .iter()
179            .map(|output| ProducerAcks {
180                producer: output.ledger.producer.clone(),
181                acks: output.acked_at.clone(),
182            })
183            .collect(),
184        commits: trackers
185            .iter()
186            .map(|tracker| CommitHistory {
187                group: tracker.group.clone(),
188                topic: tracker.topic.clone(),
189                partitions: tracker.partitions,
190                acks: tracker.history(),
191            })
192            .collect(),
193        run: Evidence {
194            contract: contract.name.clone(),
195            contract_digest: contract.digest(),
196            seed: contract.seed,
197            broker_image,
198            started_unix_ms: started,
199            ended_unix_ms: ended,
200            producers: outcome
201                .producers
202                .into_iter()
203                .map(|output| output.ledger)
204                .collect(),
205            paths,
206            ..Evidence::default()
207        },
208    })
209}
210
211/// Read every path of the restored cluster back, check each committing
212/// group's restored offsets against its last commit before the boundary, and
213/// return the evidence the judge grades.
214pub async fn read_back_restored(
215    contract: &Contract,
216    load: &LoadRecord,
217    endpoint: &Endpoint,
218    boundary: Boundary,
219    serving_unix_ms: Option<i64>,
220) -> Evidence {
221    let mut readers = Vec::new();
222    for path in &contract.paths {
223        for &isolation in &path.consumers {
224            readers.push(read_back(endpoint, path, isolation, &producer_name(&path.name)).await);
225        }
226    }
227    let mut commits = Vec::new();
228    for history in &load.commits {
229        let tracker = CommitTracker::new(
230            history.group.clone(),
231            history.topic.clone(),
232            history.partitions,
233        );
234        let acknowledged = committed_by(&history.acks, boundary.capture_started_unix_ms);
235        commits.extend(commit_checks(endpoint, &tracker, FINAL_COMMIT_CHECK, &acknowledged).await);
236    }
237    bound_evidence(load, boundary, RestoredReadings {
238        readers,
239        commits,
240        serving_unix_ms,
241        ended_unix_ms: now_ms(),
242    })
243}
244
245/// `krabka-soak-drill load`: run the load until `lost`, report each
246/// producer's totals, and write the record to `out`.
247///
248/// # Errors
249///
250/// Consumer preparation fails, the source is lost before load starts, or the
251/// record cannot be written.
252pub async fn record_load(
253    contract: &Contract,
254    endpoint: Endpoint,
255    broker_image: String,
256    lost: impl Future<Output = ()>,
257    out: &Path,
258) -> Result<LoadRecord> {
259    let record = run_load(contract, endpoint, broker_image, lost).await?;
260    for ledger in &record.run.producers {
261        tracing::info!(
262            producer = %ledger.producer,
263            attempted = ledger.attempted,
264            acknowledged = ledger.acked.len(),
265            failed = ledger.failed.len(),
266            "drill load recorded",
267        );
268    }
269    write_json(out, &record)?;
270    Ok(record)
271}
272
273/// `krabka-soak-drill readback`: read the restored cluster back against the
274/// load recorded at `load`, report each path's split at the boundary, and
275/// write the evidence to `out`.
276///
277/// # Errors
278///
279/// The load record cannot be read, or the evidence cannot be written.
280pub async fn record_readback(
281    contract: &Contract,
282    load: &Path,
283    endpoint: &Endpoint,
284    boundary: Boundary,
285    serving_unix_ms: Option<i64>,
286    out: &Path,
287) -> Result<Evidence> {
288    let record = read_load(load)?;
289    let evidence = read_back_restored(contract, &record, endpoint, boundary, serving_unix_ms).await;
290    for path in evidence
291        .recovery
292        .iter()
293        .flat_map(|recovery| &recovery.paths)
294    {
295        tracing::info!(
296            path = %path.path,
297            acknowledged = path.acknowledged,
298            before_boundary = path.before_boundary,
299            lost_beyond_boundary = path.lost_beyond_boundary.len(),
300            oldest_loss_unix_ms = ?path.oldest_loss_unix_ms,
301            "restored cluster read back",
302        );
303    }
304    write_json(out, &evidence)?;
305    Ok(evidence)
306}
307
308/// `krabka-soak-drill probe`: write one record to `topic`'s `partition` and
309/// read it back, and describe the outcome in the line the gate checks for.
310///
311/// # Errors
312///
313/// The write or the read fails.
314pub async fn probe(endpoint: &Endpoint, topic: &str, partition: i32) -> Result<String> {
315    let offset = round_trip(endpoint, topic, partition)
316        .await
317        .map_err(anyhow::Error::msg)?;
318    Ok(format!(
319        "{topic}-{partition}: wrote and read back offset {offset}"
320    ))
321}
krabka_bench_driver/soak/drill.rs

krabka_bench_driver/soak/
drill.rs
