ruststream-rdkafka 0.7.1

//! The publish policies and the live publishers they pair into, transactions included.

// Without the `testing` feature a transport enum has one variant, so a `match` on it has a single
// arm; the matches stay so the in-process arm has its place when the feature is on.
#![cfg_attr(
    not(feature = "testing"),
    allow(clippy::infallible_destructuring_match)
)]

use std::collections::HashMap;
use std::fmt;
use std::future::{Future, ready};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use rdkafka::TopicPartitionList;
use rdkafka::consumer::ConsumerGroupMetadata;
use rdkafka::producer::{FutureProducer, FutureRecord, Producer as _};
use rdkafka::util::Timeout;
#[cfg(feature = "asyncapi")]
use ruststream::asyncapi::Bindings;
use ruststream::runtime::{OutPipeline, PublishBuilder, PublishSink, Slot};
use ruststream::{
    DefaultPublish, Lend, OutgoingMessage, PairError, PublishPolicy, Publisher,
    TransactionalPublisher,
};
use tokio::sync::OnceCell;
use tokio::task;

#[cfg(feature = "testing")]
use rdkafka::types::RDKafkaType;
#[cfg(feature = "testing")]
use ruststream::Str;

use crate::broker::{ConnState, ConnectedKafkaBroker};
use crate::convert;
use crate::error::KafkaError;
#[cfg(feature = "testing")]
use crate::in_process::{Cluster, ProducerId, create_offline};

const DEFAULT_TRANSACTION_TIMEOUT: Duration = Duration::from_secs(30);

/// What one Kafka record carries beyond its topic, payload and headers.
///
/// A publish builder writes these through the steps of [`KafkaPublishSteps`]; nothing else needs
/// to name the type, except an assertion reading back what a publish carried
/// (`tb.out::<Marker>().with_options(&KafkaOptions::default().partition(3))`).
///
/// The record key is deliberately absent: it is the core's [`Partitioned`](ruststream::Partitioned)
/// contract, carried by the [`PARTITION_KEY_HEADER`](crate::PARTITION_KEY_HEADER) header so that
/// broker-agnostic code can set it, and Kafka maps it onto the native key.
///
/// # Examples
///
/// A test reading back the partition a handler's publish asked for:
///
/// ```
/// # #[cfg(all(feature = "json", feature = "testing"))]
/// # mod demo {
/// use std::error::Error;
///
/// use ruststream::testing::TestApp;
/// use ruststream_rdkafka::prelude::*;
/// use serde::{Deserialize, Serialize};
///
/// #[derive(Serialize, Deserialize, Outgoing)]
/// struct Order {
///     id: u64,
/// }
///
/// #[derive(Serialize, Outgoing)]
/// #[outgoing(name = "audit")]
/// struct Audit {
///     id: u64,
/// }
///
/// #[derive(OutSlot)]
/// #[publishes(Audit)]
/// struct Journal;
///
/// #[subscriber("orders")]
/// async fn record(
///     order: &Order,
///     Out(journal): Out<impl Publisher<Options = KafkaOptions>, Journal>,
/// ) -> HandlerOutcome {
///     let sent = journal.message(&Audit { id: order.id }).partition(3).publish().await;
///     if sent.is_err() {
///         return HandlerOutcome::retry();
///     }
///     HandlerOutcome::ack()
/// }
///
/// pub fn app() -> RustStream {
///     RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(
///         KafkaBroker::new(["kafka:9092"]).default_group("orders-svc"),
///         |b| {
///             b.include(record).out(Journal, Publish::default()).build();
///         },
///     )
/// }
///
/// pub async fn audits_go_to_partition_three() -> Result<(), Box<dyn Error + Send + Sync>> {
///     let tb = TestApp::start(app()).await?;
///
///     tb.broker::<KafkaBroker>()
///         .message(&Order { id: 7 })
///         .to("orders")
///         .publish()
///         .await?;
///
///     tb.out::<Journal>()
///         .assert_called(1)
///         .with_options(&KafkaOptions::default().partition(3));
///     tb.shutdown().await?;
///     Ok(())
/// }
/// # }
/// # #[cfg(all(feature = "json", feature = "testing"))]
/// # #[tokio::main(flavor = "multi_thread", worker_threads = 2)]
/// # async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
/// #     demo::audits_go_to_partition_three().await
/// # }
/// # #[cfg(not(all(feature = "json", feature = "testing")))]
/// # fn main() {}
/// ```
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[must_use]
pub struct KafkaOptions {
    partition: Option<i32>,
}

impl KafkaOptions {
    /// Produces the record to exactly this partition, ahead of the partitioner and the record
    /// key.
    pub const fn partition(mut self, partition: i32) -> Self {
        self.partition = Some(partition);
        self
    }

    /// The partition this record is pinned to, or `None` when the partitioner places it.
    #[must_use]
    pub const fn partition_setting(self) -> Option<i32> {
        self.partition
    }
}

/// The per-message steps this crate adds to the publish builder.
///
/// A handler body that names one imports this crate's prelude and bounds its slot on the options
/// type: `Out<impl Publisher<Options = KafkaOptions>, Marker>`. The bound on the sink is what
/// keeps these steps off a builder over another broker's publisher.
///
/// # Examples
///
/// ```
/// # #[cfg(feature = "json")]
/// # mod demo {
/// use ruststream_rdkafka::prelude::*;
/// use serde::{Deserialize, Serialize};
///
/// #[derive(Deserialize)]
/// struct Order {
///     id: u64,
/// }
///
/// #[derive(Serialize, Outgoing)]
/// struct Audit {
///     id: u64,
/// }
///
/// #[derive(OutSlot)]
/// #[publishes(Audit)]
/// struct Journal;
///
/// /// The slot is bounded on the options type, which is what puts `partition(..)` on its builder.
/// #[subscriber("orders")]
/// async fn record(
///     order: &Order,
///     Out(journal): Out<impl Publisher<Options = KafkaOptions>, Journal>,
/// ) -> HandlerOutcome {
///     let sent = journal
///         .message(&Audit { id: order.id })
///         .to("audit")
///         .partition(0)
///         .publish()
///         .await;
///     if sent.is_err() {
///         return HandlerOutcome::retry();
///     }
///     HandlerOutcome::ack()
/// }
///
/// fn app() -> RustStream {
///     RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(
///         KafkaBroker::new(["localhost:9092"]).default_group("orders-svc"),
///         |b| {
///             b.include(record).out(Journal, Publish::default()).build();
///         },
///     )
/// }
/// # }
/// # fn main() {}
/// ```
pub trait KafkaPublishSteps {
    /// Produces this one record to `partition`, ahead of the partitioner and the record key.
    #[must_use]
    fn partition(self, partition: i32) -> Self;
}

