Skip to main content

krabka_bench_driver/soak/client/
readback.rs

1//! After the load: the committed offsets, and the read-back that the
2//! reconciler compares with the producers' ledgers.
3//!
4//! The read-back is a fresh consumer with no group, assigned every partition
5//! of a path and started from the earliest offset, that reads until it has
6//! reached each partition's end. It runs once more only when the first pass
7//! found a record twice, to find every position of those records; holding a
8//! position for every record of a 72-hour run would cost more memory than
9//! the ledger it checks.
10
11use {
12    super::{
13        Endpoint,
14        load::{CommitTracker, isolation_level, key_index},
15        now_ms,
16    },
17    crate::soak::{
18        contract::{Isolation, LoadPath, Semantics},
19        evidence::{CommitCheck, DuplicateRead, ReaderLedger, parse_record_tag},
20        ledger::SeqSet,
21    },
22    bytes::Bytes,
23    krabka_client_consumer::{Consumer, ConsumerError, ConsumerRecord},
24    krabka_client_producer::{Acks, Producer, ProducerError, ProducerRecord, RecordMetadata},
25    krabka_units::prelude::*,
26    std::{
27        collections::{BTreeMap, HashMap},
28        time::{Duration, Instant},
29    },
30    tokio::time::sleep,
31};
32
33/// How long one read-back pass may take before it gives up and reports what
34/// it read. A shortfall is then a finding of the reconciler, not a hang.
35const READBACK_LIMIT: Duration = Duration::from_mins(15);
36
37/// How long a read-back pass may go without reading a record before it gives
38/// up. A cold read from the remote tier is slow, but never this slow.
39const READBACK_STALL: Duration = Duration::from_mins(2);
40
41/// How many times a committed-offset or end-offset lookup is tried.
42const LOOKUP_ATTEMPTS: u32 = 10;
43
44/// Every partition of `path`.
45#[must_use]
46pub fn partitions(path: &LoadPath) -> Vec<(String, i32)> {
47    (0..path.partitions)
48        .map(|partition| (path.name.clone(), partition))
49        .collect()
50}
51
52/// Whether every partition of `ends` has a next offset at or past its end.
53/// A partition that ends at offset zero holds nothing and needs no read.
54#[must_use]
55pub fn caught_up(next: &BTreeMap<i32, i64>, ends: &BTreeMap<i32, i64>) -> bool {
56    ends.iter()
57        .all(|(partition, end)| *end <= 0 || next.get(partition).is_some_and(|next| next >= end))
58}
59
60/// Whether no transaction is open on any partition: the last stable offset
61/// is the high watermark everywhere.
62#[must_use]
63pub fn stable(committed_ends: &BTreeMap<i32, i64>, uncommitted_ends: &BTreeMap<i32, i64>) -> bool {
64    !uncommitted_ends.is_empty() && committed_ends == uncommitted_ends
65}
66
67/// A consumer with no group, for reading and for offset lookups.
68///
69/// # Errors
70///
71/// Returns the client's error when the consumer cannot be built.
72pub async fn groupless(
73    endpoint: &Endpoint,
74    isolation: Isolation,
75) -> Result<Consumer, ConsumerError> {
76    Consumer::builder()
77        .bootstrap(endpoint.bootstrap.clone())
78        .maybe_security(endpoint.security.clone())
79        .client_id(format!("soak-readback-{isolation}"))
80        .isolation_level(isolation_level(isolation))
81        .default_api_timeout(secs(15))
82        .build()
83        .await
84}
85
86/// Build a consumer, ask it one question, and close it, retrying the whole
87/// round up to [`LOOKUP_ATTEMPTS`] times a second apart. A cluster that is
88/// still recovering from the last fault answers the first attempts with a
89/// connection error or a coordinator move.
90async fn ask_consumer<T, B, Q>(build: B, query: Q) -> Result<T, String>
91where
92    B: AsyncFn() -> Result<Consumer, ConsumerError>,
93    Q: AsyncFn(&Consumer) -> Result<T, ConsumerError>,
94{
95    let mut last = String::new();
96    for _ in 0..LOOKUP_ATTEMPTS {
97        match build().await {
98            Ok(consumer) => {
99                let answer = query(&consumer).await;
100                let _ = consumer.close().await;
101                match answer {
102                    Ok(answer) => return Ok(answer),
103                    Err(error) => last = error.to_string(),
104                }
105            }
106            Err(error) => last = error.to_string(),
107        }
108        sleep(Duration::from_secs(1)).await;
109    }
110    Err(last)
111}
112
113/// The end offset of every partition at `isolation`.
114async fn end_offsets(
115    endpoint: &Endpoint,
116    isolation: Isolation,
117    partitions: &[(String, i32)],
118) -> Result<BTreeMap<i32, i64>, String> {
119    let ends = ask_consumer(
120        async || groupless(endpoint, isolation).await,
121        async |consumer: &Consumer| consumer.end_offsets(partitions).await,
122    )
123    .await?;
124    Ok(ends
125        .into_iter()
126        .map(|((_, partition), end)| (partition, end))
127        .collect())
128}
129
130/// Wait until no transaction is open on `path`, so a `read_committed`
131/// read-back is not cut short by a transaction that its producer left open
132/// and the coordinator has yet to time out.
133pub async fn wait_stable(endpoint: &Endpoint, path: &LoadPath, limit: Time) {
134    if path.semantics != Semantics::Transactional {
135        return;
136    }
137    let partitions = partitions(path);
138    let deadline = Instant::now() + limit.to_std();
139    loop {
140        let committed = end_offsets(endpoint, Isolation::ReadCommitted, &partitions).await;
141        let uncommitted = end_offsets(endpoint, Isolation::ReadUncommitted, &partitions).await;
142        if let (Ok(committed), Ok(uncommitted)) = (&committed, &uncommitted)
143            && stable(committed, uncommitted)
144        {
145            return;
146        }
147        if Instant::now() >= deadline {
148            eprintln!(
149                "soak-contract: {}: transactions still open after {limit:?}: last stable {committed:?}, high watermark {uncommitted:?}",
150                path.name
151            );
152            return;
153        }
154        sleep(Duration::from_secs(1)).await;
155    }
156}
157
158/// Read every partition of `path` at `isolation` from the earliest offset to
159/// its end, handing each record to `visit`.
160async fn scan(
161    endpoint: &Endpoint,
162    path: &LoadPath,
163    isolation: Isolation,
164    mut visit: impl FnMut(&ConsumerRecord),
165) -> Result<(), String> {
166    let partitions = partitions(path);
167    let ends = end_offsets(endpoint, isolation, &partitions).await?;
168    let consumer = groupless(endpoint, isolation)
169        .await
170        .map_err(|error| error.to_string())?;
171    let mut consumer = consumer;
172    consumer
173        .assign(&partitions)
174        .await
175        .map_err(|error| error.to_string())?;
176    consumer
177        .seek_to_beginning(&partitions)
178        .await
179        .map_err(|error| error.to_string())?;
180    let mut next: BTreeMap<i32, i64> = BTreeMap::new();
181    let started = Instant::now();
182    let mut progressed = started;
183    let outcome = loop {
184        if caught_up(&next, &ends) || consumer.at_log_end().await {
185            break Ok(());
186        }
187        if started.elapsed() >= READBACK_LIMIT || progressed.elapsed() >= READBACK_STALL {
188            break Err(format!(
189                "{} at {isolation}: not at the end after {:?}, the last {:?} of it reading nothing: \
190                 read to {next:?} of {ends:?}",
191                path.name,
192                started.elapsed(),
193                progressed.elapsed(),
194            ));
195        }
196        match consumer.poll(secs(1)).await {
197            Ok(records) => {
198                if !records.is_empty() {
199                    progressed = Instant::now();
200                }
201                for record in &records {
202                    visit(record);
203                    let reached = next.entry(record.partition).or_insert(0);
204                    *reached = (*reached).max(record.offset + 1);
205                }
206            }
207            Err(error) => {
208                eprintln!(
209                    "soak-contract: read-back of {} at {isolation}: {error}",
210                    path.name
211                );
212                sleep(Duration::from_millis(500)).await;
213            }
214        }
215    };
216    let _ = consumer.close().await;
217    outcome
218}
219
220/// What one reader saw of `producer`'s records on `path`.
221pub async fn read_back(
222    endpoint: &Endpoint,
223    path: &LoadPath,
224    isolation: Isolation,
225    producer: &str,
226) -> ReaderLedger {
227    let mut seen = SeqSet::new();
228    let mut repeated = SeqSet::new();
229    let mut latest_by_key: BTreeMap<u64, u64> = BTreeMap::new();
230    let mut foreign = 0_u64;
231    let first = scan(endpoint, path, isolation, |record| {
232        let Some((owner, seq)) = record.value.as_deref().and_then(parse_record_tag) else {
233            foreign += 1;
234            return;
235        };
236        if owner != producer {
237            foreign += 1;
238            return;
239        }
240        if !seen.insert(seq) {
241            repeated.insert(seq);
242        }
243        if let (Some(keys), Some(key)) = (path.key_space, record.key.as_deref()) {
244            let index = key_index(key).unwrap_or_else(|| seq % keys.max(1));
245            let latest = latest_by_key.entry(index).or_insert(seq);
246            *latest = (*latest).max(seq);
247        }
248    })
249    .await;
250    if let Err(error) = first {
251        eprintln!("soak-contract: {error}");
252    }
253    if foreign > 0 {
254        eprintln!(
255            "soak-contract: {} at {isolation}: {foreign} records no soak producer of this path wrote",
256            path.name
257        );
258    }
259    let duplicates = if repeated.is_empty() {
260        Vec::new()
261    } else {
262        let mut positions: HashMap<u64, Vec<(i32, i64)>> = HashMap::new();
263        let second = scan(endpoint, path, isolation, |record| {
264            if let Some((owner, seq)) = record.value.as_deref().and_then(parse_record_tag)
265                && owner == producer
266                && repeated.contains(seq)
267            {
268                positions
269                    .entry(seq)
270                    .or_default()
271                    .push((record.partition, record.offset));
272            }
273        })
274        .await;
275        if let Err(error) = second {
276            eprintln!("soak-contract: {error}");
277        }
278        repeated
279            .iter()
280            .map(|seq| DuplicateRead {
281                seq,
282                positions: positions.remove(&seq).unwrap_or_default(),
283            })
284            .collect()
285    };
286    ReaderLedger {
287        path: path.name.clone(),
288        isolation,
289        producer: producer.to_owned(),
290        seen,
291        duplicates,
292        latest_by_key,
293    }
294}
295
296/// The committed offsets of `group` on `partitions`, as `OffsetFetch` returns
297/// them. A partition with nothing committed is absent.
298async fn committed(
299    endpoint: &Endpoint,
300    group: &str,
301    partitions: &[(String, i32)],
302) -> Result<BTreeMap<i32, i64>, String> {
303    // A consumer with a group id that never subscribes never joins, so it
304    // reads the group's offsets without becoming a member of it.
305    let build = async || {
306        Consumer::builder()
307            .bootstrap(endpoint.bootstrap.clone())
308            .maybe_security(endpoint.security.clone())
309            .client_id(format!("{group}-checker"))
310            .group_id(group.to_owned())
311            .enable_auto_commit(false)
312            .default_api_timeout(secs(15))
313            .build()
314            .await
315    };
316    let fetched = ask_consumer(build, async |consumer: &Consumer| {
317        consumer.committed(partitions).await
318    })
319    .await?;
320    Ok(fetched
321        .into_iter()
322        .filter_map(|((_, partition), offset)| offset.map(|offset| (partition, offset.offset)))
323        .collect())
324}
325
326/// One [`CommitCheck`] per partition of `acknowledged`: the last commit the
327/// coordinator acknowledged before `fault`, against what `OffsetFetch` returns
328/// now and the partition's log end.
329///
330/// A group that had acknowledged no commit yet is checked on every one of its
331/// `partitions` from zero, so each fault is checked for each group and a group
332/// that never commits is a finding rather than an absence.
333pub async fn commit_checks(
334    endpoint: &Endpoint,
335    tracker: &CommitTracker,
336    fault: &str,
337    acknowledged: &BTreeMap<i32, i64>,
338) -> Vec<CommitCheck> {
339    let (group, topic) = (tracker.group.as_str(), tracker.topic.as_str());
340    let acknowledged = if acknowledged.is_empty() {
341        (0..tracker.partitions)
342            .map(|partition| (partition, 0))
343            .collect()
344    } else {
345        acknowledged.clone()
346    };
347    let partitions: Vec<(String, i32)> = acknowledged
348        .keys()
349        .map(|partition| (topic.to_owned(), *partition))
350        .collect();
351    let fetched = committed(endpoint, group, &partitions).await;
352    if let Err(error) = &fetched {
353        eprintln!("soak-contract: OffsetFetch of {group} after {fault}: {error}");
354    }
355    let ends = end_offsets(endpoint, Isolation::ReadUncommitted, &partitions)
356        .await
357        .unwrap_or_else(|error| {
358            eprintln!("soak-contract: end offsets of {topic} after {fault}: {error}");
359            BTreeMap::new()
360        });
361    acknowledged
362        .iter()
363        .map(|(partition, acked)| CommitCheck {
364            group: group.to_owned(),
365            topic: topic.to_owned(),
366            partition: *partition,
367            fault: fault.to_owned(),
368            acknowledged: *acked,
369            fetched: fetched
370                .as_ref()
371                .ok()
372                .and_then(|fetched| fetched.get(partition).copied()),
373            log_end: ends.get(partition).copied().unwrap_or(-1),
374        })
375        .collect()
376}
377
378/// Send `record` with `producer`, then close the producer whatever the send
379/// returned.
380async fn send_once(
381    producer: Producer,
382    record: ProducerRecord,
383) -> Result<RecordMetadata, ProducerError> {
384    let sent = producer.send(record).await;
385    let _ = producer.close().await;
386    sent
387}
388
389/// Produce one record to `topic`-`partition` with `acks=all`, then read it
390/// back at the offset it was acknowledged at, and return that offset.
391///
392/// This is what "serving" means after a restore: a client of the restored
393/// cluster, authenticated as an ordinary principal, can write and then read.
394///
395/// # Errors
396///
397/// Returns why the write or the read failed, or that the record did not come
398/// back within a minute.
399pub async fn round_trip(endpoint: &Endpoint, topic: &str, partition: i32) -> Result<i64, String> {
400    let value = format!("soak-drill-probe/{};", now_ms());
401    let sent = send_once(
402        Producer::builder()
403            .bootstrap(endpoint.bootstrap.clone())
404            .maybe_security(endpoint.security.clone())
405            .client_id("soak-drill-probe".to_owned())
406            .acks(Acks::All)
407            .enable_idempotence(true)
408            .request_timeout(Duration::from_secs(10))
409            .delivery_timeout(Duration::from_secs(30))
410            .build()
411            .await
412            .map_err(|error| format!("build a producer: {error}"))?,
413        ProducerRecord {
414            topic: topic.to_owned(),
415            partition: Some(partition),
416            value: Some(Bytes::from(value.clone())),
417            ..Default::default()
418        },
419    )
420    .await
421    .map_err(|error| format!("produce to {topic}-{partition}: {error}"));
422    let offset = sent?.offset;
423    let assignment = [(topic.to_owned(), partition)];
424    let mut consumer = groupless(endpoint, Isolation::ReadUncommitted)
425        .await
426        .map_err(|error| format!("build a consumer: {error}"))?;
427    let read = async {
428        consumer
429            .assign(&assignment)
430            .await
431            .map_err(|error| error.to_string())?;
432        consumer
433            .seek(topic, partition, offset)
434            .await
435            .map_err(|error| error.to_string())?;
436        let deadline = Instant::now() + Duration::from_mins(1);
437        while Instant::now() < deadline {
438            match consumer.poll(secs(1)).await {
439                Ok(records) => {
440                    if let Some(record) = records.iter().find(|record| record.offset == offset) {
441                        return if record.value.as_deref() == Some(value.as_bytes()) {
442                            Ok(offset)
443                        } else {
444                            Err(format!(
445                                "{topic}-{partition} offset {offset} holds another record"
446                            ))
447                        };
448                    }
449                }
450                Err(error) => eprintln!("soak-drill: probe read of {topic}-{partition}: {error}"),
451            }
452        }
453        Err(format!(
454            "{topic}-{partition} offset {offset} did not come back within a minute"
455        ))
456    }
457    .await;
458    let _ = consumer.close().await;
459    read
460}
461
462#[cfg(test)]
463mod tests {
464    use {
465        super::{caught_up, stable},
466        assert2::check,
467        std::collections::BTreeMap,
468    };
469
470    fn offsets(pairs: &[(i32, i64)]) -> BTreeMap<i32, i64> { pairs.iter().copied().collect() }
471
472    #[yare::parameterized(
473        every_partition_at_its_end = { &[(0, 10), (1, 4)], &[(0, 10), (1, 4)], true },
474        one_short = { &[(0, 10), (1, 3)], &[(0, 10), (1, 4)], false },
475        one_unread = { &[(0, 10)], &[(0, 10), (1, 4)], false },
476        empty_partition_needs_no_read = { &[(0, 10)], &[(0, 10), (1, 0)], true },
477        past_the_end = { &[(0, 12)], &[(0, 10)], true },
478    )]
479    fn the_read_back_ends_when_every_partition_reached_its_end(
480        next: &[(i32, i64)],
481        ends: &[(i32, i64)],
482        expected: bool,
483    ) {
484        check!(caught_up(&offsets(next), &offsets(ends)) == expected);
485    }
486
487    #[yare::parameterized(
488        settled = { &[(0, 10), (1, 4)], &[(0, 10), (1, 4)], true },
489        open_transaction = { &[(0, 8), (1, 4)], &[(0, 10), (1, 4)], false },
490        nothing_known = { &[], &[], false },
491    )]
492    fn a_topic_is_stable_when_the_last_stable_offset_is_the_high_watermark(
493        committed: &[(i32, i64)],
494        uncommitted: &[(i32, i64)],
495        expected: bool,
496    ) {
497        check!(stable(&offsets(committed), &offsets(uncommitted)) == expected);
498    }
499}