1use std::collections::HashMap;
13use std::sync::{Arc, Mutex, Weak};
14use std::time::Duration;
15
16use futures::future::select_all;
17use rdkafka::consumer::{Consumer as _, ConsumerGroupMetadata, StreamConsumer};
18use rdkafka::{Offset, TopicPartitionList};
19use ruststream::codec::Codec;
20#[cfg(any(feature = "json", feature = "cbor", feature = "msgpack"))]
21use ruststream::codec::DefaultCodec;
22use ruststream::runtime::{
23 Outgoing, PublishContext, PublishTransform, PublishTransformIdentity, PublishTransformStack,
24 TypedPublisher,
25};
26use ruststream::{OutgoingMessage, Publisher, TransactionalPublisher as _};
27use tracing::{debug, error};
28
29use crate::error::KafkaError;
30use crate::publisher::KafkaPublisher;
31use crate::tracker::{CommitTracker, TrackingContext};
32
33const DEFAULT_COMMIT_INTERVAL: Duration = Duration::from_millis(100);
35
36#[derive(Clone)]
43pub(crate) struct EosSource {
44 tracker: Weak<CommitTracker>,
45 consumer: Weak<StreamConsumer<TrackingContext>>,
46}
47
48impl EosSource {
49 pub(crate) fn new(
50 tracker: &Arc<CommitTracker>,
51 consumer: &Arc<StreamConsumer<TrackingContext>>,
52 ) -> Self {
53 Self {
54 tracker: Arc::downgrade(tracker),
55 consumer: Arc::downgrade(consumer),
56 }
57 }
58
59 pub(crate) fn alive(&self) -> bool {
60 self.tracker.strong_count() > 0 && self.consumer.strong_count() > 0
61 }
62
63 fn upgrade(&self) -> Option<LiveSource> {
64 Some(LiveSource {
65 tracker: self.tracker.upgrade()?,
66 consumer: self.consumer.upgrade()?,
67 })
68 }
69}
70
71struct LiveSource {
73 tracker: Arc<CommitTracker>,
74 consumer: Arc<StreamConsumer<TrackingContext>>,
75}
76
77#[derive(Debug, Clone, PartialEq, Eq)]
82pub struct SourceOffset {
83 topic: String,
84 partition: i32,
85 offset: i64,
86}
87
88impl SourceOffset {
89 #[must_use]
92 pub fn new(topic: impl Into<String>, partition: i32, offset: i64) -> Self {
93 Self {
94 topic: topic.into(),
95 partition,
96 offset,
97 }
98 }
99
100 fn key(&self) -> (String, i32) {
101 (self.topic.clone(), self.partition)
102 }
103}
104
105#[derive(Debug, Clone, Copy, PartialEq, Eq)]
107enum Phase {
108 Idle,
110 Opening,
112 Open,
114 Committing,
117}
118
119#[derive(Debug)]
121struct Window {
122 phase: Phase,
123 enrolled: HashMap<(String, i32), i64>,
125 failed: bool,
127 epoch: u64,
130}
131
132struct PipelineInner {
133 publisher: KafkaPublisher,
134 id: Option<String>,
137 interval: Duration,
138 window: Mutex<Window>,
139 phase_changed: tokio::sync::Notify,
141 committed: Mutex<HashMap<(String, i32), i64>>,
144 session_low: Mutex<HashMap<(String, i32), i64>>,
147}
148
149#[derive(Clone)]
172pub struct EosPipeline {
173 inner: Arc<PipelineInner>,
174}
175
176impl std::fmt::Debug for EosPipeline {
177 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
178 f.debug_struct("EosPipeline")
179 .field("id", &self.inner.id)
180 .field("interval", &self.inner.interval)
181 .finish_non_exhaustive()
182 }
183}
184
185impl EosPipeline {
186 #[must_use]
191 pub fn new(publisher: KafkaPublisher) -> Self {
192 let id = publisher.transactional_id_str().map(str::to_owned);
193 Self {
194 inner: Arc::new(PipelineInner {
195 publisher,
196 id,
197 interval: DEFAULT_COMMIT_INTERVAL,
198 window: Mutex::new(Window {
199 phase: Phase::Idle,
200 enrolled: HashMap::new(),
201 failed: false,
202 epoch: 0,
203 }),
204 phase_changed: tokio::sync::Notify::new(),
205 committed: Mutex::new(HashMap::new()),
206 session_low: Mutex::new(HashMap::new()),
207 }),
208 }
209 }
210
211 #[must_use]
216 pub fn commit_interval(self, interval: Duration) -> Self {
217 Self {
218 inner: Arc::new(PipelineInner {
219 publisher: self.inner.publisher.clone(),
220 id: self.inner.id.clone(),
221 interval,
222 window: Mutex::new(Window {
223 phase: Phase::Idle,
224 enrolled: HashMap::new(),
225 failed: false,
226 epoch: 0,
227 }),
228 phase_changed: tokio::sync::Notify::new(),
229 committed: Mutex::new(HashMap::new()),
230 session_low: Mutex::new(HashMap::new()),
231 }),
232 }
233 }
234
235 pub async fn publish(
257 &self,
258 source: &SourceOffset,
259 msg: OutgoingMessage<'_>,
260 ) -> Result<(), KafkaError> {
261 self.admit(source).await?;
262 let sent = self.inner.publisher.publish(msg).await;
263 if sent.is_err() {
264 let mut window = self.inner.window.lock().expect("window mutex poisoned");
265 window.failed = true;
267 }
268 sent
269 }
270
271 async fn admit(&self, source: &SourceOffset) -> Result<(), KafkaError> {
274 loop {
275 let phase_changed = self.inner.phase_changed.notified();
278 let action = {
279 let mut window = self.inner.window.lock().expect("window mutex poisoned");
280 match window.phase {
281 Phase::Open => {
282 Self::enroll(&mut window, &self.inner.session_low, source);
283 Admission::Admitted
284 }
285 Phase::Idle => {
286 window.phase = Phase::Opening;
287 Admission::Opener
288 }
289 Phase::Opening => Admission::Wait,
290 Phase::Committing => {
291 let participant = window
295 .enrolled
296 .get(&source.key())
297 .is_some_and(|max| source.offset <= *max);
298 if participant {
299 Admission::Admitted
300 } else {
301 Admission::Wait
302 }
303 }
304 }
305 };
306 match action {
307 Admission::Admitted => return Ok(()),
308 Admission::Wait => {
309 phase_changed.await;
310 }
311 Admission::Opener => return self.open_window(source).await,
312 }
313 }
314 }
315
316 async fn open_window(&self, source: &SourceOffset) -> Result<(), KafkaError> {
318 let begun = self.inner.publisher.begin_transaction().await;
319 let mut window = self.inner.window.lock().expect("window mutex poisoned");
320 match begun {
321 Ok(()) => {
322 window.phase = Phase::Open;
323 window.failed = false;
324 Self::enroll(&mut window, &self.inner.session_low, source);
325 let epoch = window.epoch;
326 drop(window);
327 tokio::spawn(run_window(Arc::clone(&self.inner), epoch));
328 }
329 Err(err) => {
330 window.phase = Phase::Idle;
331 drop(window);
332 self.inner.phase_changed.notify_waiters();
333 return Err(err);
334 }
335 }
336 self.inner.phase_changed.notify_waiters();
337 Ok(())
338 }
339
340 fn enroll(
341 window: &mut Window,
342 session_low: &Mutex<HashMap<(String, i32), i64>>,
343 source: &SourceOffset,
344 ) {
345 let key = source.key();
346 session_low
347 .lock()
348 .expect("session low mutex poisoned")
349 .entry(key.clone())
350 .or_insert(source.offset);
351 let max = window.enrolled.entry(key).or_insert(source.offset);
352 if source.offset > *max {
353 *max = source.offset;
354 }
355 }
356}
357
358enum Admission {
359 Admitted,
360 Wait,
361 Opener,
362}
363
364async fn run_window(inner: Arc<PipelineInner>, epoch: u64) {
367 tokio::time::sleep(inner.interval).await;
368 let enrolled = {
369 let mut window = inner.window.lock().expect("window mutex poisoned");
370 if window.epoch != epoch || window.phase != Phase::Open {
371 return;
372 }
373 window.phase = Phase::Committing;
374 window.enrolled.clone()
375 };
376 let outcome = commit_window(&inner, &enrolled).await;
377 {
378 let mut window = inner.window.lock().expect("window mutex poisoned");
379 window.phase = Phase::Idle;
380 window.enrolled.clear();
381 window.failed = false;
382 window.epoch += 1;
383 }
384 inner.phase_changed.notify_waiters();
385 if let Err(err) = outcome {
386 error!(
387 target: "ruststream_rdkafka",
388 pipeline = inner.id.as_deref().unwrap_or("<no id>"),
389 error = %err,
390 "EOS window aborted; its sources seek back and the window redelivers",
391 );
392 }
393}
394
395async fn commit_window(
398 inner: &Arc<PipelineInner>,
399 enrolled: &HashMap<(String, i32), i64>,
400) -> Result<(), KafkaError> {
401 let id = inner.id.clone().ok_or_else(|| {
402 KafkaError::InvalidOptions(
403 "an EosPipeline publisher needs `KafkaPublisher::transactional_id`".to_owned(),
404 )
405 })?;
406 let conn = inner.publisher.shared_conn();
407 let state = conn.get().ok_or(KafkaError::NotConnected)?;
408 let sources: Vec<LiveSource> = state
409 .eos_sources(&id)
410 .iter()
411 .filter_map(EosSource::upgrade)
412 .collect();
413
414 let failed = {
415 let window = inner.window.lock().expect("window mutex poisoned");
416 window.failed
417 };
418 let ready = if failed {
419 Err(KafkaError::Publish(
420 "a publish into this window failed; the transaction is poisoned"
421 .to_owned()
422 .into(),
423 ))
424 } else {
425 wait_settled(inner, &sources, enrolled).await
426 };
427 let result = match ready {
428 Ok(()) => try_commit(inner, &sources).await,
429 Err(err) => Err(err),
430 };
431 if let Err(err) = result {
432 abort_window(inner, &sources, enrolled).await;
433 return Err(err);
434 }
435 Ok(())
436}
437
438async fn wait_settled(
441 inner: &Arc<PipelineInner>,
442 sources: &[LiveSource],
443 enrolled: &HashMap<(String, i32), i64>,
444) -> Result<(), KafkaError> {
445 let deadline = tokio::time::Instant::now() + inner.publisher.transaction_deadline();
446 loop {
447 let waiters: Vec<_> = sources
450 .iter()
451 .map(|source| Box::pin(source.tracker.advance_waiter()))
452 .collect();
453 let pending = enrolled.iter().find(|((topic, partition), max)| {
454 !sources.iter().any(|source| {
455 source
456 .tracker
457 .stored_position(topic, *partition)
458 .is_some_and(|stored| stored >= **max)
459 })
460 });
461 let Some(((topic, partition), max)) = pending else {
462 return Ok(());
463 };
464 if waiters.is_empty() {
465 return Err(KafkaError::InvalidOptions(format!(
466 "EOS pipeline has no registered sources for its id; is the subscription in \
467 `Commit::Transactional` mode with the matching pipeline id? (waiting on \
468 {topic}[{partition}] up to offset {max})",
469 )));
470 }
471 debug!(
472 target: "ruststream_rdkafka",
473 topic = %topic,
474 partition = partition,
475 up_to = max,
476 "EOS window waiting for participants to settle",
477 );
478 if tokio::time::timeout_at(deadline, select_all(waiters))
479 .await
480 .is_err()
481 {
482 return Err(KafkaError::Publish(
483 format!(
484 "EOS window stalled: {topic}[{partition}] did not settle up to offset \
485 {max} within the transaction deadline (a hung or retrying handler, or a \
486 revoked partition)",
487 )
488 .into(),
489 ));
490 }
491 }
492}
493
494async fn try_commit(inner: &Arc<PipelineInner>, sources: &[LiveSource]) -> Result<(), KafkaError> {
497 let mut sent: Vec<((String, i32), i64)> = Vec::new();
498 for source in sources {
499 let positions = source.tracker.stored_positions();
500 if positions.is_empty() {
501 continue;
502 }
503 let mut offsets = TopicPartitionList::new();
504 for ((topic, partition), stored) in &positions {
505 offsets
506 .add_partition_offset(topic, *partition, Offset::Offset(stored + 1))
507 .map_err(KafkaError::publish)?;
508 }
509 let metadata = group_metadata(source)?;
510 inner.publisher.send_offsets(offsets, metadata).await?;
511 sent.extend(positions.into_iter().map(|(key, stored)| (key, stored + 1)));
512 }
513 inner.publisher.commit().await?;
514 {
515 let mut committed = inner.committed.lock().expect("committed mutex poisoned");
516 for (key, next) in sent {
517 committed.insert(key, next);
518 }
519 }
520 Ok(())
521}
522
523fn group_metadata(source: &LiveSource) -> Result<ConsumerGroupMetadata, KafkaError> {
524 source.consumer.group_metadata().ok_or_else(|| {
525 KafkaError::Publish(
526 "the source consumer has no group metadata (not a group member yet or already \
527 closed); cannot commit its offsets transactionally"
528 .to_owned()
529 .into(),
530 )
531 })
532}
533
534async fn abort_window(
538 inner: &Arc<PipelineInner>,
539 sources: &[LiveSource],
540 enrolled: &HashMap<(String, i32), i64>,
541) {
542 if let Err(err) = inner.publisher.abort().await {
543 error!(
544 target: "ruststream_rdkafka",
545 error = %err,
546 "EOS window abort failed; the transaction resolves by its broker-side timeout",
547 );
548 }
549 let committed = inner
550 .committed
551 .lock()
552 .expect("committed mutex poisoned")
553 .clone();
554 let session_low = inner
555 .session_low
556 .lock()
557 .expect("session low mutex poisoned")
558 .clone();
559 for key @ (topic, partition) in enrolled.keys() {
560 let Some(target) = committed
561 .get(key)
562 .copied()
563 .or_else(|| session_low.get(key).copied())
564 else {
565 continue;
566 };
567 let Some(source) = sources
568 .iter()
569 .find(|source| source.tracker.covers(topic, *partition))
570 else {
571 continue;
572 };
573 if let Err(err) = source.consumer.seek(
574 topic,
575 *partition,
576 Offset::Offset(target),
577 Duration::from_secs(5),
578 ) {
579 debug!(
582 target: "ruststream_rdkafka",
583 topic = %topic,
584 partition = partition,
585 error = %err,
586 "seek-back after an aborted EOS window failed",
587 );
588 }
589 }
590}
591
592pub const EOS_SOURCE_HEADER: &str = "kafka-eos-source";
599
600pub(crate) fn encode_source(topic: &str, partition: i32, offset: i64) -> String {
601 format!("{partition}:{offset}:{topic}")
602}
603
604fn decode_source(value: &str) -> Option<SourceOffset> {
605 let mut parts = value.splitn(3, ':');
606 let partition = parts.next()?.parse().ok()?;
607 let offset = parts.next()?.parse().ok()?;
608 let topic = parts.next()?;
609 Some(SourceOffset::new(topic, partition, offset))
610}
611
612#[derive(Debug, Clone, Copy, Default)]
620pub struct EosReplies;
621
622impl<C> PublishTransform<C> for EosReplies {
623 fn apply(&self, out: &mut Outgoing<'_>, cx: &PublishContext<'_, C>) {
624 if let Some(source) = cx.headers().get(EOS_SOURCE_HEADER) {
625 let source = source.to_vec();
626 out.headers_mut().insert(EOS_SOURCE_HEADER, source);
627 }
628 }
629}
630
631impl EosPipeline {
632 #[cfg(any(feature = "json", feature = "cbor", feature = "msgpack"))]
655 #[must_use]
656 pub fn replies(
657 &self,
658 ) -> TypedPublisher<
659 Self,
660 DefaultCodec,
661 PublishTransformStack<PublishTransformIdentity, EosReplies>,
662 > {
663 TypedPublisher::new(self.clone()).transform(EosReplies)
664 }
665
666 #[must_use]
668 pub fn replies_with<C: Codec>(
669 &self,
670 codec: C,
671 ) -> TypedPublisher<Self, C, PublishTransformStack<PublishTransformIdentity, EosReplies>> {
672 TypedPublisher::with_codec(self.clone(), codec).transform(EosReplies)
673 }
674}
675
676impl Publisher for EosPipeline {
677 type Error = KafkaError;
678
679 async fn publish(&self, msg: OutgoingMessage<'_>) -> Result<(), Self::Error> {
694 let Some(source) = msg
695 .headers()
696 .get_str(EOS_SOURCE_HEADER)
697 .and_then(decode_source)
698 else {
699 return Err(KafkaError::InvalidOptions(
700 "an EOS reply carries no source coordinates: the subscription must be in \
701 `Commit::Transactional` mode for this pipeline, and the reply publisher must \
702 relay them (wire it with `EosPipeline::replies()` or add the `EosReplies` \
703 transform)"
704 .to_owned(),
705 ));
706 };
707 let mut headers = msg.headers().clone();
708 headers.remove(EOS_SOURCE_HEADER);
709 let stripped = OutgoingMessage::new(msg.name(), msg.payload()).with_headers(headers);
710 self.publish(&source, stripped).await
711 }
712}
713
714#[cfg(test)]
715mod tests {
716 use ruststream::Headers;
717
718 use super::*;
719
720 #[test]
721 fn source_header_roundtrips_topics_with_colons() {
722 let encoded = encode_source("orders:eu:v1", 3, 42);
723 let decoded = decode_source(&encoded).expect("decodes");
724 assert_eq!(decoded, SourceOffset::new("orders:eu:v1", 3, 42));
725 }
726
727 #[test]
728 fn malformed_source_headers_are_rejected() {
729 for bad in ["", "3", "3:x:orders", "x:42:orders"] {
730 assert!(decode_source(bad).is_none(), "{bad:?} must not decode");
731 }
732 }
733
734 #[tokio::test]
735 async fn reply_without_source_coordinates_fails_clearly() {
736 let pipeline = EosPipeline::new(KafkaPublisher::new(Arc::default()).transactional_id("p1"));
737 let err = Publisher::publish(&pipeline, OutgoingMessage::new("replies", b"x".as_slice()))
738 .await
739 .expect_err("a reply without the source header must fail");
740 assert!(matches!(err, KafkaError::InvalidOptions(_)));
741 assert!(err.to_string().contains("Commit::Transactional"));
742 let _ = Headers::new();
743 }
744}