impl<Sink, Body, Enc, Hdrs, Dest> KafkaPublishSteps for PublishBuilder<Sink, Body, Enc, Hdrs, Dest>
where
    Sink: PublishSink<Options = KafkaOptions>,
{
    fn partition(mut self, partition: i32) -> Self {
        self.options_mut()
            .get_or_insert_with(KafkaOptions::default)
            .partition = Some(partition);
        self
    }
}

/// The publish policy of [`KafkaPublisher`]: pure declaration, no connection, no publish
/// surface.
///
/// Constructible anywhere - in a router definition, in configuration, before startup - because
/// it holds nothing but options. The runtime pairs it with the connected broker at startup (or
/// [`ConnectedKafkaBroker::publisher`] does it by hand), and only the resulting
/// [`KafkaPublisher`] can publish.
///
/// [`transactional_id`](Self::transactional_id) is a type transition, not a flag: it yields a
/// [`KafkaTransactionalPublish`], so a plain publisher carries no transactional surface at all.
///
/// # Examples
///
/// ```
/// use std::time::Duration;
///
/// use ruststream_rdkafka::prelude::*;
/// use serde::{Deserialize, Serialize};
///
/// #[derive(Deserialize)]
/// struct Order {
///     id: u64,
/// }
///
/// #[derive(Serialize, Outgoing)]
/// #[outgoing(name = "confirmations")]
/// struct Confirmation {
///     id: u64,
/// }
///
/// #[subscriber("orders", reply)]
/// async fn confirm(order: &Order) -> Confirmation {
///     Confirmation { id: order.id }
/// }
///
/// fn app() -> RustStream {
///     RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(
///         KafkaBroker::new(["localhost:9092"]).default_group("orders-svc"),
///         |b| {
///             b.include(confirm)
///                 .out_reply(Publish::default().queue_timeout(Duration::from_secs(5)));
///         },
///     )
/// }
/// # let _ = app;
/// ```
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[must_use]
pub struct KafkaPublish {
    queue_timeout: Option<Duration>,
}

impl KafkaPublish {
    /// How long a publish may wait for space when librdkafka's local queue is full, before
    /// failing with a queue-full error. Without it a publish waits for space indefinitely,
    /// which is the natural back-pressure behavior.
    pub const fn queue_timeout(mut self, timeout: Duration) -> Self {
        self.queue_timeout = Some(timeout);
        self
    }

    /// Turns this into a transactional publish policy fenced by `id` (Kafka's
    /// `transactional.id`).
    ///
    /// The id must be stable and unique per concurrent producer: Kafka uses it to fence
    /// zombies, so two live producers sharing an id abort each other. Pair distinct policies
    /// for concurrent transactional flows, or take one publisher per source partition with
    /// [`KafkaTransactionalPublish::per_partition`].
    ///
    /// # Examples
    ///
    /// ```
    /// use ruststream_rdkafka::prelude::*;
    /// # use serde::{Deserialize, Serialize};
    /// # #[derive(Deserialize)]
    /// # struct Refund { order_id: u64, lines: u64 }
    /// # #[derive(Serialize, Outgoing)]
    /// # #[outgoing(name = "refund-lines")]
    /// # struct RefundLine { order_id: u64, line: u64 }
    /// # #[derive(OutSlot)]
    /// # #[publishes(RefundLine)]
    /// # struct Lines;
    /// #
    /// #[subscriber("refunds")]
    /// async fn refund(
    ///     order: &Refund,
    ///     Out(lines): Out<impl TransactionalPublisher, Lines>,
    /// ) -> HandlerOutcome {
    ///     if lines.begin_transaction().await.is_err() {
    ///         return HandlerOutcome::retry();
    ///     }
    ///     for line in 0..order.lines {
    ///         let entry = RefundLine {
    ///             order_id: order.order_id,
    ///             line,
    ///         };
    ///         if lines.message(&entry).publish().await.is_err() {
    ///             lines.abort().await.ok();
    ///             return HandlerOutcome::retry();
    ///         }
    ///     }
    ///     if lines.commit().await.is_err() {
    ///         return HandlerOutcome::retry();
    ///     }
    ///     HandlerOutcome::ack()
    /// }
    ///
    /// fn app() -> RustStream {
    ///     RustStream::new(AppInfo::new("refunds", "0.1.0")).with_broker(
    ///         KafkaBroker::new(["localhost:9092"]).default_group("refunds-svc"),
    ///         |b| {
    ///             b.include(refund)
    ///                 .out(Lines, Publish::default().transactional_id("refunds-svc-1"))
    ///                 .build();
    ///         },
    ///     )
    /// }
    /// # let _ = app;
    /// ```
    pub fn transactional_id(self, id: impl Into<String>) -> KafkaTransactionalPublish {
        KafkaTransactionalPublish {
            queue_timeout: self.queue_timeout,
            id: id.into(),
            transaction_timeout: DEFAULT_TRANSACTION_TIMEOUT,
        }
    }

    /// Frames every message this policy publishes in the Confluent wire format, resolving each
    /// destination topic's subject through `registry`.
    ///
    /// This is what lets a Protobuf handler stay a plain function that returns its reply: the
    /// value writes bare Protobuf through `#[wire(encode = ::prost::Message::encode)]`, and the
    /// publisher this pairs into puts the schema id and the message-index path in front of it. A
    /// value cannot do that itself - `Serialized::wire_bytes` is synchronous and has only
    /// `&self` - while the publish path knows the destination topic, which is what names the
    /// subject.
    ///
    /// Short for [`KafkaFramedPublish::over`](crate::KafkaFramedPublish::over) over the default
    /// producer settings; take that one to keep the settings a configured [`KafkaPublish`]
    /// holds. See [`KafkaFramedPublish`](crate::KafkaFramedPublish) for what it refuses, and for
    /// when it resolves a subject.
    ///
    /// # Examples
    ///
    /// ```
    /// use ruststream_rdkafka::prelude::*;
    /// use ruststream_rdkafka::SchemaRegistry;
    ///
    /// #[derive(Clone, PartialEq, prost::Message, Deserialized)]
    /// #[wire(decode = ruststream_rdkafka::protobuf::decode_confluent)]
    /// struct Order {
    ///     #[prost(int64, tag = "1")]
    ///     id: i64,
    /// }
    ///
    /// #[derive(Clone, PartialEq, prost::Message, Serialized, Outgoing)]
    /// #[wire(encode = ::prost::Message::encode)]
    /// #[outgoing(name = "confirmations")]
    /// struct Confirmation {
    ///     #[prost(int64, tag = "1")]
    ///     id: i64,
    /// }
    ///
    /// #[subscriber("orders", reply)]
    /// async fn confirm(order: &Order) -> Confirmation {
    ///     Confirmation { id: order.id }
    /// }
    ///
    /// fn app() -> RustStream {
    ///     let registry = SchemaRegistry::new("http://localhost:8081");
    ///     RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(
    ///         KafkaBroker::new(["localhost:9092"]).default_group("orders-svc"),
    ///         |b| {
    ///             b.include(confirm).out_reply(Publish::framed(&registry));
    ///         },
    ///     )
    /// }
    /// # let _ = app;
    /// ```
    #[cfg(feature = "protobuf")]
    pub fn framed(registry: &crate::SchemaRegistry) -> crate::protobuf::KafkaFramedPublish {
        crate::protobuf::KafkaFramedPublish::over(Self::default(), registry)
    }

