1use {
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
33const READBACK_LIMIT: Duration = Duration::from_mins(15);
36
37const READBACK_STALL: Duration = Duration::from_mins(2);
40
41const LOOKUP_ATTEMPTS: u32 = 10;
43
44#[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#[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#[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
67pub 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
86async 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
113async 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
130pub 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
158async 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
220pub 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
296async fn committed(
299 endpoint: &Endpoint,
300 group: &str,
301 partitions: &[(String, i32)],
302) -> Result<BTreeMap<i32, i64>, String> {
303 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
326pub 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
378async 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
389pub 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}