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/// ```
30/// use ruststream::runtime::TypedPublisher;
31/// use ruststream_rdkafka::{KafkaPublish, RoundRobin};
32///
33/// # #[cfg(feature = "json")]
34/// # fn wire() {
35/// let replies =
36///     TypedPublisher::new(KafkaPublish::default()).transform(RoundRobin::partitions(8));
37/// # let _ = replies;
38/// # }
39/// ```
40#[derive(Debug)]
41pub struct RoundRobin {
42    count: u64,
43    next: AtomicU64,
44}
45
46impl RoundRobin {
47    /// A round-robin cycle over partitions `0..count`.
48    ///
49    /// # Panics
50    ///
51    /// Panics when `count` is zero: a cycle over no partitions cannot place anything.
52    #[must_use]
53    pub fn partitions(count: i32) -> Self {
54        assert!(
55            count > 0,
56            "a round-robin cycle needs at least one partition"
57        );
58        Self {
59            #[allow(clippy::cast_sign_loss)] // asserted positive above
60            count: count as u64,
61            next: AtomicU64::new(0),
62        }
63    }
64}
65
66impl<C> PublishTransform<C> for RoundRobin {
67    fn apply(&self, out: &mut Outgoing<'_>, _cx: &PublishContext<'_, C>) {
68        if out.headers().get(PARTITION_HEADER).is_some()
69            || out.headers().get(PARTITION_KEY_HEADER).is_some()
70        {
71            return;
72        }
73        let slot = self.next.fetch_add(1, Ordering::Relaxed) % self.count;
74        out.headers_mut().insert(PARTITION_HEADER, slot.to_string());
75    }
76}
ruststream_rdkafka/distribution.rs

ruststream_rdkafka/
distribution.rs