    pub(crate) const fn queue_timeout_setting(self) -> Option<Duration> {
        self.queue_timeout
    }
}

impl PublishPolicy<ConnectedKafkaBroker> for KafkaPublish {
    type Live = KafkaPublisher;

    fn pair(
        self,
        connected: &ConnectedKafkaBroker,
    ) -> impl Future<Output = Result<Self::Live, PairError>> {
        ready(Ok(connected.publisher(self)))
    }

    /// The destination topic of this publish, which is what the channel stands for.
    #[cfg(feature = "asyncapi")]
    fn channel_bindings(&self, channel: &str) -> Bindings {
        crate::bindings::channel(channel)
    }
}

impl DefaultPublish for ConnectedKafkaBroker {
    type Policy = KafkaPublish;
}

/// A live producer handle on the broker's shared producer.
///
/// [`OutgoingMessage::name`] is the destination topic. A
/// [`PARTITION_KEY_HEADER`](crate::PARTITION_KEY_HEADER) header becomes the record's native key,
/// so Kafka routes messages that share a key to the same partition; without it the configured
/// partitioner picks one.
///
/// Each publish awaits the broker's delivery report, so an `Ok` means the cluster accepted the
/// record (durability then depends on the producer's `acks` setting, configurable through
/// [`KafkaBroker::producer_config`](crate::KafkaBroker::producer_config)).
///
/// Exists only from a connected broker, so it never sees a "not connected" state; it does alias
/// that connection and may outlive it, so after the broker shuts down every publish reports
/// [`KafkaError::Closed`]. Cheap to clone.
#[derive(Clone)]
pub struct KafkaPublisher {
    state: Arc<ConnState>,
    queue_timeout: Option<Duration>,
}

impl fmt::Debug for KafkaPublisher {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("KafkaPublisher")
            .field("queue_timeout", &self.queue_timeout)
            .finish_non_exhaustive()
    }
}

impl KafkaPublisher {
    pub(crate) const fn new(state: Arc<ConnState>, queue_timeout: Option<Duration>) -> Self {
        Self {
            state,
            queue_timeout,
        }
    }
}

/// Sends one record through `producer` and awaits its delivery report.
async fn send_via(
    producer: &FutureProducer,
    queue_timeout: Option<Duration>,
    msg: OutgoingMessage<'_>,
    options: Option<&KafkaOptions>,
) -> Result<(), KafkaError> {
    let parts = convert::headers_for_publish(msg.headers());
    let mut record = FutureRecord::<[u8], [u8]>::to(msg.name()).payload(msg.payload());
    if let Some(key) = &parts.key {
        record = record.key(key.as_ref());
    }
    if let Some(partition) = options.and_then(|options| options.partition_setting()) {
        // An explicit partition wins over the partitioner and the record key.
        record = record.partition(partition);
    }
    if let Some(headers) = parts.headers {
        record = record.headers(headers);
    }
    let queue_timeout = queue_timeout.map_or(Timeout::Never, Timeout::After);
    producer
        .send(record, queue_timeout)
        .await
        .map(|_delivery| ())
        .map_err(|(err, _record)| KafkaError::publish(err))
}

impl Publisher for KafkaPublisher {
    // librdkafka's record takes the payload as `&[u8]` and copies it into the producer queue,
    // so this transport reads the bytes and keeps nothing.
    type Payload = Lend;
    type Error = KafkaError;
    type Options = KafkaOptions;

    /// Publishes `msg` to the topic named by [`OutgoingMessage::name`] and awaits the delivery
    /// report.
    ///
    /// # Errors
    ///
    /// Returns [`KafkaError::Closed`] once the connection this handle aliases has been shut
    /// down, and [`KafkaError::Publish`] when the cluster rejects the record or the delivery
    /// times out (librdkafka's `message.timeout.ms`).
    ///
    /// # Cancel safety
    ///
    /// Not cancel safe: dropping the future may leave the record in flight, delivered or not.
    async fn publish(
        &self,
        msg: OutgoingMessage<'_>,
        options: Option<&Self::Options>,
    ) -> Result<(), Self::Error> {
        self.state.ensure_open(msg.name())?;
        #[cfg(feature = "testing")]
        if let Some(cluster) = self.state.in_process() {
            return produce_in_process(cluster, &msg, options, None);
        }
        send_via(self.state.producer()?, self.queue_timeout, msg, options).await
    }
}

/// Hands one record to the in-process cluster, as `send_via` hands it to librdkafka.
#[cfg(feature = "testing")]
fn produce_in_process(
    cluster: &Cluster,
    msg: &OutgoingMessage<'_>,
    options: Option<&KafkaOptions>,
    transaction: Option<&ProducerId>,
) -> Result<(), KafkaError> {
    cluster.produce(
        msg.name(),
        options.and_then(|options| options.partition_setting()),
        msg.payload(),
        msg.headers(),
        transaction,
    )
}

