krabka_bench_driver/soak/client/
load.rs

1//! The load: one paced producer per path and one live consumer group per path
2//! and isolation level.
3//!
4//! Every producer numbers what it sends from zero and writes
5//! [`record_tag`]`(producer, seq)` at the head of each value, so a reader can
6//! say which send a record came from. Sends are pipelined -- each tick
7//! enqueues one record and the acknowledgement is awaited beside the next
8//! ticks -- because a sequential send caps the rate at one per round trip,
9//! below the 72-hour contract's rates. Missed ticks are skipped, not caught
10//! up, so a stall during a fault never turns into a burst that inflates the
11//! steady rate after it.
12//!
13//! The transactional producer is the exception. It stops ticking while it
14//! waits for a transaction's deliveries and for its commit or abort, every
15//! transaction, on a healthy cluster as much as during a fault. Skipping the
16//! ticks missed in those waits would hold it below its rate, so it catches
17//! them up instead.
18
19use {
20    super::{
21        Endpoint, now_ms,
22        windows::{Latencies, Windows, measurement},
23    },
24    crate::{
25        scenario::{Acks as ContractAcks, Compression as ContractCompression},
26        soak::{
27            contract::{Contract, Isolation, LoadPath, Semantics},
28            evidence::{PathMeasurement, ProducerLedger, TxnLedger, TxnOutcome, record_tag},
29            ledger::SeqSet,
30            recovery::{Ack, CommitAck},
31        },
32    },
33    bytes::Bytes,
34    futures::{
35        StreamExt as _,
36        future::{OptionFuture, try_join_all},
37        stream::FuturesUnordered,
38    },
39    hdrhistogram::CreationError,
40    krabka_client_consumer::{
41        AutoOffsetReset, Consumer, ConsumerError, IsolationLevel, OffsetAndMetadata,
42    },
43    krabka_client_producer::{
44        Acks, Compression, Producer, ProducerError, ProducerRecord, RecordMetadata,
45    },
46    krabka_units::prelude::*,
47    std::{
48        collections::{BTreeMap, HashMap, VecDeque, btree_map::Entry},
49        future::Future,
50        pin::Pin,
51        str,
52        sync::{
53            Arc, Mutex,
54            atomic::{AtomicBool, Ordering},
55        },
56        time::{Duration, Instant},
57    },
58    tokio::{
59        sync::Notify,
60        task::{JoinError, JoinHandle},
61        time::{self, MissedTickBehavior, interval, sleep},
62    },
63};
64
65/// How long a client that could not be built waits before it tries again.
66const RECONNECT_BACKOFF: Duration = Duration::from_secs(1);
67
68/// What every load task shares.
69pub struct LoadContext {
70    pub contract: Contract,
71    pub endpoint: Endpoint,
72    pub windows: Arc<Windows>,
73    /// Whether each producer records when each of its sends was acknowledged.
74    /// A disaster-recovery drill needs it to place every record on one side of
75    /// the capture boundary; a 72-hour soak does not, and could not afford it.
76    record_ack_times: bool,
77    running: AtomicBool,
78    /// Every path's acknowledged records so far, for a lifecycle checkpoint
79    /// that reconciles them while the load runs. The producers' own ledgers
80    /// are only handed back when the load stops.
81    acknowledged: Mutex<BTreeMap<String, SeqSet>>,
82    abandoned: AtomicBool,
83    abandon: Notify,
84}
85
86#[bon::bon]
87impl LoadContext {
88    #[builder]
89    pub fn new(
90        contract: Contract,
91        endpoint: Endpoint,
92        windows: Arc<Windows>,
93        #[builder(default)] record_ack_times: bool,
94    ) -> Self {
95        Self {
96            contract,
97            endpoint,
98            windows,
99            record_ack_times,
100            running: AtomicBool::new(true),
101            acknowledged: Mutex::new(BTreeMap::new()),
102            abandoned: AtomicBool::new(false),
103            abandon: Notify::new(),
104        }
105    }
106
107    /// Record that `path`'s producer was acknowledged `seq`.
108    fn acknowledge(&self, path: &str, seq: u64) {
109        if let Ok(mut acknowledged) = self.acknowledged.lock() {
110            acknowledged.entry(path.to_owned()).or_default().insert(seq);
111        }
112    }
113
114    /// `path`'s acknowledged records so far.
115    pub fn acknowledged(&self, path: &str) -> SeqSet {
116        self.acknowledged
117            .lock()
118            .ok()
119            .and_then(|acknowledged| acknowledged.get(path).cloned())
120            .unwrap_or_default()
121    }
122
123    fn running(&self) -> bool { self.running.load(Ordering::Relaxed) }
124
125    fn is_abandoned(&self) -> bool { self.abandoned.load(Ordering::Relaxed) }
126
127    /// Stop every producer and consumer at its next step.
128    pub fn stop(&self) { self.running.store(false, Ordering::Relaxed); }
129
130    /// Stop every idempotent producer and every consumer now, without waiting
131    /// for a send in flight or closing a client: the cluster they talk to is
132    /// gone. A send in flight is then neither acknowledged nor failed, which
133    /// the reconciler reads as a record that may or may not have been written.
134    pub fn abandon(&self) {
135        self.stop();
136        self.abandoned.store(true, Ordering::Relaxed);
137        self.abandon.notify_waiters();
138    }
139
140    /// Resolves once [`Self::abandon`] has been called.
141    async fn abandoned(&self) {
142        let notified = self.abandon.notified();
143        if self.is_abandoned() {
144            return;
145        }
146        notified.await;
147    }
148}
149
150/// The one producer of `path`, named as the evidence names it.
151#[must_use]
152pub fn producer_name(path: &str) -> String { format!("{path}-0") }
153
154/// The consumer group of `path` read at `isolation`.
155#[must_use]
156pub fn group_name(path: &str, isolation: Isolation) -> String { format!("soak-{path}-{isolation}") }
157
158/// A record value: the tag, padded with `x` to `payload` bytes. A tag longer
159/// than the payload is sent whole.
160#[must_use]
161pub fn record_value(producer: &str, seq: u64, payload: usize) -> Bytes {
162    let mut value = record_tag(producer, seq).into_bytes();
163    if value.len() < payload {
164        value.resize(payload, b'x');
165    }
166    Bytes::from(value)
167}
168
169/// The key of record `seq` on a compacted path: `<producer>/k<seq mod keys>`.
170#[must_use]
171pub fn compacted_key(producer: &str, seq: u64, key_space: u64) -> String {
172    format!("{producer}/k{}", seq % key_space.max(1))
173}
174
175/// The key index a compacted key carries.
176#[must_use]
177pub fn key_index(key: &[u8]) -> Option<u64> {
178    let text = str::from_utf8(key).ok()?;
179    let (_, index) = text.rsplit_once("/k")?;
180    index.parse().ok()
181}
182
183/// Whether a failed compacted send has been superseded.
184///
185/// A later value of its key was already sent, so resending it would put an
186/// older value after a newer one, which is a regression the application itself
187/// would cause.
188#[must_use]
189pub fn superseded(seq: u64, key_space: u64, next_seq: u64) -> bool {
190    seq.saturating_add(key_space) < next_seq
191}
192
193/// Whether transaction `index` (counting from zero) is aborted on purpose.
194#[must_use]
195pub fn aborted_on_purpose(index: u64, abort_every: u32) -> bool {
196    (index + 1).is_multiple_of(u64::from(abort_every.max(1)))
197}
198
199/// What a producer task returns when the load stops.
200pub struct ProducerOutput {
201    pub ledger: ProducerLedger,
202    pub latencies: Latencies,
203    /// When each acknowledged send was acknowledged, in the order the acks
204    /// arrived. Empty unless the context records ack times.
205    pub acked_at: Vec<Ack>,
206}
207
208fn acks(acks: ContractAcks) -> Acks {
209    match acks {
210        ContractAcks::None => Acks::Zero,
211        ContractAcks::Leader => Acks::One,
212        ContractAcks::All => Acks::All,
213    }
214}
215
216fn compression(compression: ContractCompression) -> Compression {
217    match compression {
218        ContractCompression::None => Compression::None,
219        ContractCompression::Gzip => Compression::Gzip,
220        ContractCompression::Snappy => Compression::Snappy,
221        ContractCompression::Lz4 => Compression::Lz4,
222        ContractCompression::Zstd => Compression::Zstd,
223    }
224}
225
226/// A producer configured as the contract declares, transactional when
227/// `transactional_id` is set.
228async fn build_producer(
229    ctx: &LoadContext,
230    name: &str,
231    transactional_id: Option<String>,
232) -> Result<Producer, ProducerError> {
233    let settings = &ctx.contract.clients.producer;
234    let transaction_timeout = transactional_id
235        .as_ref()
236        .map(|_| ctx.contract.clients.transactions.timeout.to_std());
237    Producer::builder()
238        .bootstrap(ctx.endpoint.bootstrap.clone())
239        .maybe_security(ctx.endpoint.security.clone())
240        .client_id(name.to_owned())
241        .acks(acks(settings.acks))
242        .enable_idempotence(true)
243        .compression(compression(settings.compression))
244        .linger(settings.linger.to_std())
245        .batch_size(settings.batch_size.bytes_usize())
246        .request_timeout(settings.request_timeout.to_std())
247        .delivery_timeout(settings.delivery_timeout.to_std())
248        .maybe_transactional_id(transactional_id)
249        .maybe_transaction_timeout(transaction_timeout)
250        .build()
251        .await
252}
253
254/// A producer that is ready to send, or `None` once the load has stopped.
255/// A transactional one has run `InitProducerId`, which also fences and ends
256/// whatever transaction a previous producer of the same id left open.
257async fn connect_producer(ctx: &LoadContext, name: &str, transactional: bool) -> Option<Producer> {
258    let transactional_id = transactional.then(|| format!("soak-{name}"));
259    while ctx.running() {
260        match build_producer(ctx, name, transactional_id.clone()).await {
261            Ok(producer) if !transactional => return Some(producer),
262            Ok(producer) => match producer.init_transactions().await {
263                Ok(()) => return Some(producer),
264                Err(error) => {
265                    eprintln!("soak-contract: {name}: init_transactions: {error}");
266                    let _ = producer.close().await;
267                }
268            },
269            Err(error) => eprintln!("soak-contract: {name}: build producer: {error}"),
270        }
271        sleep(RECONNECT_BACKOFF).await;
272    }
273    None
274}
275
276/// One send's acknowledgement: its sequence number, which attempt it was,
277/// when it was enqueued, and how it ended.
278type Delivery =
279    Pin<Box<dyn Future<Output = (u64, u32, i64, Result<RecordMetadata, ProducerError>)> + Send>>;
280type Admission<'a> = Pin<Box<dyn Future<Output = Delivery> + Send + 'a>>;
281
282async fn enqueue(producer: &Producer, record: ProducerRecord, seq: u64, attempt: u32) -> Delivery {
283    let sent = now_ms();
284    match producer.enqueue(record).await {
285        Ok(handle) => Box::pin(async move { (seq, attempt, sent, handle.await) }),
286        Err(error) => Box::pin(async move { (seq, attempt, sent, Err(error)) }),
287    }
288}
289
290/// The pace of a pipelined producer: one tick per `period`, skipping the ticks
291/// it missed.
292fn pipelined_ticker(period: Duration) -> time::Interval {
293    let mut ticker = interval(period);
294    ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);
295    ticker
296}
297
298/// The pace of a transactional producer: one tick per `period` on average,
299/// with the ticks missed while a transaction ended fired at once.
300fn transactional_ticker(period: Duration) -> time::Interval {
301    let mut ticker = interval(period);
302    ticker.set_missed_tick_behavior(MissedTickBehavior::Burst);
303    ticker
304}
305
306fn empty_ledger(path: &LoadPath, producer: &str) -> ProducerLedger {
307    ProducerLedger {
308        path: path.name.clone(),
309        producer: producer.to_owned(),
310        semantics: path.semantics,
311        attempted: 0,
312        acked: SeqSet::new(),
313        failed: SeqSet::new(),
314        retried: SeqSet::new(),
315        transactions: Vec::new(),
316    }
317}
318
319/// The record `seq` of `path`.
320fn record(path: &LoadPath, producer: &str, seq: u64) -> ProducerRecord {
321    let key = path
322        .key_space
323        .filter(|_| path.semantics == Semantics::Compacted)
324        .map(|keys| Bytes::from(compacted_key(producer, seq, keys)));
325    ProducerRecord {
326        topic: path.name.clone(),
327        key,
328        value: Some(record_value(producer, seq, path.payload.bytes_usize())),
329        ..Default::default()
330    }
331}
332
333/// The idempotent or compacted producer of `path`, until the load stops.
334///
335/// A send that fails is sent again, up to `max_send_attempts` in all, and is
336/// recorded as failed and as retried. A compacted send whose key has a later
337/// value in flight is not sent again: see [`superseded`].
338///
339/// `latencies` is where its send latencies go.
340pub async fn run_idempotent(
341    ctx: Arc<LoadContext>,
342    path: LoadPath,
343    latencies: Latencies,
344) -> ProducerOutput {
345    let name = producer_name(&path.name);
346    let mut output = ProducerOutput {
347        ledger: empty_ledger(&path, &name),
348        latencies,
349        acked_at: Vec::new(),
350    };
351    let max_attempts = ctx.contract.clients.producer.max_send_attempts;
352    let Some(producer) = connect_producer(&ctx, &name, false).await else {
353        return output;
354    };
355    let mut ticker = pipelined_ticker(path.rate.period().to_std());
356    let mut inflight: FuturesUnordered<Delivery> = FuturesUnordered::new();
357    let mut next_seq = 0_u64;
358    let mut admission: Option<Admission<'_>> = None;
359    let mut retries = VecDeque::new();
360    loop {
361        if admission.is_none()
362            && let Some((seq, attempt)) = retries.pop_front()
363        {
364            admission = Some(Box::pin(enqueue(
365                &producer,
366                record(&path, &name, seq),
367                seq,
368                attempt,
369            )));
370        }
371        if !ctx.running() && inflight.is_empty() && admission.is_none() {
372            break;
373        }
374        tokio::select! {
375            () = ctx.abandoned() => break,
376            _ = ticker.tick(), if ctx.running() && admission.is_none() => {
377                let seq = next_seq;
378                next_seq += 1;
379                admission = Some(Box::pin(enqueue(&producer, record(&path, &name, seq), seq, 1)));
380            }
381            Some(delivery) = OptionFuture::from(admission.as_mut()), if admission.is_some() => {
382                admission = None;
383                inflight.push(delivery);
384            }
385            Some((seq, attempt, sent, result)) = inflight.next() => {
386                let at = now_ms();
387                match result {
388                    Ok(_) => {
389                        output.ledger.acked.insert(seq);
390                        ctx.acknowledge(&path.name, seq);
391                        output.latencies.record(ctx.windows.disturbed(at), Time::from_millis(at - sent));
392                        if ctx.record_ack_times {
393                            output.acked_at.push(Ack { seq, unix_ms: at });
394                        }
395                    }
396                    Err(error) => {
397                        output.ledger.failed.insert(seq);
398                        let stale = path
399                            .key_space
400                            .is_some_and(|keys| superseded(seq, keys, next_seq));
401                        if attempt < max_attempts && !stale {
402                            eprintln!("soak-contract: {name}: send {seq} attempt {attempt} failed, resending: {error}");
403                            output.ledger.retried.insert(seq);
404                            retries.push_back((seq, attempt + 1));
405                        } else {
406                            eprintln!("soak-contract: {name}: send {seq} attempt {attempt} failed: {error}");
407                        }
408                    }
409                }
410            }
411        }
412    }
413    output.ledger.attempted = next_seq;
414    drop(admission);
415    // A producer of an abandoned cluster is dropped, not closed: closing waits
416    // for the sends in flight, which can no longer be answered.
417    if !ctx.is_abandoned() {
418        let _ = producer.close().await;
419    }
420    output
421}
422
423/// The transactional producer of `path`, until the load stops.
424///
425/// Each transaction carries `records_per_transaction` paced sends, and every
426/// `abort_every`-th one is aborted on purpose. A transaction with a failed send
427/// is aborted. A commit or abort that returns an error leaves the producer in
428/// an unknown state, so it is closed and rebuilt under the same transactional
429/// id, whose `InitProducerId` settles what it left open. A commit that
430/// returned an error is recorded as [`TxnOutcome::Unknown`]; a transaction
431/// whose commit was never sent is aborted, whatever ended it.
432///
433/// `latencies` is where its send latencies go.
434pub async fn run_transactional(
435    ctx: Arc<LoadContext>,
436    path: LoadPath,
437    latencies: Latencies,
438) -> ProducerOutput {
439    let mut run = TransactionalRun {
440        name: producer_name(&path.name),
441        ledger: empty_ledger(&path, &producer_name(&path.name)),
442        latencies,
443        acked_at: Vec::new(),
444        ticker: transactional_ticker(path.rate.period().to_std()),
445        next_seq: 0,
446        index: 0,
447    };
448    if let Some(producer) = run.transact_until_stopped(&ctx, &path).await {
449        let _ = producer.close().await;
450    }
451    run.ledger.attempted = run.next_seq;
452    ProducerOutput {
453        ledger: run.ledger,
454        latencies: run.latencies,
455        acked_at: run.acked_at,
456    }
457}
458
459/// The state of a transactional producer across its transactions.
460struct TransactionalRun {
461    name: String,
462    ledger: ProducerLedger,
463    latencies: Latencies,
464    acked_at: Vec<Ack>,
465    ticker: time::Interval,
466    next_seq: u64,
467    index: u64,
468}
469
470impl TransactionalRun {
471    /// Run transactions until the load stops, reconnecting after a producer
472    /// becomes unusable, and return the producer still open at the end.
473    async fn transact_until_stopped(
474        &mut self,
475        ctx: &LoadContext,
476        path: &LoadPath,
477    ) -> Option<Producer> {
478        let mut producer: Option<Producer> = None;
479        while ctx.running() {
480            match producer.as_ref() {
481                None => producer = connect_producer(ctx, &self.name, true).await,
482                Some(current) => {
483                    if !self.transaction(ctx, path, current).await
484                        && let Some(stale) = producer.take()
485                    {
486                        let _ = stale.close().await;
487                    }
488                }
489            }
490        }
491        producer
492    }
493
494    /// Run one transaction on `producer`, and say whether the producer is
495    /// still usable.
496    async fn transaction(
497        &mut self,
498        ctx: &LoadContext,
499        path: &LoadPath,
500        producer: &Producer,
501    ) -> bool {
502        let settings = &ctx.contract.clients.transactions;
503        let name = self.name.clone();
504        let transaction = match producer.begin_transaction().await {
505            Ok(transaction) => transaction,
506            Err(error) => {
507                eprintln!("soak-contract: {name}: begin_transaction: {error}");
508                return false;
509            }
510        };
511        let start = self.next_seq;
512        let failed = self
513            .deliveries(ctx, path, |seq| {
514                enqueue(producer, record(path, &name, seq), seq, 1)
515            })
516            .await;
517        let index = self.index;
518        let (outcome, usable) = if failed || aborted_on_purpose(index, settings.abort_every) {
519            match transaction.abort().await {
520                Ok(()) => (TxnOutcome::Aborted, true),
521                Err(error) => {
522                    eprintln!("soak-contract: {name}: abort of transaction {index}: {error}");
523                    (TxnOutcome::Aborted, false)
524                }
525            }
526        } else {
527            match transaction.commit().await {
528                Ok(()) => (TxnOutcome::Committed, true),
529                Err(error) => {
530                    eprintln!("soak-contract: {name}: commit of transaction {index}: {error}");
531                    (TxnOutcome::Unknown, false)
532                }
533            }
534        };
535        self.ledger.transactions.push(TxnLedger {
536            id: index,
537            seqs: SeqSet::range(start, self.next_seq),
538            outcome,
539        });
540        self.index += 1;
541        usable
542    }
543
544    /// Observe acknowledgements alongside pacing, rather than after filling
545    /// the transaction, so the measured latency excludes the driver's wait.
546    async fn deliveries<F: Future<Output = Delivery> + Send>(
547        &mut self,
548        ctx: &LoadContext,
549        path: &LoadPath,
550        mut enqueue_record: impl FnMut(u64) -> F + Send,
551    ) -> bool {
552        let mut deliveries = FuturesUnordered::new();
553        let mut admission: Option<Pin<Box<F>>> = None;
554        let mut admitted = 0;
555        let count = ctx.contract.clients.transactions.records_per_transaction;
556        let mut failed = false;
557        loop {
558            if (admitted == count || !ctx.running()) && deliveries.is_empty() && admission.is_none()
559            {
560                break;
561            }
562            tokio::select! {
563                _ = self.ticker.tick(), if admitted < count && ctx.running() && admission.is_none() => {
564                    if !ctx.running() {
565                        continue;
566                    }
567                    let seq = self.next_seq;
568                    self.next_seq += 1;
569                    admitted += 1;
570                    admission = Some(Box::pin(enqueue_record(seq)));
571                }
572                Some(delivery) = OptionFuture::from(admission.as_mut()), if admission.is_some() => {
573                    admission = None;
574                    deliveries.push(delivery);
575                }
576                Some((seq, _, sent, result)) = deliveries.next(), if !deliveries.is_empty() => {
577                    let at = now_ms();
578                    match result {
579                        Ok(_) => {
580                            self.ledger.acked.insert(seq);
581                            ctx.acknowledge(&path.name, seq);
582                            self.latencies.record(ctx.windows.disturbed(at), Time::from_millis(at - sent));
583                            if ctx.record_ack_times {
584                                self.acked_at.push(Ack { seq, unix_ms: at });
585                            }
586                        }
587                        Err(error) => {
588                            eprintln!("soak-contract: {}: transactional send {seq} failed: {error}", self.name);
589                            self.ledger.failed.insert(seq);
590                            failed = true;
591                        }
592                    }
593                }
594                else => break,
595            }
596        }
597        failed
598    }
599}
600
601/// The last commit the coordinator acknowledged for one consumer group.
602pub struct CommitTracker {
603    pub group: String,
604    pub topic: String,
605    pub partitions: i32,
606    acknowledged: Mutex<BTreeMap<i32, i64>>,
607    /// Every acknowledged commit, in order, so a drill can ask which commit
608    /// was the last one before its capture boundary.
609    history: Mutex<Vec<CommitAck>>,
610}
611
612impl CommitTracker {
613    #[must_use]
614    pub fn new(group: String, topic: String, partitions: i32) -> Self {
615        Self {
616            group,
617            topic,
618            partitions,
619            acknowledged: Mutex::new(BTreeMap::new()),
620            history: Mutex::new(Vec::new()),
621        }
622    }
623
624    /// Every acknowledged commit so far, oldest first.
625    #[must_use]
626    pub fn history(&self) -> Vec<CommitAck> {
627        self.history
628            .lock()
629            .map(|history| history.clone())
630            .unwrap_or_default()
631    }
632
633    /// The last acknowledged commit of every partition committed so far.
634    pub fn snapshot(&self) -> BTreeMap<i32, i64> {
635        self.acknowledged
636            .lock()
637            .map(|acked| acked.clone())
638            .unwrap_or_default()
639    }
640
641    fn acknowledge(&self, positions: &BTreeMap<i32, i64>) {
642        if let Ok(mut acked) = self.acknowledged.lock() {
643            acked.extend(positions);
644        }
645        if let Ok(mut history) = self.history.lock() {
646            history.push(CommitAck {
647                unix_ms: now_ms(),
648                offsets: positions.clone(),
649            });
650        }
651    }
652}
653
654async fn build_consumer(
655    contract: &Contract,
656    endpoint: &Endpoint,
657    path: &LoadPath,
658    isolation: Isolation,
659) -> Result<Consumer, ConsumerError> {
660    let consumer = &contract.clients.consumer;
661    Consumer::builder()
662        .bootstrap(endpoint.bootstrap.clone())
663        .maybe_security(endpoint.security.clone())
664        .client_id(group_name(&path.name, isolation))
665        .group_id(group_name(&path.name, isolation))
666        .subscribe(vec![path.name.clone()])
667        .auto_offset_reset(AutoOffsetReset::Earliest)
668        .isolation_level(isolation_level(isolation))
669        .session_timeout(consumer.session_timeout)
670        .enable_auto_commit(false)
671        .default_api_timeout(consumer.session_timeout)
672        .build()
673        .await
674}
675
676#[must_use]
677pub fn isolation_level(isolation: Isolation) -> IsolationLevel {
678    match isolation {
679        Isolation::ReadCommitted => IsolationLevel::ReadCommitted,
680        Isolation::ReadUncommitted => IsolationLevel::ReadUncommitted,
681    }
682}
683
684/// The live consumer group of `path` at `isolation`, until the load stops.
685///
686/// It measures create-time-to-poll latency for every record it polls, and,
687/// with a `tracker`, commits what it has polled every `commit_interval` and
688/// records each commit the coordinator acknowledged.
689pub async fn run_consumer(
690    ctx: Arc<LoadContext>,
691    path: LoadPath,
692    isolation: Isolation,
693    tracker: Option<Arc<CommitTracker>>,
694    mut consumer: Consumer,
695    mut latencies: Latencies,
696) -> Latencies {
697    let settings = ctx.contract.clients.consumer.clone();
698    let group = group_name(&path.name, isolation);
699    let mut positions: BTreeMap<i32, i64> = BTreeMap::new();
700    let mut committed: BTreeMap<i32, i64> = BTreeMap::new();
701    let mut last_commit = Instant::now();
702    while ctx.running() {
703        let polled = tokio::select! {
704            () = ctx.abandoned() => break,
705            polled = consumer.poll(settings.poll_timeout) => polled,
706        };
707        match polled {
708            Ok(records) => {
709                let at = now_ms();
710                let disturbed = ctx.windows.disturbed(at);
711                for record in records {
712                    latencies.record(disturbed, Time::from_millis((at - record.timestamp).max(0)));
713                    let next = positions.entry(record.partition).or_insert(0);
714                    *next = (*next).max(record.offset + 1);
715                }
716            }
717            Err(error) => {
718                eprintln!("soak-contract: {group}: poll: {error}");
719                sleep(Duration::from_millis(200)).await;
720            }
721        }
722        if let Some(tracker) = &tracker
723            && last_commit.elapsed() >= settings.commit_interval.to_std()
724            && positions != committed
725        {
726            last_commit = Instant::now();
727            let offsets: HashMap<(String, i32), OffsetAndMetadata> = positions
728                .iter()
729                .map(|(partition, offset)| {
730                    (
731                        (path.name.clone(), *partition),
732                        OffsetAndMetadata::new(*offset),
733                    )
734                })
735                .collect();
736            let commit = tokio::select! {
737                () = ctx.abandoned() => break,
738                commit = consumer.commit_offsets_sync(offsets) => commit,
739            };
740            match commit {
741                Ok(()) => {
742                    tracker.acknowledge(&positions);
743                    committed.clone_from(&positions);
744                }
745                Err(error) => eprintln!("soak-contract: {group}: commit: {error}"),
746            }
747        }
748    }
749    if !ctx.is_abandoned() {
750        let _ = consumer.close().await;
751    }
752    latencies
753}
754
755/// Every task of a running load: one producer per path, one consumer group
756/// per path and isolation level, and the commit tracker of every group that
757/// commits.
758pub struct RunningLoad {
759    pub ctx: Arc<LoadContext>,
760    pub trackers: Vec<Arc<CommitTracker>>,
761    producers: Vec<JoinHandle<ProducerOutput>>,
762    consumers: Vec<(String, JoinHandle<Latencies>)>,
763}
764
765/// What a load left behind once every task ended.
766pub struct LoadOutcome {
767    pub producers: Vec<ProducerOutput>,
768    /// Every consumer group's latency, merged per path.
769    pub consume_latency: BTreeMap<String, Latencies>,
770}
771
772/// Consumers whose group assignment completed before the measured load.
773pub struct PreparedLoad {
774    consumers: Vec<(LoadPath, Isolation, Consumer)>,
775}
776
777/// Join every consumer group before starting the measured load and its fault
778/// clock, so coordinator startup does not age the first produced records.
779///
780/// # Errors
781///
782/// Returns a consumer startup error before any producer starts.
783pub async fn prepare(
784    contract: &Contract,
785    endpoint: &Endpoint,
786) -> Result<PreparedLoad, ConsumerError> {
787    let consumers = contract.paths.iter().flat_map(|path| {
788        path.consumers.iter().map(move |&isolation| async move {
789            let consumer = build_consumer(contract, endpoint, path, isolation).await?;
790            Ok::<_, ConsumerError>((path.clone(), isolation, consumer))
791        })
792    });
793    Ok(PreparedLoad {
794        consumers: try_join_all(consumers).await?,
795    })
796}
797
798/// Start the load `ctx` declares. It runs until [`LoadContext::stop`] or
799/// [`LoadContext::abandon`], and [`RunningLoad::finish`] then collects it.
800///
801/// # Errors
802///
803/// The [`CreationError`] building a latency histogram returned. Nothing has
804/// started then.
805pub fn start(ctx: &Arc<LoadContext>, prepared: PreparedLoad) -> Result<RunningLoad, CreationError> {
806    let contract = &ctx.contract;
807    let mut producer_latencies = Vec::with_capacity(contract.paths.len());
808    for _ in &contract.paths {
809        producer_latencies.push(Latencies::new()?);
810    }
811    let mut consumer_latencies = Vec::with_capacity(prepared.consumers.len());
812    for _ in &prepared.consumers {
813        consumer_latencies.push(Latencies::new()?);
814    }
815    let producers = contract
816        .paths
817        .iter()
818        .zip(producer_latencies)
819        .map(|(path, latencies)| match path.semantics {
820            Semantics::Transactional => {
821                tokio::spawn(run_transactional(Arc::clone(ctx), path.clone(), latencies))
822            }
823            Semantics::Idempotent | Semantics::Compacted => {
824                tokio::spawn(run_idempotent(Arc::clone(ctx), path.clone(), latencies))
825            }
826        })
827        .collect();
828    let mut trackers = Vec::new();
829    let mut consumers = Vec::new();
830    for ((path, isolation, consumer), latencies) in
831        prepared.consumers.into_iter().zip(consumer_latencies)
832    {
833        let tracker = path.commit_offsets.then(|| {
834            Arc::new(CommitTracker::new(
835                group_name(&path.name, isolation),
836                path.name.clone(),
837                path.partitions,
838            ))
839        });
840        if let Some(tracker) = &tracker {
841            trackers.push(Arc::clone(tracker));
842        }
843        consumers.push((
844            path.name.clone(),
845            tokio::spawn(run_consumer(
846                Arc::clone(ctx),
847                path,
848                isolation,
849                tracker,
850                consumer,
851                latencies,
852            )),
853        ));
854    }
855    Ok(RunningLoad {
856        ctx: Arc::clone(ctx),
857        trackers,
858        producers,
859        consumers,
860    })
861}
862
863impl RunningLoad {
864    /// Wait for every task of a load that was stopped or abandoned.
865    ///
866    /// # Errors
867    ///
868    /// When a load task panicked or was cancelled: its ledger is gone, and a
869    /// run without it cannot be judged.
870    pub async fn finish(self) -> Result<LoadOutcome, JoinError> {
871        let mut producers = Vec::new();
872        for task in self.producers {
873            producers.push(task.await?);
874        }
875        let mut consume_latency: BTreeMap<String, Latencies> = BTreeMap::new();
876        for (path, task) in self.consumers {
877            let latencies = task.await?;
878            match consume_latency.entry(path) {
879                Entry::Occupied(mut merged) => merged.get_mut().merge(&latencies),
880                Entry::Vacant(first) => {
881                    first.insert(latencies);
882                }
883            }
884        }
885        Ok(LoadOutcome {
886            producers,
887            consume_latency,
888        })
889    }
890}
891
892impl LoadOutcome {
893    /// Each path's measurement over `steady` time.
894    ///
895    /// # Errors
896    ///
897    /// The [`CreationError`] building the empty latencies of a path with no
898    /// samples returned.
899    pub fn measurements(
900        &self,
901        contract: &Contract,
902        steady: Time,
903    ) -> Result<Vec<PathMeasurement>, CreationError> {
904        let empty = Latencies::new()?;
905        Ok(contract
906            .paths
907            .iter()
908            .map(|path| {
909                let produce = self
910                    .producers
911                    .iter()
912                    .find(|output| output.ledger.path == path.name)
913                    .map_or(&empty, |output| &output.latencies);
914                measurement(
915                    &path.name,
916                    produce,
917                    self.consume_latency.get(&path.name).unwrap_or(&empty),
918                    steady,
919                )
920            })
921            .collect())
922    }
923}
924
925#[cfg(test)]
926mod tests {
927    use {
928        super::{
929            Delivery, Endpoint, Latencies, LoadContext, ProducerError, RecordMetadata,
930            TransactionalRun, Windows, aborted_on_purpose, compacted_key, empty_ledger, key_index,
931            now_ms, pipelined_ticker, record_value, superseded, transactional_ticker,
932        },
933        crate::soak::{
934            contract::{Semantics, fixture},
935            evidence::parse_record_tag,
936            ledger::SeqSet,
937        },
938        assert2::check,
939        futures::FutureExt as _,
940        std::{
941            sync::{Arc, Mutex},
942            time::Duration,
943        },
944        tokio::time::{Instant, Interval, sleep, sleep_until},
945    };
946
947    const PERIOD: Duration = Duration::from_millis(10);
948
949    #[yare::parameterized(
950        acknowledged = { false, false },
951        rejected = { true, false },
952        stalled_admission = { false, true },
953    )]
954    #[test_macro(tokio::test(start_paused = true))]
955    async fn transaction_deliveries_are_observed_before_the_next_paced_send(
956        reject_last: bool,
957        stall_admission: bool,
958    ) {
959        let mut contract = fixture::self_check();
960        contract.clients.transactions.records_per_transaction = 3;
961        let path = contract
962            .paths
963            .iter()
964            .find(|path| path.semantics == Semantics::Transactional)
965            .unwrap()
966            .clone();
967        let ctx = LoadContext::builder()
968            .contract(contract)
969            .endpoint(Endpoint::plaintext("unused:9092"))
970            .windows(Arc::new(Windows::new(Vec::new())))
971            .build();
972        let mut run = TransactionalRun {
973            name: "paced".into(),
974            ledger: empty_ledger(&path, "paced"),
975            latencies: Latencies::new().unwrap(),
976            acked_at: Vec::new(),
977            ticker: transactional_ticker(PERIOD),
978            next_seq: 0,
979            index: 0,
980        };
981        let start = Instant::now();
982        let observed = Arc::new(Mutex::new(Vec::new()));
983        let failed = run
984            .deliveries(&ctx, &path, |seq| {
985                let observed = Arc::clone(&observed);
986                let delay = if stall_admission && seq == 1 {
987                    PERIOD * 3
988                } else {
989                    Duration::ZERO
990                };
991                sleep(delay).map(move |()| {
992                    let ready_at = Instant::now() + PERIOD / 2;
993                    let delivery: Delivery = Box::pin(async move {
994                        sleep_until(ready_at).await;
995                        observed.lock().unwrap().push(start.elapsed());
996                        let result = if reject_last && seq == 2 {
997                            Err(ProducerError::Server(19))
998                        } else {
999                            Ok(RecordMetadata {
1000                                topic: "paced".into(),
1001                                partition: 0,
1002                                offset: i64::try_from(seq).unwrap(),
1003                                timestamp_ms: 0,
1004                                serialized_key_size: -1,
1005                                serialized_value_size: 1,
1006                            })
1007                        };
1008                        (seq, 1, now_ms(), result)
1009                    });
1010                    delivery
1011                })
1012            })
1013            .await;
1014        check!(
1015            *observed.lock().unwrap()
1016                == if stall_admission {
1017                    vec![PERIOD / 2, PERIOD * 4 + PERIOD / 2, PERIOD * 4 + PERIOD / 2]
1018                } else {
1019                    vec![PERIOD / 2, PERIOD + PERIOD / 2, PERIOD * 2 + PERIOD / 2]
1020                }
1021        );
1022        check!(failed == reject_last);
1023        check!(run.ledger.acked == SeqSet::range(0, if reject_last { 2 } else { 3 }));
1024        check!(
1025            run.ledger.failed
1026                == if reject_last {
1027                    SeqSet::range(2, 3)
1028                } else {
1029                    SeqSet::new()
1030                }
1031        );
1032        check!(run.latencies.steady_count == if reject_last { 2 } else { 3 });
1033    }
1034
1035    /// How long ten ticks take when the ticker stalls for five periods after
1036    /// the first, as a transactional producer does while a transaction ends.
1037    async fn ten_ticks_across_a_stall(mut ticker: Interval) -> Duration {
1038        let start = Instant::now();
1039        ticker.tick().await;
1040        sleep(PERIOD * 5).await;
1041        for _ in 1..10 {
1042            ticker.tick().await;
1043        }
1044        start.elapsed()
1045    }
1046
1047    #[yare::parameterized(
1048        transactional_keeps_its_average_rate = { transactional_ticker, PERIOD * 9 },
1049        pipelined_drops_what_it_missed = { pipelined_ticker, PERIOD * 13 },
1050    )]
1051    #[test_macro(tokio::test(start_paused = true))]
1052    async fn only_a_pipelined_producer_drops_the_ticks_a_stall_missed(
1053        ticker: fn(Duration) -> Interval,
1054        took: Duration,
1055    ) {
1056        check!(ten_ticks_across_a_stall(ticker(PERIOD)).await == took);
1057    }
1058
1059    #[yare::parameterized(
1060        padded = { 64, 64 },
1061        tag_longer_than_payload = { 4, 13 },
1062    )]
1063    fn a_value_carries_its_tag_padded_to_the_payload(payload: usize, length: usize) {
1064        let value = record_value("path-0", 12_345, payload);
1065        check!(value.len() == length);
1066        check!(parse_record_tag(&value) == Some(("path-0", 12_345)));
1067    }
1068
1069    #[test]
1070    fn a_compacted_key_names_its_index() {
1071        let key = compacted_key("classic-compacted-0", 130, 64);
1072        check!(key == "classic-compacted-0/k2");
1073        check!(key_index(key.as_bytes()) == Some(2));
1074        check!(key_index(b"no-index") == None);
1075    }
1076
1077    #[yare::parameterized(
1078        nothing_later_sent = { 10, 64, 70, false },
1079        next_value_of_the_key_not_yet_sent = { 10, 64, 74, false },
1080        next_value_of_the_key_sent = { 10, 64, 75, true },
1081    )]
1082    fn a_failed_compacted_send_is_superseded_once_its_key_moved_on(
1083        seq: u64,
1084        keys: u64,
1085        next: u64,
1086        expected: bool,
1087    ) {
1088        check!(superseded(seq, keys, next) == expected);
1089    }
1090
1091    #[yare::parameterized(
1092        first = { 0, false },
1093        third = { 2, false },
1094        fourth = { 3, true },
1095        eighth = { 7, true },
1096    )]
1097    fn every_fourth_transaction_is_aborted(index: u64, aborted: bool) {
1098        check!(aborted_on_purpose(index, 4) == aborted);
1099    }
1100}
krabka_bench_driver/soak/client/load.rs

krabka_bench_driver/soak/client/
load.rs
