Skip to main content

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(&registry));
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}