/// The publish policy of [`KafkaTransactionalPublisher`]: the transactional mode as its own
/// type, reached from [`KafkaPublish::transactional_id`].
///
/// Pairing it is where Kafka does the real work: the transactional producer is created and its
/// transactions initialized (the call that fences earlier producers with the same id), so a
/// misconfigured transactional id fails at startup rather than at the first transaction.
///
/// # Examples
///
/// ```
/// use std::time::Duration;
///
/// use ruststream_rdkafka::prelude::*;
/// # use serde::{Deserialize, Serialize};
/// # #[derive(Deserialize)]
/// # struct Refund { order_id: u64, lines: u64 }
/// # #[derive(Serialize, Outgoing)]
/// # #[outgoing(name = "refund-lines")]
/// # struct RefundLine { order_id: u64, line: u64 }
/// # #[derive(OutSlot)]
/// # #[publishes(RefundLine)]
/// # struct Lines;
/// #
/// #[subscriber("refunds")]
/// async fn refund(
///     order: &Refund,
///     Out(lines): Out<impl TransactionalPublisher, Lines>,
/// ) -> HandlerOutcome {
///     if lines.begin_transaction().await.is_err() {
///         return HandlerOutcome::retry();
///     }
///     for line in 0..order.lines {
///         let entry = RefundLine {
///             order_id: order.order_id,
///             line,
///         };
///         if lines.message(&entry).publish().await.is_err() {
///             lines.abort().await.ok();
///             return HandlerOutcome::retry();
///         }
///     }
///     if lines.commit().await.is_err() {
///         return HandlerOutcome::retry();
///     }
///     HandlerOutcome::ack()
/// }
///
/// fn app() -> RustStream {
///     let lines = Publish::default()
///         .transactional_id("refunds-svc-1")
///         .transaction_timeout(Duration::from_secs(10));
///     RustStream::new(AppInfo::new("refunds", "0.1.0")).with_broker(
///         KafkaBroker::new(["localhost:9092"]).default_group("refunds-svc"),
///         |b| {
///             b.include(refund).out(Lines, lines).build();
///         },
///     )
/// }
/// # let _ = app;
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
#[must_use]
pub struct KafkaTransactionalPublish {
    queue_timeout: Option<Duration>,
    id: String,
    transaction_timeout: Duration,
}

impl KafkaTransactionalPublish {
    /// See [`KafkaPublish::queue_timeout`].
    pub const fn queue_timeout(mut self, timeout: Duration) -> Self {
        self.queue_timeout = Some(timeout);
        self
    }

    /// How long transaction control calls (`init`, `commit`, `abort`) may block before
    /// reporting failure. Defaults to 30 seconds; this is the call deadline handed to
    /// librdkafka, not its `transaction.timeout.ms` (reachable through
    /// [`KafkaBroker::producer_config`](crate::KafkaBroker::producer_config)).
    pub const fn transaction_timeout(mut self, timeout: Duration) -> Self {
        self.transaction_timeout = timeout;
        self
    }

    /// Turns this into a per-partition policy: the id becomes the base of one transactional id
    /// per source partition (`"{base}-p{partition}"`), pairing into
    /// [`TransactionalPartitions`].
    ///
    /// # Examples
    ///
    /// ```
    /// use ruststream_rdkafka::prelude::*;
    /// # use serde::{Deserialize, Serialize};
    /// # #[derive(Deserialize)]
    /// # struct Invoice { id: u64, lines: u64 }
    /// # #[derive(Serialize, Outgoing)]
    /// # #[outgoing(name = "invoice-lines")]
    /// # struct InvoiceLine { invoice_id: u64, line: u64 }
    /// # #[subscriber(KafkaTopic::new("invoices").commit(Commit::Tracked), workers(4, by_key))]
    /// # async fn bill(
    /// #     invoice: &Invoice,
    /// #     Ctx(partition): Ctx<Partition>,
    /// #     Out(lanes): Out<impl PartitionLanes>,
    /// # ) -> HandlerOutcome {
    /// #     let Ok(lane) = lanes.for_partition(partition).await else {
    /// #         return HandlerOutcome::retry();
    /// #     };
    /// #     if lane.begin_transaction().await.is_err() {
    /// #         return HandlerOutcome::retry();
    /// #     }
    /// #     for line in 0..invoice.lines {
    /// #         let entry = InvoiceLine { invoice_id: invoice.id, line };
    /// #         if lane.message(&entry).publish().await.is_err() {
    /// #             lane.abort().await.ok();
    /// #             return HandlerOutcome::retry();
    /// #         }
    /// #     }
    /// #     if lane.commit().await.is_err() {
    /// #         return HandlerOutcome::retry();
    /// #     }
    /// #     HandlerOutcome::ack()
    /// # }
    ///
    /// fn app() -> RustStream {
    ///     RustStream::new(AppInfo::new("billing", "0.1.0")).with_broker(
    ///         KafkaBroker::new(["localhost:9092"]).default_group("billing-svc"),
    ///         |b| {
    ///             // One transactional id per source partition: "billing-svc-1-p0", "-p1", ...
    ///             b.include(bill)
    ///                 .out(
    ///                     DefaultSlot,
    ///                     Publish::default().transactional_id("billing-svc-1").per_partition(),
    ///                 )
    ///                 .build();
    ///         },
    ///     )
    /// }
    /// # let _ = app;
    /// ```
    pub fn per_partition(self) -> KafkaPartitionedPublish {
        KafkaPartitionedPublish { template: self }
    }

    /// The transactional id this policy fences with.
    #[must_use]
    pub fn id(&self) -> &str {
        &self.id
    }

    pub(crate) fn with_id(&self, id: String) -> Self {
        Self {
            queue_timeout: self.queue_timeout,
            id,
            transaction_timeout: self.transaction_timeout,
        }
    }
}

impl PublishPolicy<ConnectedKafkaBroker> for KafkaTransactionalPublish {
    type Live = KafkaTransactionalPublisher;

    async fn pair(self, connected: &ConnectedKafkaBroker) -> Result<Self::Live, PairError> {
        connected
            .transactional_publisher(self)
            .await
            .map_err(PairError::new)
    }

    /// The destination topic of this publish, which is what the channel stands for.
    #[cfg(feature = "asyncapi")]
    fn channel_bindings(&self, channel: &str) -> Bindings {
        crate::bindings::channel(channel)
    }
}

impl ConnectedKafkaBroker {
    /// A live transactional publisher: creates the transactional producer from the broker's
    /// resolved producer configuration and initializes its transactions.
    ///
    /// Async because the initialization is real work (it fences earlier producers holding the
    /// same transactional id); it runs once, when the publisher comes alive.
    ///
    /// # Errors
    ///
    /// Returns [`KafkaError::Closed`] once the connection has been shut down and
    /// [`KafkaError::Publish`] when the producer cannot be created or the initialization fails
    /// within the policy's [`transaction_timeout`](KafkaTransactionalPublish::transaction_timeout).
    ///
    /// # Examples
    ///
    /// A migration tool that copies a batch of records all-or-nothing:
    ///
    /// ```no_run
    /// use std::error::Error;
    ///
    /// use ruststream::Broker;
    /// use ruststream_rdkafka::prelude::*;
    /// use serde::Serialize;
    ///
    /// #[derive(Serialize, Outgoing)]
    /// #[outgoing(name = "accounts-v2")]
    /// struct Account {
    ///     id: u64,
    /// }
    ///
    /// async fn migrate(accounts: &[Account]) -> Result<(), Box<dyn Error + Send + Sync>> {
    ///     let connected = KafkaBroker::new(["localhost:9092"]).connect().await?;
    ///     let publisher = connected
    ///         .transactional_publisher(Publish::default().transactional_id("accounts-migration"))
    ///         .await?;
    ///     publisher.begin_transaction().await?;
    ///     for account in accounts {
    ///         if let Err(err) = publisher.message(account).publish().await {
    ///             publisher.abort().await?;
    ///             return Err(err.into());
    ///         }
    ///     }
    ///     publisher.commit().await?;
    ///     Ok(())
    /// }
    /// # let _ = migrate;
    /// ```
    pub async fn transactional_publisher(
        &self,
        policy: KafkaTransactionalPublish,
    ) -> Result<KafkaTransactionalPublisher, KafkaError> {
        open_transactional(self.state(), &policy).await
    }
}

