ruststream_rdkafka/publisher.rs
1//! The publish policies and the live publishers they pair into, transactions included.
2
3use std::collections::HashMap;
4use std::fmt;
5use std::sync::{Arc, Mutex};
6use std::time::Duration;
7
8use rdkafka::TopicPartitionList;
9use rdkafka::consumer::ConsumerGroupMetadata;
10use rdkafka::producer::{FutureProducer, FutureRecord, Producer as _};
11use rdkafka::util::Timeout;
12use ruststream::{
13 DefaultPublish, OutgoingMessage, PairError, PublishPolicy, Publisher, TransactionalPublisher,
14};
15use tokio::sync::OnceCell;
16use tokio::task;
17
18use crate::broker::{ConnState, ConnectedKafkaBroker, EarlyConn};
19use crate::convert;
20use crate::error::KafkaError;
21
22const DEFAULT_TRANSACTION_TIMEOUT: Duration = Duration::from_secs(30);
23
24/// The publish policy of [`KafkaPublisher`]: pure declaration, no connection, no publish
25/// surface.
26///
27/// Constructible anywhere - in a router definition, in configuration, before startup - because
28/// it holds nothing but options. The runtime pairs it with the connected broker at startup (or
29/// [`ConnectedKafkaBroker::publisher`] does it by hand), and only the resulting
30/// [`KafkaPublisher`] can publish.
31///
32/// [`transactional_id`](Self::transactional_id) is a type transition, not a flag: it yields a
33/// [`KafkaTransactionalPublish`], so a plain publisher carries no transactional surface at all.
34///
35/// # Examples
36///
37/// ```
38/// use std::time::Duration;
39///
40/// use ruststream_rdkafka::KafkaPublish;
41///
42/// let policy = KafkaPublish::default().queue_timeout(Duration::from_secs(5));
43/// # let _ = policy;
44/// ```
45#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
46#[must_use]
47pub struct KafkaPublish {
48 queue_timeout: Option<Duration>,
49}
50
51impl KafkaPublish {
52 /// How long a publish may wait for space when librdkafka's local queue is full, before
53 /// failing with a queue-full error. Without it a publish waits for space indefinitely,
54 /// which is the natural back-pressure behavior.
55 pub const fn queue_timeout(mut self, timeout: Duration) -> Self {
56 self.queue_timeout = Some(timeout);
57 self
58 }
59
60 /// Turns this into a transactional publish policy fenced by `id` (Kafka's
61 /// `transactional.id`).
62 ///
63 /// The id must be stable and unique per concurrent producer: Kafka uses it to fence
64 /// zombies, so two live producers sharing an id abort each other. Pair distinct policies
65 /// for concurrent transactional flows, or take one publisher per source partition with
66 /// [`KafkaTransactionalPublish::per_partition`].
67 ///
68 /// # Examples
69 ///
70 /// ```
71 /// use ruststream_rdkafka::KafkaPublish;
72 ///
73 /// let policy = KafkaPublish::default().transactional_id("orders-svc-1");
74 /// # let _ = policy;
75 /// ```
76 pub fn transactional_id(self, id: impl Into<String>) -> KafkaTransactionalPublish {
77 KafkaTransactionalPublish {
78 queue_timeout: self.queue_timeout,
79 id: id.into(),
80 transaction_timeout: DEFAULT_TRANSACTION_TIMEOUT,
81 }
82 }
83
84 pub(crate) const fn queue_timeout_setting(self) -> Option<Duration> {
85 self.queue_timeout
86 }
87}
88
89impl PublishPolicy<ConnectedKafkaBroker> for KafkaPublish {
90 type Live = KafkaPublisher;
91
92 async fn pair(self, connected: &ConnectedKafkaBroker) -> Result<Self::Live, PairError> {
93 Ok(connected.publisher(self))
94 }
95}
96
97impl DefaultPublish for ConnectedKafkaBroker {
98 type Policy = KafkaPublish;
99}
100
101/// A live producer handle on the broker's shared producer.
102///
103/// [`OutgoingMessage::name`] is the destination topic. A
104/// [`PARTITION_KEY_HEADER`](crate::PARTITION_KEY_HEADER) header becomes the record's native key,
105/// so Kafka routes messages that share a key to the same partition; without it the configured
106/// partitioner picks one.
107///
108/// Each publish awaits the broker's delivery report, so an `Ok` means the cluster accepted the
109/// record (durability then depends on the producer's `acks` setting, configurable through
110/// [`KafkaBroker::producer_config`](crate::KafkaBroker::producer_config)).
111///
112/// Exists only from a connected broker, so it never sees a "not connected" state; it does alias
113/// that connection and may outlive it, so after the broker shuts down every publish reports
114/// [`KafkaError::Closed`]. Cheap to clone.
115#[derive(Clone)]
116pub struct KafkaPublisher {
117 state: Arc<ConnState>,
118 queue_timeout: Option<Duration>,
119}
120
121impl fmt::Debug for KafkaPublisher {
122 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
123 f.debug_struct("KafkaPublisher")
124 .field("queue_timeout", &self.queue_timeout)
125 .finish_non_exhaustive()
126 }
127}
128
129impl KafkaPublisher {
130 pub(crate) const fn new(state: Arc<ConnState>, queue_timeout: Option<Duration>) -> Self {
131 Self {
132 state,
133 queue_timeout,
134 }
135 }
136}
137
138/// Sends one record through `producer` and awaits its delivery report.
139async fn send_via(
140 producer: &FutureProducer,
141 queue_timeout: Option<Duration>,
142 msg: OutgoingMessage<'_>,
143) -> Result<(), KafkaError> {
144 let parts = convert::headers_for_publish(msg.headers())?;
145 let mut record = FutureRecord::<[u8], [u8]>::to(msg.name()).payload(msg.payload());
146 if let Some(key) = &parts.key {
147 record = record.key(key.as_ref());
148 }
149 if let Some(partition) = parts.partition {
150 // An explicit partition wins over the partitioner and the record key.
151 record = record.partition(partition);
152 }
153 if let Some(headers) = parts.headers {
154 record = record.headers(headers);
155 }
156 let queue_timeout = queue_timeout.map_or(Timeout::Never, Timeout::After);
157 producer
158 .send(record, queue_timeout)
159 .await
160 .map(|_delivery| ())
161 .map_err(|(err, _record)| KafkaError::publish(err))
162}
163
164impl Publisher for KafkaPublisher {
165 type Error = KafkaError;
166
167 /// Publishes `msg` to the topic named by [`OutgoingMessage::name`] and awaits the delivery
168 /// report.
169 ///
170 /// # Errors
171 ///
172 /// Returns [`KafkaError::Closed`] once the connection this handle aliases has been shut
173 /// down, and [`KafkaError::Publish`] when the cluster rejects the record or the delivery
174 /// times out (librdkafka's `message.timeout.ms`).
175 ///
176 /// # Cancel safety
177 ///
178 /// Not cancel safe: dropping the future may leave the record in flight, delivered or not.
179 async fn publish(&self, msg: OutgoingMessage<'_>) -> Result<(), Self::Error> {
180 self.state.ensure_open(msg.name())?;
181 send_via(self.state.producer(), self.queue_timeout, msg).await
182 }
183}
184
185/// A publisher minted from an unconnected [`KafkaBroker`](crate::KafkaBroker), for the one
186/// wiring that cannot take a policy.
187///
188/// [`BrokerScope::retry_via`](ruststream::runtime::BrokerScope::retry_via) - the deferred
189/// republish behind `retry_after`, which Kafka relies on because it has no native delayed
190/// redelivery - is configured while the app builder runs and takes a live [`Publisher`], since
191/// a publisher that cannot send would be a lie. This type is the sanctioned exception that
192/// makes the pair work on a lazy-connect broker: it holds the cell
193/// [`Broker::connect`](ruststream::Broker::connect) fills, not a connection. The policy path
194/// ([`KafkaPublish`] and its transitions) is untouched and stays connection-free by
195/// construction.
196///
197/// Its two runtime checks are the aliasing rule the broker contract deliberately keeps
198/// dynamic: a handle that predates the connection, or outlives it, must surface an error rather
199/// than silently succeed. Publishing before `connect` reports [`KafkaError::NotConnected`];
200/// publishing after the connected broker shut down reports [`KafkaError::Closed`].
201///
202/// # Examples
203///
204/// ```no_run
205/// use ruststream_rdkafka::KafkaBroker;
206///
207/// let broker = KafkaBroker::new(["localhost:9092"]);
208/// let retries = broker.retry_publisher();
209/// // ... `b.retry_via(retries)` while the app builder runs.
210/// # let _ = retries;
211/// ```
212#[derive(Clone)]
213pub struct KafkaRetryPublisher {
214 conn: EarlyConn,
215}
216
217impl fmt::Debug for KafkaRetryPublisher {
218 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
219 f.debug_struct("KafkaRetryPublisher")
220 .field("connected", &self.conn.get().is_some())
221 .finish_non_exhaustive()
222 }
223}
224
225impl KafkaRetryPublisher {
226 pub(crate) const fn new(conn: EarlyConn) -> Self {
227 Self { conn }
228 }
229}
230
231impl Publisher for KafkaRetryPublisher {
232 type Error = KafkaError;
233
234 /// Publishes `msg` through the broker's shared producer and awaits the delivery report.
235 ///
236 /// # Errors
237 ///
238 /// Returns [`KafkaError::NotConnected`] before the broker connects,
239 /// [`KafkaError::Closed`] once it has shut down, and [`KafkaError::Publish`] when the
240 /// cluster rejects the record or the delivery times out.
241 ///
242 /// # Cancel safety
243 ///
244 /// Not cancel safe: dropping the future may leave the record in flight, delivered or not.
245 async fn publish(&self, msg: OutgoingMessage<'_>) -> Result<(), Self::Error> {
246 let state = self.conn.get().ok_or_else(|| KafkaError::NotConnected {
247 topic: msg.name().to_owned(),
248 })?;
249 state.ensure_open(msg.name())?;
250 send_via(state.producer(), None, msg).await
251 }
252}
253
254/// The publish policy of [`KafkaTransactionalPublisher`]: the transactional mode as its own
255/// type, reached from [`KafkaPublish::transactional_id`].
256///
257/// Pairing it is where Kafka does the real work: the transactional producer is created and its
258/// transactions initialized (the call that fences earlier producers with the same id), so a
259/// misconfigured transactional id fails at startup rather than at the first transaction.
260///
261/// # Examples
262///
263/// ```
264/// use std::time::Duration;
265///
266/// use ruststream_rdkafka::KafkaPublish;
267///
268/// let policy = KafkaPublish::default()
269/// .transactional_id("orders-svc-1")
270/// .transaction_timeout(Duration::from_secs(10));
271/// # let _ = policy;
272/// ```
273#[derive(Debug, Clone, PartialEq, Eq)]
274#[must_use]
275pub struct KafkaTransactionalPublish {
276 queue_timeout: Option<Duration>,
277 id: String,
278 transaction_timeout: Duration,
279}
280
281impl KafkaTransactionalPublish {
282 /// See [`KafkaPublish::queue_timeout`].
283 pub const fn queue_timeout(mut self, timeout: Duration) -> Self {
284 self.queue_timeout = Some(timeout);
285 self
286 }
287
288 /// How long transaction control calls (`init`, `commit`, `abort`) may block before
289 /// reporting failure. Defaults to 30 seconds; this is the call deadline handed to
290 /// librdkafka, not its `transaction.timeout.ms` (reachable through
291 /// [`KafkaBroker::producer_config`](crate::KafkaBroker::producer_config)).
292 pub const fn transaction_timeout(mut self, timeout: Duration) -> Self {
293 self.transaction_timeout = timeout;
294 self
295 }
296
297 /// Turns this into a per-partition policy: the id becomes the base of one transactional id
298 /// per source partition (`"{base}-p{partition}"`), pairing into
299 /// [`TransactionalPartitions`].
300 ///
301 /// # Examples
302 ///
303 /// ```
304 /// use ruststream_rdkafka::KafkaPublish;
305 ///
306 /// let policy = KafkaPublish::default()
307 /// .transactional_id("billing-svc-1")
308 /// .per_partition();
309 /// # let _ = policy;
310 /// ```
311 pub fn per_partition(self) -> KafkaPartitionedPublish {
312 KafkaPartitionedPublish { template: self }
313 }
314
315 /// The transactional id this policy fences with.
316 #[must_use]
317 pub fn id(&self) -> &str {
318 &self.id
319 }
320
321 fn with_id(&self, id: String) -> Self {
322 Self {
323 queue_timeout: self.queue_timeout,
324 id,
325 transaction_timeout: self.transaction_timeout,
326 }
327 }
328}
329
330impl PublishPolicy<ConnectedKafkaBroker> for KafkaTransactionalPublish {
331 type Live = KafkaTransactionalPublisher;
332
333 async fn pair(self, connected: &ConnectedKafkaBroker) -> Result<Self::Live, PairError> {
334 connected
335 .transactional_publisher(self)
336 .await
337 .map_err(PairError::new)
338 }
339}
340
341impl ConnectedKafkaBroker {
342 /// A live transactional publisher: creates the transactional producer from the broker's
343 /// resolved producer configuration and initializes its transactions.
344 ///
345 /// Async because the initialization is real work (it fences earlier producers holding the
346 /// same transactional id); it runs once, when the publisher comes alive.
347 ///
348 /// # Errors
349 ///
350 /// Returns [`KafkaError::Closed`] once the connection has been shut down and
351 /// [`KafkaError::Publish`] when the producer cannot be created or the initialization fails
352 /// within the policy's [`transaction_timeout`](KafkaTransactionalPublish::transaction_timeout).
353 ///
354 /// # Examples
355 ///
356 /// ```no_run
357 /// use ruststream::Broker;
358 /// use ruststream_rdkafka::{KafkaBroker, KafkaPublish};
359 ///
360 /// # async fn demo() -> Result<(), ruststream_rdkafka::KafkaError> {
361 /// let connected = KafkaBroker::new(["localhost:9092"]).connect().await?;
362 /// let publisher = connected
363 /// .transactional_publisher(KafkaPublish::default().transactional_id("orders-svc-1"))
364 /// .await?;
365 /// # let _ = publisher;
366 /// # Ok(())
367 /// # }
368 /// ```
369 pub async fn transactional_publisher(
370 &self,
371 policy: KafkaTransactionalPublish,
372 ) -> Result<KafkaTransactionalPublisher, KafkaError> {
373 open_transactional(self.state(), &policy).await
374 }
375}
376
377/// Creates and initializes one transactional producer for `policy`.
378pub(crate) async fn open_transactional(
379 state: &Arc<ConnState>,
380 policy: &KafkaTransactionalPublish,
381) -> Result<KafkaTransactionalPublisher, KafkaError> {
382 state.ensure_open(policy.id())?;
383 let mut config = state.producer_config().clone();
384 config.set("transactional.id", policy.id());
385 let producer: FutureProducer = config.create().map_err(KafkaError::publish)?;
386 // init_transactions blocks (it fences earlier producers with this id), so it runs on the
387 // blocking pool.
388 let init = producer.clone();
389 let timeout = policy.transaction_timeout;
390 task::spawn_blocking(move || init.init_transactions(timeout))
391 .await
392 .map_err(|err| KafkaError::Publish(Box::new(err)))?
393 .map_err(KafkaError::publish)?;
394 Ok(KafkaTransactionalPublisher {
395 inner: Arc::new(TxInner {
396 state: Arc::clone(state),
397 producer,
398 queue_timeout: policy.queue_timeout,
399 timeout,
400 id: policy.id.clone(),
401 open: Mutex::new(false),
402 }),
403 })
404}
405
406struct TxInner {
407 state: Arc<ConnState>,
408 producer: FutureProducer,
409 queue_timeout: Option<Duration>,
410 timeout: Duration,
411 id: String,
412 /// Whether a transaction is currently open. Interleaving `publish` with
413 /// `begin_transaction`/`commit` from concurrent tasks is not supported: which side of the
414 /// transaction boundary a concurrent publish lands on would be a race either way.
415 open: Mutex<bool>,
416}
417
418/// A live publisher that produces inside Kafka transactions.
419///
420/// Records between `begin_transaction` and `commit` become visible atomically (readers on
421/// Kafka's default `read_committed` isolation see all of them or none); `abort` discards them
422/// broker-side.
423///
424/// Its transactional producer is created and initialized when the
425/// [`KafkaTransactionalPublish`] policy pairs, so nothing is lazy here: the handle is fenced
426/// from the moment it exists. Clones share one producer and one transaction state.
427#[derive(Clone)]
428pub struct KafkaTransactionalPublisher {
429 inner: Arc<TxInner>,
430}
431
432impl fmt::Debug for KafkaTransactionalPublisher {
433 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
434 f.debug_struct("KafkaTransactionalPublisher")
435 .field("id", &self.inner.id)
436 .field("timeout", &self.inner.timeout)
437 .finish_non_exhaustive()
438 }
439}
440
441impl KafkaTransactionalPublisher {
442 /// The transactional id fencing this publisher.
443 #[must_use]
444 pub fn id(&self) -> &str {
445 &self.inner.id
446 }
447
448 pub(crate) fn deadline(&self) -> Duration {
449 self.inner.timeout
450 }
451
452 pub(crate) fn state(&self) -> &Arc<ConnState> {
453 &self.inner.state
454 }
455
456 fn is_open(&self) -> bool {
457 *self
458 .inner
459 .open
460 .lock()
461 .expect("transaction state mutex poisoned")
462 }
463
464 fn set_open(&self, open: bool) {
465 *self
466 .inner
467 .open
468 .lock()
469 .expect("transaction state mutex poisoned") = open;
470 }
471
472 fn no_transaction(&self) -> KafkaError {
473 KafkaError::NoTransaction {
474 id: self.inner.id.clone(),
475 }
476 }
477
478 /// Adds consumed source offsets (and their group's metadata) to the open transaction, so
479 /// they commit atomically with the records published into it. The EOS pipeline's commit
480 /// path; must run between `begin_transaction` and `commit`.
481 pub(crate) async fn send_offsets(
482 &self,
483 offsets: TopicPartitionList,
484 metadata: ConsumerGroupMetadata,
485 ) -> Result<(), KafkaError> {
486 self.inner.state.ensure_open(&self.inner.id)?;
487 if !self.is_open() {
488 return Err(self.no_transaction());
489 }
490 let producer = self.inner.producer.clone();
491 let timeout = self.inner.timeout;
492 task::spawn_blocking(move || {
493 producer.send_offsets_to_transaction(&offsets, &metadata, timeout)
494 })
495 .await
496 .map_err(|err| KafkaError::Publish(Box::new(err)))?
497 .map_err(KafkaError::publish)
498 }
499}
500
501impl Publisher for KafkaTransactionalPublisher {
502 type Error = KafkaError;
503
504 /// Publishes `msg` to the topic named by [`OutgoingMessage::name`]. Inside an open
505 /// transaction the record joins it; otherwise it goes out through the broker's shared plain
506 /// producer.
507 ///
508 /// # Errors
509 ///
510 /// Returns [`KafkaError::Closed`] once the connection this handle aliases has been shut
511 /// down and [`KafkaError::Publish`] when the cluster rejects the record or the delivery
512 /// times out.
513 ///
514 /// # Cancel safety
515 ///
516 /// Not cancel safe: dropping the future may leave the record in flight, delivered or not.
517 async fn publish(&self, msg: OutgoingMessage<'_>) -> Result<(), Self::Error> {
518 self.inner.state.ensure_open(msg.name())?;
519 if self.is_open() {
520 return send_via(&self.inner.producer, self.inner.queue_timeout, msg).await;
521 }
522 send_via(self.inner.state.producer(), self.inner.queue_timeout, msg).await
523 }
524}
525
526impl TransactionalPublisher for KafkaTransactionalPublisher {
527 /// Begins a Kafka transaction.
528 ///
529 /// One producer runs one transaction at a time, so beginning while one is open is an
530 /// error, not a queue: a second begin means two flows share one publisher, and silently
531 /// merging their messages into one transaction would commit one flow's records with the
532 /// other's. Concurrent transactional flows use distinct publishers (see
533 /// [`TransactionalPartitions`]).
534 ///
535 /// # Errors
536 ///
537 /// Returns [`KafkaError::TransactionBusy`] when a transaction is already open on this
538 /// publisher (or a clone sharing it), [`KafkaError::Closed`] after the broker shut down,
539 /// and [`KafkaError::Publish`] when the begin call fails.
540 // The guard intentionally spans the begin call: check-and-begin must be atomic so two
541 // concurrent begins cannot both pass the check.
542 #[allow(clippy::significant_drop_tightening)]
543 async fn begin_transaction(&self) -> Result<(), Self::Error> {
544 self.inner.state.ensure_open(&self.inner.id)?;
545 let mut open = self
546 .inner
547 .open
548 .lock()
549 .expect("transaction state mutex poisoned");
550 if *open {
551 return Err(KafkaError::TransactionBusy {
552 id: self.inner.id.clone(),
553 });
554 }
555 // A rejected begin leaves the open transaction untouched, per the trait contract.
556 self.inner
557 .producer
558 .begin_transaction()
559 .map_err(KafkaError::publish)?;
560 *open = true;
561 Ok(())
562 }
563
564 /// Commits the open transaction, making its records visible atomically.
565 ///
566 /// # Errors
567 ///
568 /// Returns [`KafkaError::NoTransaction`] when no transaction is open on this publisher,
569 /// [`KafkaError::Closed`] after the broker shut down, and [`KafkaError::Publish`] when the
570 /// commit fails. librdkafka distinguishes retriable failures from ones requiring an abort;
571 /// after an error the transaction's state is unresolved, so treat the publisher as needing
572 /// an [`abort`](TransactionalPublisher::abort) or replacement.
573 async fn commit(&self) -> Result<(), Self::Error> {
574 self.inner.state.ensure_open(&self.inner.id)?;
575 if !self.is_open() {
576 return Err(self.no_transaction());
577 }
578 let producer = self.inner.producer.clone();
579 let timeout = self.inner.timeout;
580 task::spawn_blocking(move || producer.commit_transaction(timeout))
581 .await
582 .map_err(|err| KafkaError::Publish(Box::new(err)))?
583 .map_err(KafkaError::publish)?;
584 self.set_open(false);
585 Ok(())
586 }
587
588 /// Aborts the open transaction, discarding its records broker-side.
589 ///
590 /// # Errors
591 ///
592 /// Returns [`KafkaError::NoTransaction`] when no transaction is open on this publisher,
593 /// [`KafkaError::Closed`] after the broker shut down, and [`KafkaError::Publish`] when the
594 /// abort fails.
595 async fn abort(&self) -> Result<(), Self::Error> {
596 self.inner.state.ensure_open(&self.inner.id)?;
597 if !self.is_open() {
598 return Err(self.no_transaction());
599 }
600 let producer = self.inner.producer.clone();
601 let timeout = self.inner.timeout;
602 let aborted = task::spawn_blocking(move || producer.abort_transaction(timeout))
603 .await
604 .map_err(|err| KafkaError::Publish(Box::new(err)))?;
605 // The transaction is over either way: a failed abort resolves broker-side by its own
606 // timeout, and leaving the handle "open" would wedge it permanently.
607 self.set_open(false);
608 aborted.map_err(KafkaError::publish)
609 }
610}
611
612/// The publish policy of [`TransactionalPartitions`], reached from
613/// [`KafkaTransactionalPublish::per_partition`].
614///
615/// # Examples
616///
617/// ```
618/// use ruststream_rdkafka::KafkaPublish;
619///
620/// let policy = KafkaPublish::default()
621/// .transactional_id("billing-svc-1")
622/// .per_partition();
623/// # let _ = policy;
624/// ```
625#[derive(Debug, Clone, PartialEq, Eq)]
626#[must_use]
627pub struct KafkaPartitionedPublish {
628 template: KafkaTransactionalPublish,
629}
630
631impl PublishPolicy<ConnectedKafkaBroker> for KafkaPartitionedPublish {
632 type Live = TransactionalPartitions;
633
634 async fn pair(self, connected: &ConnectedKafkaBroker) -> Result<Self::Live, PairError> {
635 Ok(TransactionalPartitions {
636 inner: Arc::new(PartitionsInner {
637 state: Arc::clone(connected.state()),
638 template: self.template,
639 publishers: Mutex::new(HashMap::new()),
640 }),
641 })
642 }
643}
644
645/// Transactional publishers, one per source partition, materialized on first use.
646///
647/// Kafka permits one open transaction per producer and one live producer per transactional id
648/// (initializing a second fences the first), so concurrent transactional handlers need one
649/// producer each. The source partition is the natural scope: under the default
650/// [`LaneKey::Partition`](crate::LaneKey::Partition) worker pool a partition's deliveries
651/// process serially on one lane, so a publisher per partition gives every lane an independent
652/// transaction with no coordination. The id set (`"{base}-p{partition}"`) follows the topic's
653/// partitions rather than the worker count: changing `workers(n)` neither changes the ids nor
654/// weakens zombie fencing - the scheme Kafka Streams uses for its per-task producers.
655///
656/// Not for [`LaneKey::RecordKey`](crate::LaneKey::RecordKey) pools: record-key lanes spread
657/// one partition across lanes, so two lanes would share a partition's publisher and collide
658/// on its single transaction ([`KafkaError::TransactionBusy`]).
659///
660/// Clones share the cache, so one injected handle serves every handler invocation.
661#[derive(Clone)]
662pub struct TransactionalPartitions {
663 inner: Arc<PartitionsInner>,
664}
665
666struct PartitionsInner {
667 state: Arc<ConnState>,
668 template: KafkaTransactionalPublish,
669 /// The per-partition publishers. The set of partitions is only known as deliveries arrive,
670 /// so materialization stays lazy here (a cell per partition, so two lanes racing the same
671 /// partition initialize one producer, not two).
672 publishers: Mutex<HashMap<i32, Arc<OnceCell<KafkaTransactionalPublisher>>>>,
673}
674
675impl fmt::Debug for TransactionalPartitions {
676 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
677 f.debug_struct("TransactionalPartitions")
678 .field("id_base", &self.inner.template.id())
679 .finish_non_exhaustive()
680 }
681}
682
683impl TransactionalPartitions {
684 /// The publisher owning `partition`'s transactional id, created and initialized on first
685 /// use.
686 ///
687 /// `partition` is the delivery's source partition (`KafkaContext`'s `Partition` field in a
688 /// handler); passing anything else still works but forfeits the serialization argument
689 /// that makes the per-partition scope safe.
690 ///
691 /// # Errors
692 ///
693 /// Returns [`KafkaError::Closed`] once the broker has shut down and
694 /// [`KafkaError::Publish`] when the partition's producer cannot be created or its
695 /// transactions cannot be initialized.
696 ///
697 /// # Panics
698 ///
699 /// Panics when the internal cache mutex is poisoned, which requires a prior panic while
700 /// materializing a publisher (an invariant violation, not an operational failure).
701 pub async fn for_partition(
702 &self,
703 partition: i32,
704 ) -> Result<KafkaTransactionalPublisher, KafkaError> {
705 let cell = {
706 let mut publishers = self
707 .inner
708 .publishers
709 .lock()
710 .expect("partition publisher cache mutex poisoned");
711 Arc::clone(publishers.entry(partition).or_default())
712 };
713 let policy = self
714 .inner
715 .template
716 .with_id(format!("{}-p{partition}", self.inner.template.id()));
717 cell.get_or_try_init(|| open_transactional(&self.inner.state, &policy))
718 .await
719 .cloned()
720 }
721}
722
723#[cfg(test)]
724mod tests {
725 use crate::broker::KafkaBroker;
726
727 use super::*;
728
729 #[tokio::test]
730 async fn the_early_publisher_errors_before_connect() {
731 // No I/O anywhere: the cell is simply still empty.
732 let publisher = KafkaBroker::new(["localhost:9092"]).retry_publisher();
733 let err = publisher
734 .publish(OutgoingMessage::new("orders", b"deferred".as_slice()))
735 .await
736 .expect_err("publishing before connect must error");
737 assert!(
738 matches!(&err, KafkaError::NotConnected { topic } if topic == "orders"),
739 "the error must name the topic it could not reach, got: {err}",
740 );
741 }
742
743 #[test]
744 fn transactional_id_is_a_type_transition() {
745 let plain = KafkaPublish::default().queue_timeout(Duration::from_secs(1));
746 let transactional = plain.transactional_id("svc-1");
747 assert_eq!(transactional.id(), "svc-1");
748 assert_eq!(transactional.queue_timeout, plain.queue_timeout_setting());
749 assert_eq!(
750 transactional.transaction_timeout,
751 DEFAULT_TRANSACTION_TIMEOUT
752 );
753 }
754
755 #[test]
756 fn per_partition_derives_ids_from_the_base() {
757 let policy = KafkaPublish::default()
758 .transactional_id("svc-1")
759 .transaction_timeout(Duration::from_secs(5));
760 let derived = policy.with_id(format!("{}-p3", policy.id()));
761 assert_eq!(derived.id(), "svc-1-p3");
762 assert_eq!(derived.transaction_timeout, Duration::from_secs(5));
763 }
764}