ruststream_rdkafka/
distribution.rs

1//! Producer-side distribution policies built on the explicit-partition header.
2//!
3//! librdkafka's partitioner families (random, consistent, murmur2, fnv1a) cannot express
4//! per-message round-robin, and keyless distribution may batch-stick to one partition. For
5//! workloads with long, near-constant per-message processing times that unevenness turns into
6//! one hot consumer and idle peers; [`RoundRobin`] stamps each outgoing reply with the next
7//! partition in the cycle instead.
8
9use std::sync::atomic::{AtomicU64, Ordering};
10
11use ruststream::runtime::{Outgoing, PublishContext, PublishTransform};
12
13use crate::message::{PARTITION_HEADER, PARTITION_KEY_HEADER};
14
15/// A [`PublishTransform`] distributing replies round-robin across the first `count` partitions.
16///
17/// Each reply gets the [`PARTITION_HEADER`] of an incrementing counter modulo `count`, so the
18/// publisher targets partitions 0..count in a cycle, one message each - the evenest possible
19/// spread for long, near-constant-cost messages. A reply that already carries an explicit
20/// partition or a record key is left alone: keys exist for ordering, and overriding either
21/// would silently break the caller's placement.
22///
23/// The count is explicit on purpose (cheap and predictable); it must match the destination
24/// topic's partition count, or the tail partitions simply receive nothing (a smaller count)
25/// or publishes fail (a larger one).
26///
27/// # Examples
28///
29/// ```no_run
30/// use ruststream::runtime::TypedPublisher;
31/// use ruststream_rdkafka::{KafkaBroker, RoundRobin};
32///
33/// # #[cfg(feature = "json")]
34/// # fn wire(broker: &KafkaBroker) {
35/// let replies = TypedPublisher::new(broker.publisher()).transform(RoundRobin::partitions(8));
36/// # let _ = replies;
37/// # }
38/// ```
39#[derive(Debug)]
40pub struct RoundRobin {
41    count: u64,
42    next: AtomicU64,
43}
44
45impl RoundRobin {
46    /// A round-robin cycle over partitions `0..count`.
47    ///
48    /// # Panics
49    ///
50    /// Panics when `count` is zero: a cycle over no partitions cannot place anything.
51    #[must_use]
52    pub fn partitions(count: i32) -> Self {
53        assert!(
54            count > 0,
55            "a round-robin cycle needs at least one partition"
56        );
57        Self {
58            #[allow(clippy::cast_sign_loss)] // asserted positive above
59            count: count as u64,
60            next: AtomicU64::new(0),
61        }
62    }
63}
64
65impl<C> PublishTransform<C> for RoundRobin {
66    fn apply(&self, out: &mut Outgoing<'_>, _cx: &PublishContext<'_, C>) {
67        if out.headers().get(PARTITION_HEADER).is_some()
68            || out.headers().get(PARTITION_KEY_HEADER).is_some()
69        {
70            return;
71        }
72        let slot = self.next.fetch_add(1, Ordering::Relaxed) % self.count;
73        out.headers_mut().insert(PARTITION_HEADER, slot.to_string());
74    }
75}
ruststream_rdkafka/distribution.rs

ruststream_rdkafka/
distribution.rs