/// Creates and initializes one transactional producer for `policy`.
pub(crate) async fn open_transactional(
    state: &Arc<ConnState>,
    policy: &KafkaTransactionalPublish,
) -> Result<KafkaTransactionalPublisher, KafkaError> {
    state.ensure_open(policy.id())?;
    let mut config = state.producer_config().clone();
    config.set("transactional.id", policy.id());
    #[cfg(feature = "testing")]
    if let Some(cluster) = state.in_process() {
        // The configuration is validated as the live producer's creation would validate it, and
        // the pairing fences every earlier one with this id, as `init_transactions` does.
        create_offline(&config, RDKafkaType::RD_KAFKA_PRODUCER).map_err(KafkaError::publish)?;
        let producer = TxProducer::InProcess {
            cluster: Arc::clone(cluster),
            id: cluster.init_transactions(policy.id()),
        };
        return Ok(KafkaTransactionalPublisher::with(state, producer, policy));
    }
    let producer: FutureProducer = config.create().map_err(KafkaError::publish)?;
    // init_transactions blocks (it fences earlier producers with this id), so it runs on the
    // blocking pool.
    let init = producer.clone();
    let timeout = policy.transaction_timeout;
    task::spawn_blocking(move || init.init_transactions(timeout))
        .await
        .map_err(|err| KafkaError::Publish(Box::new(err)))?
        .map_err(KafkaError::publish)?;
    Ok(KafkaTransactionalPublisher::with(
        state,
        TxProducer::Kafka(producer),
        policy,
    ))
}

/// The producer a transactional publisher runs its transactions on: a librdkafka producer, or,
/// under the `testing` feature, a producer id paired with the in-process cluster.
///
/// Without the feature there is one variant, so the type is the producer itself.
enum TxProducer {
    Kafka(FutureProducer),
    #[cfg(feature = "testing")]
    InProcess {
        cluster: Arc<Cluster>,
        id: ProducerId,
    },
}

#[cfg(not(feature = "testing"))]
const _: () = assert!(size_of::<TxProducer>() == size_of::<FutureProducer>());

struct TxInner {
    state: Arc<ConnState>,
    producer: TxProducer,
    queue_timeout: Option<Duration>,
    timeout: Duration,
    id: String,
    /// Whether a transaction is currently open. Interleaving `publish` with
    /// `begin_transaction`/`commit` from concurrent tasks is not supported: which side of the
    /// transaction boundary a concurrent publish lands on would be a race either way.
    open: Mutex<bool>,
}

/// A live publisher that produces inside Kafka transactions.
///
/// Records between `begin_transaction` and `commit` become visible atomically (readers on
/// Kafka's default `read_committed` isolation see all of them or none); `abort` discards them
/// broker-side.
///
/// Its transactional producer is created and initialized when the
/// [`KafkaTransactionalPublish`] policy pairs, so nothing is lazy here: the handle is fenced
/// from the moment it exists. Clones share one producer and one transaction state.
///
/// Outside an open transaction the handle publishes like a plain one. A second
/// `begin_transaction` reports [`KafkaError::TransactionBusy`], and a `commit` or `abort` with
/// nothing open reports [`KafkaError::NoTransaction`].
///
/// # Examples
///
/// A handler receives the publisher through an `Out` slot bounded by the core's
/// [`TransactionalPublisher`] capability, so its body names nothing of this crate's:
///
/// ```
/// # #[cfg(feature = "json")]
/// # mod demo {
/// use ruststream_rdkafka::prelude::*;
/// use serde::{Deserialize, Serialize};
///
/// #[derive(Deserialize)]
/// struct Refund {
///     order_id: u64,
///     lines: u64,
/// }
///
/// #[derive(Serialize, Outgoing)]
/// #[outgoing(name = "refund-lines")]
/// struct RefundLine {
///     order_id: u64,
///     line: u64,
/// }
///
/// #[derive(OutSlot)]
/// #[publishes(RefundLine)]
/// struct Lines;
///
/// #[subscriber("refunds")]
/// async fn refund(
///     order: &Refund,
///     Out(lines): Out<impl TransactionalPublisher, Lines>,
/// ) -> HandlerOutcome {
///     if lines.begin_transaction().await.is_err() {
///         return HandlerOutcome::retry();
///     }
///     for line in 0..order.lines {
///         let entry = RefundLine {
///             order_id: order.order_id,
///             line,
///         };
///         if lines.message(&entry).publish().await.is_err() {
///             lines.abort().await.ok();
///             return HandlerOutcome::retry();
///         }
///     }
///     if lines.commit().await.is_err() {
///         return HandlerOutcome::retry();
///     }
///     HandlerOutcome::ack()
/// }
///
/// fn app() -> RustStream {
///     RustStream::new(AppInfo::new("refunds", "0.1.0")).with_broker(
///         KafkaBroker::new(["localhost:9092"]).default_group("refunds-svc"),
///         |b| {
///             b.include(refund)
///                 .out(Lines, Publish::default().transactional_id("refunds-svc-1"))
///                 .build();
///         },
///     )
/// }
/// # }
/// # fn main() {}
/// ```
#[derive(Clone)]
pub struct KafkaTransactionalPublisher {
    inner: Arc<TxInner>,
}

impl fmt::Debug for KafkaTransactionalPublisher {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("KafkaTransactionalPublisher")
            .field("id", &self.inner.id)
            .field("timeout", &self.inner.timeout)
            .finish_non_exhaustive()
    }
}

impl KafkaTransactionalPublisher {
    fn with(
        state: &Arc<ConnState>,
        producer: TxProducer,
        policy: &KafkaTransactionalPublish,
    ) -> Self {
        Self {
            inner: Arc::new(TxInner {
                state: Arc::clone(state),
                producer,
                queue_timeout: policy.queue_timeout,
                timeout: policy.transaction_timeout,
                id: policy.id.clone(),
                open: Mutex::new(false),
            }),
        }
    }

    /// The transactional id fencing this publisher.
    #[must_use]
    pub fn id(&self) -> &str {
        &self.inner.id
    }

    pub(crate) fn deadline(&self) -> Duration {
        self.inner.timeout
    }

