ruststream_rdkafka/publisher.rs
1//! The publish policies and the live publishers they pair into, transactions included.
2
3// Without the `testing` feature a transport enum has one variant, so a `match` on it has a single
4// arm; the matches stay so the in-process arm has its place when the feature is on.
5#![cfg_attr(
6 not(feature = "testing"),
7 allow(clippy::infallible_destructuring_match)
8)]
9
10use std::collections::HashMap;
11use std::fmt;
12use std::future::{Future, ready};
13use std::sync::{Arc, Mutex};
14use std::time::Duration;
15
16use rdkafka::TopicPartitionList;
17use rdkafka::consumer::ConsumerGroupMetadata;
18use rdkafka::producer::{FutureProducer, FutureRecord, Producer as _};
19use rdkafka::util::Timeout;
20#[cfg(feature = "asyncapi")]
21use ruststream::asyncapi::Bindings;
22use ruststream::runtime::{OutPipeline, PublishBuilder, PublishSink, Slot};
23use ruststream::{
24 DefaultPublish, Lend, OutgoingMessage, PairError, PublishPolicy, Publisher,
25 TransactionalPublisher,
26};
27use tokio::sync::OnceCell;
28use tokio::task;
29
30#[cfg(feature = "testing")]
31use rdkafka::types::RDKafkaType;
32#[cfg(feature = "testing")]
33use ruststream::Str;
34
35use crate::broker::{ConnState, ConnectedKafkaBroker};
36use crate::convert;
37use crate::error::KafkaError;
38#[cfg(feature = "testing")]
39use crate::in_process::{Cluster, ProducerId, create_offline};
40
41const DEFAULT_TRANSACTION_TIMEOUT: Duration = Duration::from_secs(30);
42
43/// What one Kafka record carries beyond its topic, payload and headers.
44///
45/// A publish builder writes these through the steps of [`KafkaPublishSteps`]; nothing else needs
46/// to name the type, except an assertion reading back what a publish carried
47/// (`tb.out::<Marker>().with_options(&KafkaOptions::default().partition(3))`).
48///
49/// The record key is deliberately absent: it is the core's [`Partitioned`](ruststream::Partitioned)
50/// contract, carried by the [`PARTITION_KEY_HEADER`](crate::PARTITION_KEY_HEADER) header so that
51/// broker-agnostic code can set it, and Kafka maps it onto the native key.
52///
53/// # Examples
54///
55/// A test reading back the partition a handler's publish asked for:
56///
57/// ```
58/// # #[cfg(all(feature = "json", feature = "testing"))]
59/// # mod demo {
60/// use std::error::Error;
61///
62/// use ruststream::testing::TestApp;
63/// use ruststream_rdkafka::prelude::*;
64/// use serde::{Deserialize, Serialize};
65///
66/// #[derive(Serialize, Deserialize, Outgoing)]
67/// struct Order {
68/// id: u64,
69/// }
70///
71/// #[derive(Serialize, Outgoing)]
72/// #[outgoing(name = "audit")]
73/// struct Audit {
74/// id: u64,
75/// }
76///
77/// #[derive(OutSlot)]
78/// #[publishes(Audit)]
79/// struct Journal;
80///
81/// #[subscriber("orders")]
82/// async fn record(
83/// order: &Order,
84/// Out(journal): Out<impl Publisher<Options = KafkaOptions>, Journal>,
85/// ) -> HandlerOutcome {
86/// let sent = journal.message(&Audit { id: order.id }).partition(3).publish().await;
87/// if sent.is_err() {
88/// return HandlerOutcome::retry();
89/// }
90/// HandlerOutcome::ack()
91/// }
92///
93/// pub fn app() -> RustStream {
94/// RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(
95/// KafkaBroker::new(["kafka:9092"]).default_group("orders-svc"),
96/// |b| {
97/// b.include(record).out(Journal, Publish::default()).build();
98/// },
99/// )
100/// }
101///
102/// pub async fn audits_go_to_partition_three() -> Result<(), Box<dyn Error + Send + Sync>> {
103/// let tb = TestApp::start(app()).await?;
104///
105/// tb.broker::<KafkaBroker>()
106/// .message(&Order { id: 7 })
107/// .to("orders")
108/// .publish()
109/// .await?;
110///
111/// tb.out::<Journal>()
112/// .assert_called(1)
113/// .with_options(&KafkaOptions::default().partition(3));
114/// tb.shutdown().await?;
115/// Ok(())
116/// }
117/// # }
118/// # #[cfg(all(feature = "json", feature = "testing"))]
119/// # #[tokio::main(flavor = "multi_thread", worker_threads = 2)]
120/// # async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
121/// # demo::audits_go_to_partition_three().await
122/// # }
123/// # #[cfg(not(all(feature = "json", feature = "testing")))]
124/// # fn main() {}
125/// ```
126#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
127#[must_use]
128pub struct KafkaOptions {
129 partition: Option<i32>,
130}
131
132impl KafkaOptions {
133 /// Produces the record to exactly this partition, ahead of the partitioner and the record
134 /// key.
135 pub const fn partition(mut self, partition: i32) -> Self {
136 self.partition = Some(partition);
137 self
138 }
139
140 /// The partition this record is pinned to, or `None` when the partitioner places it.
141 #[must_use]
142 pub const fn partition_setting(self) -> Option<i32> {
143 self.partition
144 }
145}
146
147/// The per-message steps this crate adds to the publish builder.
148///
149/// A handler body that names one imports this crate's prelude and bounds its slot on the options
150/// type: `Out<impl Publisher<Options = KafkaOptions>, Marker>`. The bound on the sink is what
151/// keeps these steps off a builder over another broker's publisher.
152///
153/// # Examples
154///
155/// ```
156/// # #[cfg(feature = "json")]
157/// # mod demo {
158/// use ruststream_rdkafka::prelude::*;
159/// use serde::{Deserialize, Serialize};
160///
161/// #[derive(Deserialize)]
162/// struct Order {
163/// id: u64,
164/// }
165///
166/// #[derive(Serialize, Outgoing)]
167/// struct Audit {
168/// id: u64,
169/// }
170///
171/// #[derive(OutSlot)]
172/// #[publishes(Audit)]
173/// struct Journal;
174///
175/// /// The slot is bounded on the options type, which is what puts `partition(..)` on its builder.
176/// #[subscriber("orders")]
177/// async fn record(
178/// order: &Order,
179/// Out(journal): Out<impl Publisher<Options = KafkaOptions>, Journal>,
180/// ) -> HandlerOutcome {
181/// let sent = journal
182/// .message(&Audit { id: order.id })
183/// .to("audit")
184/// .partition(0)
185/// .publish()
186/// .await;
187/// if sent.is_err() {
188/// return HandlerOutcome::retry();
189/// }
190/// HandlerOutcome::ack()
191/// }
192///
193/// fn app() -> RustStream {
194/// RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(
195/// KafkaBroker::new(["localhost:9092"]).default_group("orders-svc"),
196/// |b| {
197/// b.include(record).out(Journal, Publish::default()).build();
198/// },
199/// )
200/// }
201/// # }
202/// # fn main() {}
203/// ```
204pub trait KafkaPublishSteps {
205 /// Produces this one record to `partition`, ahead of the partitioner and the record key.
206 #[must_use]
207 fn partition(self, partition: i32) -> Self;
208}
209
210impl<Sink, Body, Enc, Hdrs, Dest> KafkaPublishSteps for PublishBuilder<Sink, Body, Enc, Hdrs, Dest>
211where
212 Sink: PublishSink<Options = KafkaOptions>,
213{
214 fn partition(mut self, partition: i32) -> Self {
215 self.options_mut()
216 .get_or_insert_with(KafkaOptions::default)
217 .partition = Some(partition);
218 self
219 }
220}
221
222/// The publish policy of [`KafkaPublisher`]: pure declaration, no connection, no publish
223/// surface.
224///
225/// Constructible anywhere - in a router definition, in configuration, before startup - because
226/// it holds nothing but options. The runtime pairs it with the connected broker at startup (or
227/// [`ConnectedKafkaBroker::publisher`] does it by hand), and only the resulting
228/// [`KafkaPublisher`] can publish.
229///
230/// [`transactional_id`](Self::transactional_id) is a type transition, not a flag: it yields a
231/// [`KafkaTransactionalPublish`], so a plain publisher carries no transactional surface at all.
232///
233/// # Examples
234///
235/// ```
236/// use std::time::Duration;
237///
238/// use ruststream_rdkafka::prelude::*;
239/// use serde::{Deserialize, Serialize};
240///
241/// #[derive(Deserialize)]
242/// struct Order {
243/// id: u64,
244/// }
245///
246/// #[derive(Serialize, Outgoing)]
247/// #[outgoing(name = "confirmations")]
248/// struct Confirmation {
249/// id: u64,
250/// }
251///
252/// #[subscriber("orders", reply)]
253/// async fn confirm(order: &Order) -> Confirmation {
254/// Confirmation { id: order.id }
255/// }
256///
257/// fn app() -> RustStream {
258/// RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(
259/// KafkaBroker::new(["localhost:9092"]).default_group("orders-svc"),
260/// |b| {
261/// b.include(confirm)
262/// .out_reply(Publish::default().queue_timeout(Duration::from_secs(5)));
263/// },
264/// )
265/// }
266/// # let _ = app;
267/// ```
268#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
269#[must_use]
270pub struct KafkaPublish {
271 queue_timeout: Option<Duration>,
272}
273
274impl KafkaPublish {
275 /// How long a publish may wait for space when librdkafka's local queue is full, before
276 /// failing with a queue-full error. Without it a publish waits for space indefinitely,
277 /// which is the natural back-pressure behavior.
278 pub const fn queue_timeout(mut self, timeout: Duration) -> Self {
279 self.queue_timeout = Some(timeout);
280 self
281 }
282
283 /// Turns this into a transactional publish policy fenced by `id` (Kafka's
284 /// `transactional.id`).
285 ///
286 /// The id must be stable and unique per concurrent producer: Kafka uses it to fence
287 /// zombies, so two live producers sharing an id abort each other. Pair distinct policies
288 /// for concurrent transactional flows, or take one publisher per source partition with
289 /// [`KafkaTransactionalPublish::per_partition`].
290 ///
291 /// # Examples
292 ///
293 /// ```
294 /// use ruststream_rdkafka::prelude::*;
295 /// # use serde::{Deserialize, Serialize};
296 /// # #[derive(Deserialize)]
297 /// # struct Refund { order_id: u64, lines: u64 }
298 /// # #[derive(Serialize, Outgoing)]
299 /// # #[outgoing(name = "refund-lines")]
300 /// # struct RefundLine { order_id: u64, line: u64 }
301 /// # #[derive(OutSlot)]
302 /// # #[publishes(RefundLine)]
303 /// # struct Lines;
304 /// #
305 /// #[subscriber("refunds")]
306 /// async fn refund(
307 /// order: &Refund,
308 /// Out(lines): Out<impl TransactionalPublisher, Lines>,
309 /// ) -> HandlerOutcome {
310 /// if lines.begin_transaction().await.is_err() {
311 /// return HandlerOutcome::retry();
312 /// }
313 /// for line in 0..order.lines {
314 /// let entry = RefundLine {
315 /// order_id: order.order_id,
316 /// line,
317 /// };
318 /// if lines.message(&entry).publish().await.is_err() {
319 /// lines.abort().await.ok();
320 /// return HandlerOutcome::retry();
321 /// }
322 /// }
323 /// if lines.commit().await.is_err() {
324 /// return HandlerOutcome::retry();
325 /// }
326 /// HandlerOutcome::ack()
327 /// }
328 ///
329 /// fn app() -> RustStream {
330 /// RustStream::new(AppInfo::new("refunds", "0.1.0")).with_broker(
331 /// KafkaBroker::new(["localhost:9092"]).default_group("refunds-svc"),
332 /// |b| {
333 /// b.include(refund)
334 /// .out(Lines, Publish::default().transactional_id("refunds-svc-1"))
335 /// .build();
336 /// },
337 /// )
338 /// }
339 /// # let _ = app;
340 /// ```
341 pub fn transactional_id(self, id: impl Into<String>) -> KafkaTransactionalPublish {
342 KafkaTransactionalPublish {
343 queue_timeout: self.queue_timeout,
344 id: id.into(),
345 transaction_timeout: DEFAULT_TRANSACTION_TIMEOUT,
346 }
347 }
348
349 /// Frames every message this policy publishes in the Confluent wire format, resolving each
350 /// destination topic's subject through `registry`.
351 ///
352 /// This is what lets a Protobuf handler stay a plain function that returns its reply: the
353 /// value writes bare Protobuf through `#[wire(encode = ::prost::Message::encode)]`, and the
354 /// publisher this pairs into puts the schema id and the message-index path in front of it. A
355 /// value cannot do that itself - `Serialized::wire_bytes` is synchronous and has only
356 /// `&self` - while the publish path knows the destination topic, which is what names the
357 /// subject.
358 ///
359 /// Short for [`KafkaFramedPublish::over`](crate::KafkaFramedPublish::over) over the default
360 /// producer settings; take that one to keep the settings a configured [`KafkaPublish`]
361 /// holds. See [`KafkaFramedPublish`](crate::KafkaFramedPublish) for what it refuses, and for
362 /// when it resolves a subject.
363 ///
364 /// # Examples
365 ///
366 /// ```
367 /// use ruststream_rdkafka::prelude::*;
368 /// use ruststream_rdkafka::SchemaRegistry;
369 ///
370 /// #[derive(Clone, PartialEq, prost::Message, Deserialized)]
371 /// #[wire(decode = ruststream_rdkafka::protobuf::decode_confluent)]
372 /// struct Order {
373 /// #[prost(int64, tag = "1")]
374 /// id: i64,
375 /// }
376 ///
377 /// #[derive(Clone, PartialEq, prost::Message, Serialized, Outgoing)]
378 /// #[wire(encode = ::prost::Message::encode)]
379 /// #[outgoing(name = "confirmations")]
380 /// struct Confirmation {
381 /// #[prost(int64, tag = "1")]
382 /// id: i64,
383 /// }
384 ///
385 /// #[subscriber("orders", reply)]
386 /// async fn confirm(order: &Order) -> Confirmation {
387 /// Confirmation { id: order.id }
388 /// }
389 ///
390 /// fn app() -> RustStream {
391 /// let registry = SchemaRegistry::new("http://localhost:8081");
392 /// RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(
393 /// KafkaBroker::new(["localhost:9092"]).default_group("orders-svc"),
394 /// |b| {
395 /// b.include(confirm).out_reply(Publish::framed(®istry));
396 /// },
397 /// )
398 /// }
399 /// # let _ = app;
400 /// ```
401 #[cfg(feature = "protobuf")]
402 pub fn framed(registry: &crate::SchemaRegistry) -> crate::protobuf::KafkaFramedPublish {
403 crate::protobuf::KafkaFramedPublish::over(Self::default(), registry)
404 }
405
406 pub(crate) const fn queue_timeout_setting(self) -> Option<Duration> {
407 self.queue_timeout
408 }
409}
410
411impl PublishPolicy<ConnectedKafkaBroker> for KafkaPublish {
412 type Live = KafkaPublisher;
413
414 fn pair(
415 self,
416 connected: &ConnectedKafkaBroker,
417 ) -> impl Future<Output = Result<Self::Live, PairError>> {
418 ready(Ok(connected.publisher(self)))
419 }
420
421 /// The destination topic of this publish, which is what the channel stands for.
422 #[cfg(feature = "asyncapi")]
423 fn channel_bindings(&self, channel: &str) -> Bindings {
424 crate::bindings::channel(channel)
425 }
426}
427
428impl DefaultPublish for ConnectedKafkaBroker {
429 type Policy = KafkaPublish;
430}
431
432/// A live producer handle on the broker's shared producer.
433///
434/// [`OutgoingMessage::name`] is the destination topic. A
435/// [`PARTITION_KEY_HEADER`](crate::PARTITION_KEY_HEADER) header becomes the record's native key,
436/// so Kafka routes messages that share a key to the same partition; without it the configured
437/// partitioner picks one.
438///
439/// Each publish awaits the broker's delivery report, so an `Ok` means the cluster accepted the
440/// record (durability then depends on the producer's `acks` setting, configurable through
441/// [`KafkaBroker::producer_config`](crate::KafkaBroker::producer_config)).
442///
443/// Exists only from a connected broker, so it never sees a "not connected" state; it does alias
444/// that connection and may outlive it, so after the broker shuts down every publish reports
445/// [`KafkaError::Closed`]. Cheap to clone.
446#[derive(Clone)]
447pub struct KafkaPublisher {
448 state: Arc<ConnState>,
449 queue_timeout: Option<Duration>,
450}
451
452impl fmt::Debug for KafkaPublisher {
453 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
454 f.debug_struct("KafkaPublisher")
455 .field("queue_timeout", &self.queue_timeout)
456 .finish_non_exhaustive()
457 }
458}
459
460impl KafkaPublisher {
461 pub(crate) const fn new(state: Arc<ConnState>, queue_timeout: Option<Duration>) -> Self {
462 Self {
463 state,
464 queue_timeout,
465 }
466 }
467}
468
469/// Sends one record through `producer` and awaits its delivery report.
470async fn send_via(
471 producer: &FutureProducer,
472 queue_timeout: Option<Duration>,
473 msg: OutgoingMessage<'_>,
474 options: Option<&KafkaOptions>,
475) -> Result<(), KafkaError> {
476 let parts = convert::headers_for_publish(msg.headers());
477 let mut record = FutureRecord::<[u8], [u8]>::to(msg.name()).payload(msg.payload());
478 if let Some(key) = &parts.key {
479 record = record.key(key.as_ref());
480 }
481 if let Some(partition) = options.and_then(|options| options.partition_setting()) {
482 // An explicit partition wins over the partitioner and the record key.
483 record = record.partition(partition);
484 }
485 if let Some(headers) = parts.headers {
486 record = record.headers(headers);
487 }
488 let queue_timeout = queue_timeout.map_or(Timeout::Never, Timeout::After);
489 producer
490 .send(record, queue_timeout)
491 .await
492 .map(|_delivery| ())
493 .map_err(|(err, _record)| KafkaError::publish(err))
494}
495
496impl Publisher for KafkaPublisher {
497 // librdkafka's record takes the payload as `&[u8]` and copies it into the producer queue,
498 // so this transport reads the bytes and keeps nothing.
499 type Payload = Lend;
500 type Error = KafkaError;
501 type Options = KafkaOptions;
502
503 /// Publishes `msg` to the topic named by [`OutgoingMessage::name`] and awaits the delivery
504 /// report.
505 ///
506 /// # Errors
507 ///
508 /// Returns [`KafkaError::Closed`] once the connection this handle aliases has been shut
509 /// down, and [`KafkaError::Publish`] when the cluster rejects the record or the delivery
510 /// times out (librdkafka's `message.timeout.ms`).
511 ///
512 /// # Cancel safety
513 ///
514 /// Not cancel safe: dropping the future may leave the record in flight, delivered or not.
515 async fn publish(
516 &self,
517 msg: OutgoingMessage<'_>,
518 options: Option<&Self::Options>,
519 ) -> Result<(), Self::Error> {
520 self.state.ensure_open(msg.name())?;
521 #[cfg(feature = "testing")]
522 if let Some(cluster) = self.state.in_process() {
523 return produce_in_process(cluster, &msg, options, None);
524 }
525 send_via(self.state.producer()?, self.queue_timeout, msg, options).await
526 }
527}
528
529/// Hands one record to the in-process cluster, as `send_via` hands it to librdkafka.
530#[cfg(feature = "testing")]
531fn produce_in_process(
532 cluster: &Cluster,
533 msg: &OutgoingMessage<'_>,
534 options: Option<&KafkaOptions>,
535 transaction: Option<&ProducerId>,
536) -> Result<(), KafkaError> {
537 cluster.produce(
538 msg.name(),
539 options.and_then(|options| options.partition_setting()),
540 msg.payload(),
541 msg.headers(),
542 transaction,
543 )
544}
545
546/// The publish policy of [`KafkaTransactionalPublisher`]: the transactional mode as its own
547/// type, reached from [`KafkaPublish::transactional_id`].
548///
549/// Pairing it is where Kafka does the real work: the transactional producer is created and its
550/// transactions initialized (the call that fences earlier producers with the same id), so a
551/// misconfigured transactional id fails at startup rather than at the first transaction.
552///
553/// # Examples
554///
555/// ```
556/// use std::time::Duration;
557///
558/// use ruststream_rdkafka::prelude::*;
559/// # use serde::{Deserialize, Serialize};
560/// # #[derive(Deserialize)]
561/// # struct Refund { order_id: u64, lines: u64 }
562/// # #[derive(Serialize, Outgoing)]
563/// # #[outgoing(name = "refund-lines")]
564/// # struct RefundLine { order_id: u64, line: u64 }
565/// # #[derive(OutSlot)]
566/// # #[publishes(RefundLine)]
567/// # struct Lines;
568/// #
569/// #[subscriber("refunds")]
570/// async fn refund(
571/// order: &Refund,
572/// Out(lines): Out<impl TransactionalPublisher, Lines>,
573/// ) -> HandlerOutcome {
574/// if lines.begin_transaction().await.is_err() {
575/// return HandlerOutcome::retry();
576/// }
577/// for line in 0..order.lines {
578/// let entry = RefundLine {
579/// order_id: order.order_id,
580/// line,
581/// };
582/// if lines.message(&entry).publish().await.is_err() {
583/// lines.abort().await.ok();
584/// return HandlerOutcome::retry();
585/// }
586/// }
587/// if lines.commit().await.is_err() {
588/// return HandlerOutcome::retry();
589/// }
590/// HandlerOutcome::ack()
591/// }
592///
593/// fn app() -> RustStream {
594/// let lines = Publish::default()
595/// .transactional_id("refunds-svc-1")
596/// .transaction_timeout(Duration::from_secs(10));
597/// RustStream::new(AppInfo::new("refunds", "0.1.0")).with_broker(
598/// KafkaBroker::new(["localhost:9092"]).default_group("refunds-svc"),
599/// |b| {
600/// b.include(refund).out(Lines, lines).build();
601/// },
602/// )
603/// }
604/// # let _ = app;
605/// ```
606#[derive(Debug, Clone, PartialEq, Eq)]
607#[must_use]
608pub struct KafkaTransactionalPublish {
609 queue_timeout: Option<Duration>,
610 id: String,
611 transaction_timeout: Duration,
612}
613
614impl KafkaTransactionalPublish {
615 /// See [`KafkaPublish::queue_timeout`].
616 pub const fn queue_timeout(mut self, timeout: Duration) -> Self {
617 self.queue_timeout = Some(timeout);
618 self
619 }
620
621 /// How long transaction control calls (`init`, `commit`, `abort`) may block before
622 /// reporting failure. Defaults to 30 seconds; this is the call deadline handed to
623 /// librdkafka, not its `transaction.timeout.ms` (reachable through
624 /// [`KafkaBroker::producer_config`](crate::KafkaBroker::producer_config)).
625 pub const fn transaction_timeout(mut self, timeout: Duration) -> Self {
626 self.transaction_timeout = timeout;
627 self
628 }
629
630 /// Turns this into a per-partition policy: the id becomes the base of one transactional id
631 /// per source partition (`"{base}-p{partition}"`), pairing into
632 /// [`TransactionalPartitions`].
633 ///
634 /// # Examples
635 ///
636 /// ```
637 /// use ruststream_rdkafka::prelude::*;
638 /// # use serde::{Deserialize, Serialize};
639 /// # #[derive(Deserialize)]
640 /// # struct Invoice { id: u64, lines: u64 }
641 /// # #[derive(Serialize, Outgoing)]
642 /// # #[outgoing(name = "invoice-lines")]
643 /// # struct InvoiceLine { invoice_id: u64, line: u64 }
644 /// # #[subscriber(KafkaTopic::new("invoices").commit(Commit::Tracked), workers(4, by_key))]
645 /// # async fn bill(
646 /// # invoice: &Invoice,
647 /// # Ctx(partition): Ctx<Partition>,
648 /// # Out(lanes): Out<impl PartitionLanes>,
649 /// # ) -> HandlerOutcome {
650 /// # let Ok(lane) = lanes.for_partition(partition).await else {
651 /// # return HandlerOutcome::retry();
652 /// # };
653 /// # if lane.begin_transaction().await.is_err() {
654 /// # return HandlerOutcome::retry();
655 /// # }
656 /// # for line in 0..invoice.lines {
657 /// # let entry = InvoiceLine { invoice_id: invoice.id, line };
658 /// # if lane.message(&entry).publish().await.is_err() {
659 /// # lane.abort().await.ok();
660 /// # return HandlerOutcome::retry();
661 /// # }
662 /// # }
663 /// # if lane.commit().await.is_err() {
664 /// # return HandlerOutcome::retry();
665 /// # }
666 /// # HandlerOutcome::ack()
667 /// # }
668 ///
669 /// fn app() -> RustStream {
670 /// RustStream::new(AppInfo::new("billing", "0.1.0")).with_broker(
671 /// KafkaBroker::new(["localhost:9092"]).default_group("billing-svc"),
672 /// |b| {
673 /// // One transactional id per source partition: "billing-svc-1-p0", "-p1", ...
674 /// b.include(bill)
675 /// .out(
676 /// DefaultSlot,
677 /// Publish::default().transactional_id("billing-svc-1").per_partition(),
678 /// )
679 /// .build();
680 /// },
681 /// )
682 /// }
683 /// # let _ = app;
684 /// ```
685 pub fn per_partition(self) -> KafkaPartitionedPublish {
686 KafkaPartitionedPublish { template: self }
687 }
688
689 /// The transactional id this policy fences with.
690 #[must_use]
691 pub fn id(&self) -> &str {
692 &self.id
693 }
694
695 pub(crate) fn with_id(&self, id: String) -> Self {
696 Self {
697 queue_timeout: self.queue_timeout,
698 id,
699 transaction_timeout: self.transaction_timeout,
700 }
701 }
702}
703
704impl PublishPolicy<ConnectedKafkaBroker> for KafkaTransactionalPublish {
705 type Live = KafkaTransactionalPublisher;
706
707 async fn pair(self, connected: &ConnectedKafkaBroker) -> Result<Self::Live, PairError> {
708 connected
709 .transactional_publisher(self)
710 .await
711 .map_err(PairError::new)
712 }
713
714 /// The destination topic of this publish, which is what the channel stands for.
715 #[cfg(feature = "asyncapi")]
716 fn channel_bindings(&self, channel: &str) -> Bindings {
717 crate::bindings::channel(channel)
718 }
719}
720
721impl ConnectedKafkaBroker {
722 /// A live transactional publisher: creates the transactional producer from the broker's
723 /// resolved producer configuration and initializes its transactions.
724 ///
725 /// Async because the initialization is real work (it fences earlier producers holding the
726 /// same transactional id); it runs once, when the publisher comes alive.
727 ///
728 /// # Errors
729 ///
730 /// Returns [`KafkaError::Closed`] once the connection has been shut down and
731 /// [`KafkaError::Publish`] when the producer cannot be created or the initialization fails
732 /// within the policy's [`transaction_timeout`](KafkaTransactionalPublish::transaction_timeout).
733 ///
734 /// # Examples
735 ///
736 /// A migration tool that copies a batch of records all-or-nothing:
737 ///
738 /// ```no_run
739 /// use std::error::Error;
740 ///
741 /// use ruststream::Broker;
742 /// use ruststream_rdkafka::prelude::*;
743 /// use serde::Serialize;
744 ///
745 /// #[derive(Serialize, Outgoing)]
746 /// #[outgoing(name = "accounts-v2")]
747 /// struct Account {
748 /// id: u64,
749 /// }
750 ///
751 /// async fn migrate(accounts: &[Account]) -> Result<(), Box<dyn Error + Send + Sync>> {
752 /// let connected = KafkaBroker::new(["localhost:9092"]).connect().await?;
753 /// let publisher = connected
754 /// .transactional_publisher(Publish::default().transactional_id("accounts-migration"))
755 /// .await?;
756 /// publisher.begin_transaction().await?;
757 /// for account in accounts {
758 /// if let Err(err) = publisher.message(account).publish().await {
759 /// publisher.abort().await?;
760 /// return Err(err.into());
761 /// }
762 /// }
763 /// publisher.commit().await?;
764 /// Ok(())
765 /// }
766 /// # let _ = migrate;
767 /// ```
768 pub async fn transactional_publisher(
769 &self,
770 policy: KafkaTransactionalPublish,
771 ) -> Result<KafkaTransactionalPublisher, KafkaError> {
772 open_transactional(self.state(), &policy).await
773 }
774}
775
776/// Creates and initializes one transactional producer for `policy`.
777pub(crate) async fn open_transactional(
778 state: &Arc<ConnState>,
779 policy: &KafkaTransactionalPublish,
780) -> Result<KafkaTransactionalPublisher, KafkaError> {
781 state.ensure_open(policy.id())?;
782 let mut config = state.producer_config().clone();
783 config.set("transactional.id", policy.id());
784 #[cfg(feature = "testing")]
785 if let Some(cluster) = state.in_process() {
786 // The configuration is validated as the live producer's creation would validate it, and
787 // the pairing fences every earlier one with this id, as `init_transactions` does.
788 create_offline(&config, RDKafkaType::RD_KAFKA_PRODUCER).map_err(KafkaError::publish)?;
789 let producer = TxProducer::InProcess {
790 cluster: Arc::clone(cluster),
791 id: cluster.init_transactions(policy.id()),
792 };
793 return Ok(KafkaTransactionalPublisher::with(state, producer, policy));
794 }
795 let producer: FutureProducer = config.create().map_err(KafkaError::publish)?;
796 // init_transactions blocks (it fences earlier producers with this id), so it runs on the
797 // blocking pool.
798 let init = producer.clone();
799 let timeout = policy.transaction_timeout;
800 task::spawn_blocking(move || init.init_transactions(timeout))
801 .await
802 .map_err(|err| KafkaError::Publish(Box::new(err)))?
803 .map_err(KafkaError::publish)?;
804 Ok(KafkaTransactionalPublisher::with(
805 state,
806 TxProducer::Kafka(producer),
807 policy,
808 ))
809}
810
811/// The producer a transactional publisher runs its transactions on: a librdkafka producer, or,
812/// under the `testing` feature, a producer id paired with the in-process cluster.
813///
814/// Without the feature there is one variant, so the type is the producer itself.
815enum TxProducer {
816 Kafka(FutureProducer),
817 #[cfg(feature = "testing")]
818 InProcess {
819 cluster: Arc<Cluster>,
820 id: ProducerId,
821 },
822}
823
824#[cfg(not(feature = "testing"))]
825const _: () = assert!(size_of::<TxProducer>() == size_of::<FutureProducer>());
826
827struct TxInner {
828 state: Arc<ConnState>,
829 producer: TxProducer,
830 queue_timeout: Option<Duration>,
831 timeout: Duration,
832 id: String,
833 /// Whether a transaction is currently open. Interleaving `publish` with
834 /// `begin_transaction`/`commit` from concurrent tasks is not supported: which side of the
835 /// transaction boundary a concurrent publish lands on would be a race either way.
836 open: Mutex<bool>,
837}
838
839/// A live publisher that produces inside Kafka transactions.
840///
841/// Records between `begin_transaction` and `commit` become visible atomically (readers on
842/// Kafka's default `read_committed` isolation see all of them or none); `abort` discards them
843/// broker-side.
844///
845/// Its transactional producer is created and initialized when the
846/// [`KafkaTransactionalPublish`] policy pairs, so nothing is lazy here: the handle is fenced
847/// from the moment it exists. Clones share one producer and one transaction state.
848///
849/// Outside an open transaction the handle publishes like a plain one. A second
850/// `begin_transaction` reports [`KafkaError::TransactionBusy`], and a `commit` or `abort` with
851/// nothing open reports [`KafkaError::NoTransaction`].
852///
853/// # Examples
854///
855/// A handler receives the publisher through an `Out` slot bounded by the core's
856/// [`TransactionalPublisher`] capability, so its body names nothing of this crate's:
857///
858/// ```
859/// # #[cfg(feature = "json")]
860/// # mod demo {
861/// use ruststream_rdkafka::prelude::*;
862/// use serde::{Deserialize, Serialize};
863///
864/// #[derive(Deserialize)]
865/// struct Refund {
866/// order_id: u64,
867/// lines: u64,
868/// }
869///
870/// #[derive(Serialize, Outgoing)]
871/// #[outgoing(name = "refund-lines")]
872/// struct RefundLine {
873/// order_id: u64,
874/// line: u64,
875/// }
876///
877/// #[derive(OutSlot)]
878/// #[publishes(RefundLine)]
879/// struct Lines;
880///
881/// #[subscriber("refunds")]
882/// async fn refund(
883/// order: &Refund,
884/// Out(lines): Out<impl TransactionalPublisher, Lines>,
885/// ) -> HandlerOutcome {
886/// if lines.begin_transaction().await.is_err() {
887/// return HandlerOutcome::retry();
888/// }
889/// for line in 0..order.lines {
890/// let entry = RefundLine {
891/// order_id: order.order_id,
892/// line,
893/// };
894/// if lines.message(&entry).publish().await.is_err() {
895/// lines.abort().await.ok();
896/// return HandlerOutcome::retry();
897/// }
898/// }
899/// if lines.commit().await.is_err() {
900/// return HandlerOutcome::retry();
901/// }
902/// HandlerOutcome::ack()
903/// }
904///
905/// fn app() -> RustStream {
906/// RustStream::new(AppInfo::new("refunds", "0.1.0")).with_broker(
907/// KafkaBroker::new(["localhost:9092"]).default_group("refunds-svc"),
908/// |b| {
909/// b.include(refund)
910/// .out(Lines, Publish::default().transactional_id("refunds-svc-1"))
911/// .build();
912/// },
913/// )
914/// }
915/// # }
916/// # fn main() {}
917/// ```
918#[derive(Clone)]
919pub struct KafkaTransactionalPublisher {
920 inner: Arc<TxInner>,
921}
922
923impl fmt::Debug for KafkaTransactionalPublisher {
924 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
925 f.debug_struct("KafkaTransactionalPublisher")
926 .field("id", &self.inner.id)
927 .field("timeout", &self.inner.timeout)
928 .finish_non_exhaustive()
929 }
930}
931
932impl KafkaTransactionalPublisher {
933 fn with(
934 state: &Arc<ConnState>,
935 producer: TxProducer,
936 policy: &KafkaTransactionalPublish,
937 ) -> Self {
938 Self {
939 inner: Arc::new(TxInner {
940 state: Arc::clone(state),
941 producer,
942 queue_timeout: policy.queue_timeout,
943 timeout: policy.transaction_timeout,
944 id: policy.id.clone(),
945 open: Mutex::new(false),
946 }),
947 }
948 }
949
950 /// The transactional id fencing this publisher.
951 #[must_use]
952 pub fn id(&self) -> &str {
953 &self.inner.id
954 }
955
956 pub(crate) fn deadline(&self) -> Duration {
957 self.inner.timeout
958 }
959
960 pub(crate) fn state(&self) -> &Arc<ConnState> {
961 &self.inner.state
962 }
963
964 fn is_open(&self) -> bool {
965 *self
966 .inner
967 .open
968 .lock()
969 .expect("transaction state mutex poisoned")
970 }
971
972 fn set_open(&self, open: bool) {
973 *self
974 .inner
975 .open
976 .lock()
977 .expect("transaction state mutex poisoned") = open;
978 }
979
980 fn no_transaction(&self) -> KafkaError {
981 KafkaError::NoTransaction {
982 id: self.inner.id.clone(),
983 }
984 }
985
986 /// Adds consumed source offsets (and their group's metadata) to the open transaction, so
987 /// they commit atomically with the records published into it. The EOS pipeline's commit
988 /// path; must run between `begin_transaction` and `commit`.
989 pub(crate) async fn send_offsets(
990 &self,
991 offsets: TopicPartitionList,
992 metadata: ConsumerGroupMetadata,
993 ) -> Result<(), KafkaError> {
994 self.inner.state.ensure_open(&self.inner.id)?;
995 if !self.is_open() {
996 return Err(self.no_transaction());
997 }
998 let producer = match &self.inner.producer {
999 TxProducer::Kafka(producer) => producer.clone(),
1000 #[cfg(feature = "testing")]
1001 TxProducer::InProcess { .. } => {
1002 return Err(KafkaError::InvalidOptions(
1003 "librdkafka offsets reached an in-process transaction".to_owned(),
1004 ));
1005 }
1006 };
1007 let timeout = self.inner.timeout;
1008 task::spawn_blocking(move || {
1009 producer.send_offsets_to_transaction(&offsets, &metadata, timeout)
1010 })
1011 .await
1012 .map_err(|err| KafkaError::Publish(Box::new(err)))?
1013 .map_err(KafkaError::publish)
1014 }
1015}
1016
1017#[cfg(feature = "testing")]
1018impl KafkaTransactionalPublisher {
1019 /// Adds consumed positions to the open transaction on the in-process cluster, to commit into
1020 /// `group` with it: the in-process counterpart of [`send_offsets`](Self::send_offsets).
1021 pub(crate) fn send_offsets_in_process(
1022 &self,
1023 group: &str,
1024 positions: &[((Str, i32), i64)],
1025 ) -> Result<(), KafkaError> {
1026 self.inner.state.ensure_open(&self.inner.id)?;
1027 if !self.is_open() {
1028 return Err(self.no_transaction());
1029 }
1030 match &self.inner.producer {
1031 TxProducer::InProcess { cluster, id } => cluster.send_offsets(id, group, positions),
1032 TxProducer::Kafka(_) => Err(KafkaError::InvalidOptions(
1033 "in-process offsets reached a librdkafka transaction".to_owned(),
1034 )),
1035 }
1036 }
1037}
1038
1039impl Publisher for KafkaTransactionalPublisher {
1040 // The same record as the plain publisher's, inside a transaction.
1041 type Payload = Lend;
1042 type Error = KafkaError;
1043 type Options = KafkaOptions;
1044
1045 /// Publishes `msg` to the topic named by [`OutgoingMessage::name`]. Inside an open
1046 /// transaction the record joins it; otherwise it goes out through the broker's shared plain
1047 /// producer.
1048 ///
1049 /// # Errors
1050 ///
1051 /// Returns [`KafkaError::Closed`] once the connection this handle aliases has been shut
1052 /// down and [`KafkaError::Publish`] when the cluster rejects the record or the delivery
1053 /// times out.
1054 ///
1055 /// # Cancel safety
1056 ///
1057 /// Not cancel safe: dropping the future may leave the record in flight, delivered or not.
1058 async fn publish(
1059 &self,
1060 msg: OutgoingMessage<'_>,
1061 options: Option<&Self::Options>,
1062 ) -> Result<(), Self::Error> {
1063 self.inner.state.ensure_open(msg.name())?;
1064 let producer = match &self.inner.producer {
1065 TxProducer::Kafka(producer) => producer,
1066 #[cfg(feature = "testing")]
1067 TxProducer::InProcess { cluster, id } => {
1068 let transaction = self.is_open().then_some(id);
1069 return produce_in_process(cluster, &msg, options, transaction);
1070 }
1071 };
1072 if self.is_open() {
1073 return send_via(producer, self.inner.queue_timeout, msg, options).await;
1074 }
1075 send_via(
1076 self.inner.state.producer()?,
1077 self.inner.queue_timeout,
1078 msg,
1079 options,
1080 )
1081 .await
1082 }
1083}
1084
1085impl TransactionalPublisher for KafkaTransactionalPublisher {
1086 /// Begins a Kafka transaction.
1087 ///
1088 /// One producer runs one transaction at a time, so beginning while one is open is an
1089 /// error, not a queue: a second begin means two flows share one publisher, and silently
1090 /// merging their messages into one transaction would commit one flow's records with the
1091 /// other's. Concurrent transactional flows use distinct publishers (see
1092 /// [`TransactionalPartitions`]).
1093 ///
1094 /// # Errors
1095 ///
1096 /// Returns [`KafkaError::TransactionBusy`] when a transaction is already open on this
1097 /// publisher (or a clone sharing it), [`KafkaError::Closed`] after the broker shut down,
1098 /// and [`KafkaError::Publish`] when the begin call fails.
1099 // The guard intentionally spans the begin call: check-and-begin must be atomic so two
1100 // concurrent begins cannot both pass the check.
1101 fn begin_transaction(&self) -> impl Future<Output = Result<(), Self::Error>> {
1102 if let Err(err) = self.inner.state.ensure_open(&self.inner.id) {
1103 return ready(Err(err));
1104 }
1105 {
1106 let mut open = self
1107 .inner
1108 .open
1109 .lock()
1110 .expect("transaction state mutex poisoned");
1111 if *open {
1112 return ready(Err(KafkaError::TransactionBusy {
1113 id: self.inner.id.clone(),
1114 }));
1115 }
1116 // A rejected begin leaves the open transaction untouched, per the trait contract.
1117 let begun = match &self.inner.producer {
1118 TxProducer::Kafka(producer) => {
1119 producer.begin_transaction().map_err(KafkaError::publish)
1120 }
1121 #[cfg(feature = "testing")]
1122 TxProducer::InProcess { cluster, id } => {
1123 cluster.begin(id, self.inner.state.runtime())
1124 }
1125 };
1126 if let Err(err) = begun {
1127 return ready(Err(err));
1128 }
1129 *open = true;
1130 }
1131 ready(Ok(()))
1132 }
1133
1134 /// Commits the open transaction, making its records visible atomically.
1135 ///
1136 /// # Errors
1137 ///
1138 /// Returns [`KafkaError::NoTransaction`] when no transaction is open on this publisher,
1139 /// [`KafkaError::Closed`] after the broker shut down, and [`KafkaError::Publish`] when the
1140 /// commit fails. librdkafka distinguishes retriable failures from ones requiring an abort;
1141 /// after an error the transaction's state is unresolved, so treat the publisher as needing
1142 /// an [`abort`](TransactionalPublisher::abort) or replacement.
1143 async fn commit(&self) -> Result<(), Self::Error> {
1144 self.inner.state.ensure_open(&self.inner.id)?;
1145 if !self.is_open() {
1146 return Err(self.no_transaction());
1147 }
1148 let producer = match &self.inner.producer {
1149 TxProducer::Kafka(producer) => producer.clone(),
1150 #[cfg(feature = "testing")]
1151 TxProducer::InProcess { cluster, id } => {
1152 // As on the live path, only a commit that succeeded closes the handle's
1153 // transaction; after a failure the caller aborts.
1154 cluster.commit(id)?;
1155 self.set_open(false);
1156 return Ok(());
1157 }
1158 };
1159 let timeout = self.inner.timeout;
1160 task::spawn_blocking(move || producer.commit_transaction(timeout))
1161 .await
1162 .map_err(|err| KafkaError::Publish(Box::new(err)))?
1163 .map_err(KafkaError::publish)?;
1164 self.set_open(false);
1165 Ok(())
1166 }
1167
1168 /// Aborts the open transaction, discarding its records broker-side.
1169 ///
1170 /// # Errors
1171 ///
1172 /// Returns [`KafkaError::NoTransaction`] when no transaction is open on this publisher,
1173 /// [`KafkaError::Closed`] after the broker shut down, and [`KafkaError::Publish`] when the
1174 /// abort fails.
1175 async fn abort(&self) -> Result<(), Self::Error> {
1176 self.inner.state.ensure_open(&self.inner.id)?;
1177 if !self.is_open() {
1178 return Err(self.no_transaction());
1179 }
1180 let producer = match &self.inner.producer {
1181 TxProducer::Kafka(producer) => producer.clone(),
1182 #[cfg(feature = "testing")]
1183 TxProducer::InProcess { cluster, id } => {
1184 let aborted = cluster.abort(id);
1185 self.set_open(false);
1186 return aborted;
1187 }
1188 };
1189 let timeout = self.inner.timeout;
1190 let aborted = task::spawn_blocking(move || producer.abort_transaction(timeout))
1191 .await
1192 .map_err(|err| KafkaError::Publish(Box::new(err)))?;
1193 // The transaction is over either way: a failed abort resolves broker-side by its own
1194 // timeout, and leaving the handle "open" would wedge it permanently.
1195 self.set_open(false);
1196 aborted.map_err(KafkaError::publish)
1197 }
1198}
1199
1200/// The publish policy of [`TransactionalPartitions`], reached from
1201/// [`KafkaTransactionalPublish::per_partition`].
1202///
1203/// # Examples
1204///
1205/// ```
1206/// use ruststream_rdkafka::prelude::*;
1207/// # use serde::{Deserialize, Serialize};
1208/// # #[derive(Deserialize)]
1209/// # struct Invoice { id: u64, lines: u64 }
1210/// # #[derive(Serialize, Outgoing)]
1211/// # #[outgoing(name = "invoice-lines")]
1212/// # struct InvoiceLine { invoice_id: u64, line: u64 }
1213/// # #[subscriber(KafkaTopic::new("invoices").commit(Commit::Tracked), workers(4, by_key))]
1214/// # async fn bill(
1215/// # invoice: &Invoice,
1216/// # Ctx(partition): Ctx<Partition>,
1217/// # Out(lanes): Out<impl PartitionLanes>,
1218/// # ) -> HandlerOutcome {
1219/// # let Ok(lane) = lanes.for_partition(partition).await else {
1220/// # return HandlerOutcome::retry();
1221/// # };
1222/// # if lane.begin_transaction().await.is_err() {
1223/// # return HandlerOutcome::retry();
1224/// # }
1225/// # for line in 0..invoice.lines {
1226/// # let entry = InvoiceLine { invoice_id: invoice.id, line };
1227/// # if lane.message(&entry).publish().await.is_err() {
1228/// # lane.abort().await.ok();
1229/// # return HandlerOutcome::retry();
1230/// # }
1231/// # }
1232/// # if lane.commit().await.is_err() {
1233/// # return HandlerOutcome::retry();
1234/// # }
1235/// # HandlerOutcome::ack()
1236/// # }
1237///
1238/// fn app() -> RustStream {
1239/// RustStream::new(AppInfo::new("billing", "0.1.0")).with_broker(
1240/// KafkaBroker::new(["localhost:9092"]).default_group("billing-svc"),
1241/// |b| {
1242/// // One transactional id per source partition: "billing-svc-1-p0", "-p1", ...
1243/// b.include(bill)
1244/// .out(
1245/// DefaultSlot,
1246/// Publish::default().transactional_id("billing-svc-1").per_partition(),
1247/// )
1248/// .build();
1249/// },
1250/// )
1251/// }
1252/// # let _ = app;
1253/// ```
1254#[derive(Debug, Clone, PartialEq, Eq)]
1255#[must_use]
1256pub struct KafkaPartitionedPublish {
1257 template: KafkaTransactionalPublish,
1258}
1259
1260impl PublishPolicy<ConnectedKafkaBroker> for KafkaPartitionedPublish {
1261 type Live = TransactionalPartitions;
1262
1263 fn pair(
1264 self,
1265 connected: &ConnectedKafkaBroker,
1266 ) -> impl Future<Output = Result<Self::Live, PairError>> {
1267 ready(Ok(TransactionalPartitions {
1268 inner: Arc::new(PartitionsInner {
1269 state: Arc::clone(connected.state()),
1270 template: self.template,
1271 publishers: Mutex::new(HashMap::new()),
1272 }),
1273 }))
1274 }
1275
1276 /// The destination topic of this publish, which is what the channel stands for.
1277 #[cfg(feature = "asyncapi")]
1278 fn channel_bindings(&self, channel: &str) -> Bindings {
1279 crate::bindings::channel(channel)
1280 }
1281}
1282
1283/// Transactional publishers, one per source partition, materialized on first use.
1284///
1285/// Kafka permits one open transaction per producer and one live producer per transactional id
1286/// (initializing a second fences the first), so concurrent transactional handlers need one
1287/// producer each. The source partition is the natural scope: under the default
1288/// [`LaneKey::Partition`](crate::LaneKey::Partition) worker pool a partition's deliveries
1289/// process serially on one lane, so a publisher per partition gives every lane an independent
1290/// transaction with no coordination. The id set (`"{base}-p{partition}"`) follows the topic's
1291/// partitions rather than the worker count: changing `workers(n)` neither changes the ids nor
1292/// weakens zombie fencing - the scheme Kafka Streams uses for its per-task producers.
1293///
1294/// Not for [`LaneKey::RecordKey`](crate::LaneKey::RecordKey) pools: record-key lanes spread
1295/// one partition across lanes, so two lanes would share a partition's publisher and collide
1296/// on its single transaction ([`KafkaError::TransactionBusy`]).
1297///
1298/// Clones share the cache, so one injected handle serves every handler invocation.
1299#[derive(Clone)]
1300pub struct TransactionalPartitions {
1301 inner: Arc<PartitionsInner>,
1302}
1303
1304struct PartitionsInner {
1305 state: Arc<ConnState>,
1306 template: KafkaTransactionalPublish,
1307 /// The per-partition publishers. The set of partitions is only known as deliveries arrive,
1308 /// so materialization stays lazy here (a cell per partition, so two lanes racing the same
1309 /// partition initialize one producer, not two).
1310 publishers: Mutex<HashMap<i32, Arc<OnceCell<KafkaTransactionalPublisher>>>>,
1311}
1312
1313impl fmt::Debug for TransactionalPartitions {
1314 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1315 f.debug_struct("TransactionalPartitions")
1316 .field("id_base", &self.inner.template.id())
1317 .finish_non_exhaustive()
1318 }
1319}
1320
1321impl TransactionalPartitions {
1322 /// The publisher owning `partition`'s transactional id, created and initialized on first
1323 /// use.
1324 ///
1325 /// `partition` is the delivery's source partition (`KafkaContext`'s `Partition` field in a
1326 /// handler); passing anything else still works but forfeits the serialization argument
1327 /// that makes the per-partition scope safe.
1328 ///
1329 /// # Errors
1330 ///
1331 /// Returns [`KafkaError::Closed`] once the broker has shut down and
1332 /// [`KafkaError::Publish`] when the partition's producer cannot be created or its
1333 /// transactions cannot be initialized.
1334 ///
1335 /// # Panics
1336 ///
1337 /// Panics when the internal cache mutex is poisoned, which requires a prior panic while
1338 /// materializing a publisher (an invariant violation, not an operational failure).
1339 pub async fn for_partition(
1340 &self,
1341 partition: i32,
1342 ) -> Result<KafkaTransactionalPublisher, KafkaError> {
1343 let cell = {
1344 let mut publishers = self
1345 .inner
1346 .publishers
1347 .lock()
1348 .expect("partition publisher cache mutex poisoned");
1349 Arc::clone(publishers.entry(partition).or_default())
1350 };
1351 let policy = self
1352 .inner
1353 .template
1354 .with_id(format!("{}-p{partition}", self.inner.template.id()));
1355 cell.get_or_try_init(|| open_transactional(&self.inner.state, &policy))
1356 .await
1357 .cloned()
1358 }
1359}
1360
1361/// The capability of handing out one transactional publisher per source partition.
1362///
1363/// A handler that drives per-partition transactions names this in its `Out` slot
1364/// (`Out(lanes): Out<impl PartitionLanes>`) instead of a publisher type: the concrete value is
1365/// [`TransactionalPartitions`], inferred from the [`KafkaTransactionalPublish::per_partition`]
1366/// policy attached at the include site.
1367///
1368/// # Test capture
1369///
1370/// This is a router, not a publisher: it hands out a publisher of its own rather than sending a
1371/// message. What a lane then publishes leaves through that publisher, so it lands in the
1372/// broker's publish log and not in the slot's test record - the same boundary a settled owned
1373/// transaction's buffer has. Assert on the publish log for lane traffic, and keep the slot
1374/// record for handlers that publish through the slot itself.
1375///
1376/// # Examples
1377///
1378/// ```
1379/// use ruststream_rdkafka::prelude::*;
1380/// use serde::{Deserialize, Serialize};
1381///
1382/// #[derive(Deserialize)]
1383/// struct Invoice {
1384/// id: u64,
1385/// lines: u64,
1386/// }
1387///
1388/// #[derive(Serialize, Outgoing)]
1389/// #[outgoing(name = "invoice-lines")]
1390/// struct InvoiceLine {
1391/// invoice_id: u64,
1392/// line: u64,
1393/// }
1394///
1395/// /// Each lane runs the transactions of its own source partition, so four workers never share
1396/// /// one producer.
1397/// #[subscriber(KafkaTopic::new("invoices").commit(Commit::Tracked), workers(4, by_key))]
1398/// async fn bill(
1399/// invoice: &Invoice,
1400/// Ctx(partition): Ctx<Partition>,
1401/// Out(lanes): Out<impl PartitionLanes>,
1402/// ) -> HandlerOutcome {
1403/// let Ok(lane) = lanes.for_partition(partition).await else {
1404/// return HandlerOutcome::retry();
1405/// };
1406/// if lane.begin_transaction().await.is_err() {
1407/// return HandlerOutcome::retry();
1408/// }
1409/// for line in 0..invoice.lines {
1410/// let entry = InvoiceLine { invoice_id: invoice.id, line };
1411/// if lane.message(&entry).publish().await.is_err() {
1412/// lane.abort().await.ok();
1413/// return HandlerOutcome::retry();
1414/// }
1415/// }
1416/// if lane.commit().await.is_err() {
1417/// return HandlerOutcome::retry();
1418/// }
1419/// HandlerOutcome::ack()
1420/// }
1421/// ```
1422#[diagnostic::on_unimplemented(
1423 message = "`{Self}` does not hand out per-partition transactional publishers",
1424 note = "for an `Out<impl PartitionLanes, _>` slot, attach a \
1425 `KafkaPublish::default().transactional_id(..).per_partition()` policy"
1426)]
1427pub trait PartitionLanes: Send + Sync {
1428 /// The transactional publisher a lane is handed.
1429 ///
1430 /// An associated type rather than the concrete [`KafkaTransactionalPublisher`]: a handler
1431 /// bounds its slot `Out<impl PartitionLanes>` on the capability, and the policy attached at
1432 /// the mount site decides the publisher.
1433 type Publisher: TransactionalPublisher<Error = KafkaError>;
1434
1435 /// The publisher owning `partition`'s transactional id.
1436 ///
1437 /// # Errors
1438 ///
1439 /// See [`TransactionalPartitions::for_partition`].
1440 fn for_partition(
1441 &self,
1442 partition: i32,
1443 ) -> impl Future<Output = Result<Self::Publisher, KafkaError>> + Send;
1444}
1445
1446impl PartitionLanes for TransactionalPartitions {
1447 type Publisher = KafkaTransactionalPublisher;
1448
1449 fn for_partition(
1450 &self,
1451 partition: i32,
1452 ) -> impl Future<Output = Result<KafkaTransactionalPublisher, KafkaError>> + Send {
1453 Self::for_partition(self, partition)
1454 }
1455}
1456
1457// The arena entry a handler's `Out` parameter binds. `Slot`'s `Deref` already routes method
1458// calls to the wired value, but only this impl lets a body hand the entry to a function generic
1459// over the capability (`fn issue<L: PartitionLanes>(lanes: &L, ..)`), which is the whole point of
1460// bounding the slot with a trait instead of a concrete type. Both paths reach the same unwrapped
1461// value, which is what puts a lane's publishes outside the slot's test capture (see the trait's
1462// documentation).
1463impl<M, L, EncodeCodec, Pipe, Body> PartitionLanes for Slot<M, L, EncodeCodec, Pipe, Body>
1464where
1465 L: PartitionLanes,
1466 EncodeCodec: Send + Sync,
1467 // The entry's publish path, over the wired value whose per-record settings its transforms
1468 // write. A lane's traffic never travels it (it leaves through the unwrapped value), but
1469 // naming the bound the mount site's entry already carries keeps this impl on exactly the
1470 // slots the runtime builds.
1471 Pipe: OutPipeline<L>,
1472{
1473 type Publisher = L::Publisher;
1474
1475 fn for_partition(
1476 &self,
1477 partition: i32,
1478 ) -> impl Future<Output = Result<Self::Publisher, KafkaError>> + Send {
1479 (**self).for_partition(partition)
1480 }
1481}
1482
1483#[cfg(test)]
1484mod tests {
1485 use super::*;
1486
1487 #[test]
1488 fn transactional_id_is_a_type_transition() {
1489 let plain = KafkaPublish::default().queue_timeout(Duration::from_secs(1));
1490 let transactional = plain.transactional_id("svc-1");
1491 assert_eq!(transactional.id(), "svc-1");
1492 assert_eq!(transactional.queue_timeout, plain.queue_timeout_setting());
1493 assert_eq!(
1494 transactional.transaction_timeout,
1495 DEFAULT_TRANSACTION_TIMEOUT
1496 );
1497 }
1498
1499 #[test]
1500 fn per_partition_derives_ids_from_the_base() {
1501 let policy = KafkaPublish::default()
1502 .transactional_id("svc-1")
1503 .transaction_timeout(Duration::from_secs(5));
1504 let derived = policy.with_id(format!("{}-p3", policy.id()));
1505 assert_eq!(derived.id(), "svc-1-p3");
1506 assert_eq!(derived.transaction_timeout, Duration::from_secs(5));
1507 }
1508}