    pub(crate) fn state(&self) -> &Arc<ConnState> {
        &self.inner.state
    }

    fn is_open(&self) -> bool {
        *self
            .inner
            .open
            .lock()
            .expect("transaction state mutex poisoned")
    }

    fn set_open(&self, open: bool) {
        *self
            .inner
            .open
            .lock()
            .expect("transaction state mutex poisoned") = open;
    }

    fn no_transaction(&self) -> KafkaError {
        KafkaError::NoTransaction {
            id: self.inner.id.clone(),
        }
    }

    /// Adds consumed source offsets (and their group's metadata) to the open transaction, so
    /// they commit atomically with the records published into it. The EOS pipeline's commit
    /// path; must run between `begin_transaction` and `commit`.
    pub(crate) async fn send_offsets(
        &self,
        offsets: TopicPartitionList,
        metadata: ConsumerGroupMetadata,
    ) -> Result<(), KafkaError> {
        self.inner.state.ensure_open(&self.inner.id)?;
        if !self.is_open() {
            return Err(self.no_transaction());
        }
        let producer = match &self.inner.producer {
            TxProducer::Kafka(producer) => producer.clone(),
            #[cfg(feature = "testing")]
            TxProducer::InProcess { .. } => {
                return Err(KafkaError::InvalidOptions(
                    "librdkafka offsets reached an in-process transaction".to_owned(),
                ));
            }
        };
        let timeout = self.inner.timeout;
        task::spawn_blocking(move || {
            producer.send_offsets_to_transaction(&offsets, &metadata, timeout)
        })
        .await
        .map_err(|err| KafkaError::Publish(Box::new(err)))?
        .map_err(KafkaError::publish)
    }
}

#[cfg(feature = "testing")]
impl KafkaTransactionalPublisher {
    /// Adds consumed positions to the open transaction on the in-process cluster, to commit into
    /// `group` with it: the in-process counterpart of [`send_offsets`](Self::send_offsets).
    pub(crate) fn send_offsets_in_process(
        &self,
        group: &str,
        positions: &[((Str, i32), i64)],
    ) -> Result<(), KafkaError> {
        self.inner.state.ensure_open(&self.inner.id)?;
        if !self.is_open() {
            return Err(self.no_transaction());
        }
        match &self.inner.producer {
            TxProducer::InProcess { cluster, id } => cluster.send_offsets(id, group, positions),
            TxProducer::Kafka(_) => Err(KafkaError::InvalidOptions(
                "in-process offsets reached a librdkafka transaction".to_owned(),
            )),
        }
    }
}

impl Publisher for KafkaTransactionalPublisher {
    // The same record as the plain publisher's, inside a transaction.
    type Payload = Lend;
    type Error = KafkaError;
    type Options = KafkaOptions;

    /// Publishes `msg` to the topic named by [`OutgoingMessage::name`]. Inside an open
    /// transaction the record joins it; otherwise it goes out through the broker's shared plain
    /// producer.
    ///
    /// # Errors
    ///
    /// Returns [`KafkaError::Closed`] once the connection this handle aliases has been shut
    /// down and [`KafkaError::Publish`] when the cluster rejects the record or the delivery
    /// times out.
    ///
    /// # Cancel safety
    ///
    /// Not cancel safe: dropping the future may leave the record in flight, delivered or not.
    async fn publish(
        &self,
        msg: OutgoingMessage<'_>,
        options: Option<&Self::Options>,
    ) -> Result<(), Self::Error> {
        self.inner.state.ensure_open(msg.name())?;
        let producer = match &self.inner.producer {
            TxProducer::Kafka(producer) => producer,
            #[cfg(feature = "testing")]
            TxProducer::InProcess { cluster, id } => {
                let transaction = self.is_open().then_some(id);
                return produce_in_process(cluster, &msg, options, transaction);
            }
        };
        if self.is_open() {
            return send_via(producer, self.inner.queue_timeout, msg, options).await;
        }
        send_via(
            self.inner.state.producer()?,
            self.inner.queue_timeout,
            msg,
            options,
        )
        .await
    }
}

impl TransactionalPublisher for KafkaTransactionalPublisher {
    /// Begins a Kafka transaction.
    ///
    /// One producer runs one transaction at a time, so beginning while one is open is an
    /// error, not a queue: a second begin means two flows share one publisher, and silently
    /// merging their messages into one transaction would commit one flow's records with the
    /// other's. Concurrent transactional flows use distinct publishers (see
    /// [`TransactionalPartitions`]).
    ///
    /// # Errors
    ///
    /// Returns [`KafkaError::TransactionBusy`] when a transaction is already open on this
    /// publisher (or a clone sharing it), [`KafkaError::Closed`] after the broker shut down,
    /// and [`KafkaError::Publish`] when the begin call fails.
    // The guard intentionally spans the begin call: check-and-begin must be atomic so two
    // concurrent begins cannot both pass the check.
    fn begin_transaction(&self) -> impl Future<Output = Result<(), Self::Error>> {
        if let Err(err) = self.inner.state.ensure_open(&self.inner.id) {
            return ready(Err(err));
        }
        {
            let mut open = self
                .inner
                .open
                .lock()
                .expect("transaction state mutex poisoned");
            if *open {
                return ready(Err(KafkaError::TransactionBusy {
                    id: self.inner.id.clone(),
                }));
            }
            // A rejected begin leaves the open transaction untouched, per the trait contract.
            let begun = match &self.inner.producer {
                TxProducer::Kafka(producer) => {
                    producer.begin_transaction().map_err(KafkaError::publish)
                }
                #[cfg(feature = "testing")]
                TxProducer::InProcess { cluster, id } => {
                    cluster.begin(id, self.inner.state.runtime())
                }
            };
            if let Err(err) = begun {
                return ready(Err(err));
            }
            *open = true;
        }
        ready(Ok(()))
    }

    /// Commits the open transaction, making its records visible atomically.
    ///
    /// # Errors
    ///
    /// Returns [`KafkaError::NoTransaction`] when no transaction is open on this publisher,
    /// [`KafkaError::Closed`] after the broker shut down, and [`KafkaError::Publish`] when the
    /// commit fails. librdkafka distinguishes retriable failures from ones requiring an abort;
    /// after an error the transaction's state is unresolved, so treat the publisher as needing
    /// an [`abort`](TransactionalPublisher::abort) or replacement.
    async fn commit(&self) -> Result<(), Self::Error> {
        self.inner.state.ensure_open(&self.inner.id)?;
        if !self.is_open() {
            return Err(self.no_transaction());
        }
        let producer = match &self.inner.producer {
            TxProducer::Kafka(producer) => producer.clone(),
            #[cfg(feature = "testing")]
            TxProducer::InProcess { cluster, id } => {
                // As on the live path, only a commit that succeeded closes the handle's
                // transaction; after a failure the caller aborts.
                cluster.commit(id)?;
                self.set_open(false);
                return Ok(());
            }
        };
        let timeout = self.inner.timeout;
        task::spawn_blocking(move || producer.commit_transaction(timeout))
            .await
            .map_err(|err| KafkaError::Publish(Box::new(err)))?
            .map_err(KafkaError::publish)?;
        self.set_open(false);
        Ok(())
    }

    /// Aborts the open transaction, discarding its records broker-side.
    ///
    /// # Errors
    ///
    /// Returns [`KafkaError::NoTransaction`] when no transaction is open on this publisher,
    /// [`KafkaError::Closed`] after the broker shut down, and [`KafkaError::Publish`] when the
    /// abort fails.
    async fn abort(&self) -> Result<(), Self::Error> {
        self.inner.state.ensure_open(&self.inner.id)?;
        if !self.is_open() {
            return Err(self.no_transaction());
        }
        let producer = match &self.inner.producer {
            TxProducer::Kafka(producer) => producer.clone(),
            #[cfg(feature = "testing")]
            TxProducer::InProcess { cluster, id } => {
                let aborted = cluster.abort(id);
                self.set_open(false);
                return aborted;
            }
        };
        let timeout = self.inner.timeout;
        let aborted = task::spawn_blocking(move || producer.abort_transaction(timeout))
            .await
            .map_err(|err| KafkaError::Publish(Box::new(err)))?;
        // The transaction is over either way: a failed abort resolves broker-side by its own
        // timeout, and leaving the handle "open" would wedge it permanently.
        self.set_open(false);
        aborted.map_err(KafkaError::publish)
    }
}

/// The publish policy of [`TransactionalPartitions`], reached from
/// [`KafkaTransactionalPublish::per_partition`].
///
/// # Examples
///
/// ```
/// use ruststream_rdkafka::prelude::*;
/// # use serde::{Deserialize, Serialize};
/// # #[derive(Deserialize)]
/// # struct Invoice { id: u64, lines: u64 }
/// # #[derive(Serialize, Outgoing)]
/// # #[outgoing(name = "invoice-lines")]
/// # struct InvoiceLine { invoice_id: u64, line: u64 }
/// # #[subscriber(KafkaTopic::new("invoices").commit(Commit::Tracked), workers(4, by_key))]
/// # async fn bill(
/// #     invoice: &Invoice,
/// #     Ctx(partition): Ctx<Partition>,
/// #     Out(lanes): Out<impl PartitionLanes>,
/// # ) -> HandlerOutcome {
/// #     let Ok(lane) = lanes.for_partition(partition).await else {
/// #         return HandlerOutcome::retry();
/// #     };
/// #     if lane.begin_transaction().await.is_err() {
/// #         return HandlerOutcome::retry();
/// #     }
/// #     for line in 0..invoice.lines {
/// #         let entry = InvoiceLine { invoice_id: invoice.id, line };
/// #         if lane.message(&entry).publish().await.is_err() {
/// #             lane.abort().await.ok();
/// #             return HandlerOutcome::retry();
/// #         }
/// #     }
/// #     if lane.commit().await.is_err() {
/// #         return HandlerOutcome::retry();
/// #     }
/// #     HandlerOutcome::ack()
/// # }
///
/// fn app() -> RustStream {
///     RustStream::new(AppInfo::new("billing", "0.1.0")).with_broker(
///         KafkaBroker::new(["localhost:9092"]).default_group("billing-svc"),
///         |b| {
///             // One transactional id per source partition: "billing-svc-1-p0", "-p1", ...
///             b.include(bill)
///                 .out(
///                     DefaultSlot,
///                     Publish::default().transactional_id("billing-svc-1").per_partition(),
///                 )
///                 .build();
///         },
///     )
/// }
/// # let _ = app;
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
#[must_use]
pub struct KafkaPartitionedPublish {
    template: KafkaTransactionalPublish,
}

impl PublishPolicy<ConnectedKafkaBroker> for KafkaPartitionedPublish {
    type Live = TransactionalPartitions;

    fn pair(
        self,
        connected: &ConnectedKafkaBroker,
    ) -> impl Future<Output = Result<Self::Live, PairError>> {
        ready(Ok(TransactionalPartitions {
            inner: Arc::new(PartitionsInner {
                state: Arc::clone(connected.state()),
                template: self.template,
                publishers: Mutex::new(HashMap::new()),
            }),
        }))
    }

    /// The destination topic of this publish, which is what the channel stands for.
    #[cfg(feature = "asyncapi")]
    fn channel_bindings(&self, channel: &str) -> Bindings {
        crate::bindings::channel(channel)
    }
}

/// Transactional publishers, one per source partition, materialized on first use.
///
/// Kafka permits one open transaction per producer and one live producer per transactional id
/// (initializing a second fences the first), so concurrent transactional handlers need one
/// producer each. The source partition is the natural scope: under the default
/// [`LaneKey::Partition`](crate::LaneKey::Partition) worker pool a partition's deliveries
/// process serially on one lane, so a publisher per partition gives every lane an independent
/// transaction with no coordination. The id set (`"{base}-p{partition}"`) follows the topic's
/// partitions rather than the worker count: changing `workers(n)` neither changes the ids nor
/// weakens zombie fencing - the scheme Kafka Streams uses for its per-task producers.
///
/// Not for [`LaneKey::RecordKey`](crate::LaneKey::RecordKey) pools: record-key lanes spread
/// one partition across lanes, so two lanes would share a partition's publisher and collide
/// on its single transaction ([`KafkaError::TransactionBusy`]).
///
/// Clones share the cache, so one injected handle serves every handler invocation.
#[derive(Clone)]
pub struct TransactionalPartitions {
    inner: Arc<PartitionsInner>,
}

struct PartitionsInner {
    state: Arc<ConnState>,
    template: KafkaTransactionalPublish,
    /// The per-partition publishers. The set of partitions is only known as deliveries arrive,
    /// so materialization stays lazy here (a cell per partition, so two lanes racing the same
    /// partition initialize one producer, not two).
    publishers: Mutex<HashMap<i32, Arc<OnceCell<KafkaTransactionalPublisher>>>>,
}

impl fmt::Debug for TransactionalPartitions {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("TransactionalPartitions")
            .field("id_base", &self.inner.template.id())
            .finish_non_exhaustive()
    }
}

impl TransactionalPartitions {
    /// The publisher owning `partition`'s transactional id, created and initialized on first
    /// use.
    ///
    /// `partition` is the delivery's source partition (`KafkaContext`'s `Partition` field in a
    /// handler); passing anything else still works but forfeits the serialization argument
    /// that makes the per-partition scope safe.
    ///
    /// # Errors
    ///
    /// Returns [`KafkaError::Closed`] once the broker has shut down and
    /// [`KafkaError::Publish`] when the partition's producer cannot be created or its
    /// transactions cannot be initialized.
    ///
    /// # Panics
    ///
    /// Panics when the internal cache mutex is poisoned, which requires a prior panic while
    /// materializing a publisher (an invariant violation, not an operational failure).
    pub async fn for_partition(
        &self,
        partition: i32,
    ) -> Result<KafkaTransactionalPublisher, KafkaError> {
        let cell = {
            let mut publishers = self
                .inner
                .publishers
                .lock()
                .expect("partition publisher cache mutex poisoned");
            Arc::clone(publishers.entry(partition).or_default())
        };
        let policy = self
            .inner
            .template
            .with_id(format!("{}-p{partition}", self.inner.template.id()));
        cell.get_or_try_init(|| open_transactional(&self.inner.state, &policy))
            .await
            .cloned()
    }
}

/// The capability of handing out one transactional publisher per source partition.
///
/// A handler that drives per-partition transactions names this in its `Out` slot
/// (`Out(lanes): Out<impl PartitionLanes>`) instead of a publisher type: the concrete value is
/// [`TransactionalPartitions`], inferred from the [`KafkaTransactionalPublish::per_partition`]
/// policy attached at the include site.
///
/// # Test capture
///
/// This is a router, not a publisher: it hands out a publisher of its own rather than sending a
/// message. What a lane then publishes leaves through that publisher, so it lands in the
/// broker's publish log and not in the slot's test record - the same boundary a settled owned
/// transaction's buffer has. Assert on the publish log for lane traffic, and keep the slot
/// record for handlers that publish through the slot itself.
///
/// # Examples
///
/// ```
/// use ruststream_rdkafka::prelude::*;
/// use serde::{Deserialize, Serialize};
///
/// #[derive(Deserialize)]
/// struct Invoice {
///     id: u64,
///     lines: u64,
/// }
///
/// #[derive(Serialize, Outgoing)]
/// #[outgoing(name = "invoice-lines")]
/// struct InvoiceLine {
///     invoice_id: u64,
///     line: u64,
/// }
///
/// /// Each lane runs the transactions of its own source partition, so four workers never share
/// /// one producer.
/// #[subscriber(KafkaTopic::new("invoices").commit(Commit::Tracked), workers(4, by_key))]
/// async fn bill(
///     invoice: &Invoice,
///     Ctx(partition): Ctx<Partition>,
///     Out(lanes): Out<impl PartitionLanes>,
/// ) -> HandlerOutcome {
///     let Ok(lane) = lanes.for_partition(partition).await else {
///         return HandlerOutcome::retry();
///     };
///     if lane.begin_transaction().await.is_err() {
///         return HandlerOutcome::retry();
///     }
///     for line in 0..invoice.lines {
///         let entry = InvoiceLine { invoice_id: invoice.id, line };
///         if lane.message(&entry).publish().await.is_err() {
///             lane.abort().await.ok();
///             return HandlerOutcome::retry();
///         }
///     }
///     if lane.commit().await.is_err() {
///         return HandlerOutcome::retry();
///     }
///     HandlerOutcome::ack()
/// }
/// ```
#[diagnostic::on_unimplemented(
    message = "`{Self}` does not hand out per-partition transactional publishers",
    note = "for an `Out<impl PartitionLanes, _>` slot, attach a \
            `KafkaPublish::default().transactional_id(..).per_partition()` policy"
)]
pub trait PartitionLanes: Send + Sync {
    /// The transactional publisher a lane is handed.
    ///
    /// An associated type rather than the concrete [`KafkaTransactionalPublisher`]: a handler
    /// bounds its slot `Out<impl PartitionLanes>` on the capability, and the policy attached at
    /// the mount site decides the publisher.
    type Publisher: TransactionalPublisher<Error = KafkaError>;

    /// The publisher owning `partition`'s transactional id.
    ///
    /// # Errors
    ///
    /// See [`TransactionalPartitions::for_partition`].
    fn for_partition(
        &self,
        partition: i32,
    ) -> impl Future<Output = Result<Self::Publisher, KafkaError>> + Send;
}

impl PartitionLanes for TransactionalPartitions {
    type Publisher = KafkaTransactionalPublisher;

    fn for_partition(
        &self,
        partition: i32,
    ) -> impl Future<Output = Result<KafkaTransactionalPublisher, KafkaError>> + Send {
        Self::for_partition(self, partition)
    }
}

// The arena entry a handler's `Out` parameter binds. `Slot`'s `Deref` already routes method
// calls to the wired value, but only this impl lets a body hand the entry to a function generic
// over the capability (`fn issue<L: PartitionLanes>(lanes: &L, ..)`), which is the whole point of
// bounding the slot with a trait instead of a concrete type. Both paths reach the same unwrapped
// value, which is what puts a lane's publishes outside the slot's test capture (see the trait's
// documentation).
impl<M, L, EncodeCodec, Pipe, Body> PartitionLanes for Slot<M, L, EncodeCodec, Pipe, Body>
where
    L: PartitionLanes,
    EncodeCodec: Send + Sync,
    // The entry's publish path, over the wired value whose per-record settings its transforms
    // write. A lane's traffic never travels it (it leaves through the unwrapped value), but
    // naming the bound the mount site's entry already carries keeps this impl on exactly the
    // slots the runtime builds.
    Pipe: OutPipeline<L>,
{
    type Publisher = L::Publisher;

    fn for_partition(
        &self,
        partition: i32,
    ) -> impl Future<Output = Result<Self::Publisher, KafkaError>> + Send {
        (**self).for_partition(partition)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn transactional_id_is_a_type_transition() {
        let plain = KafkaPublish::default().queue_timeout(Duration::from_secs(1));
        let transactional = plain.transactional_id("svc-1");
        assert_eq!(transactional.id(), "svc-1");
        assert_eq!(transactional.queue_timeout, plain.queue_timeout_setting());
        assert_eq!(
            transactional.transaction_timeout,
            DEFAULT_TRANSACTION_TIMEOUT
        );
    }

    #[test]
    fn per_partition_derives_ids_from_the_base() {
        let policy = KafkaPublish::default()
            .transactional_id("svc-1")
            .transaction_timeout(Duration::from_secs(5));
        let derived = policy.with_id(format!("{}-p3", policy.id()));
        assert_eq!(derived.id(), "svc-1-p3");
        assert_eq!(derived.transaction_timeout, Duration::from_secs(5));
    }
}