ruststream_rdkafka/
distribution.rs

1//! Producer-side distribution policies built on the per-record partition setting.
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`] pins each outgoing record to the next
7//! partition in the cycle instead.
8
9use std::sync::atomic::{AtomicU64, Ordering};
10
11use ruststream::runtime::{ContextKind, Outgoing, PublishTransform, Reads};
12
13use crate::message::PARTITION_KEY_HEADER;
14use crate::publisher::KafkaOptions;
15
16/// A [`PublishTransform`] distributing records round-robin across the first `count` partitions.
17///
18/// Each record is pinned to an incrementing counter modulo `count`, so the publisher targets
19/// partitions 0..count in a cycle, one message each - the evenest possible spread for long,
20/// near-constant-cost messages. A record that already names a partition or carries a record key
21/// is left alone: keys exist for ordering, and overriding either would silently break the
22/// caller's placement.
23///
24/// The cycle writes the same per-record setting the builder's
25/// [`partition`](crate::KafkaPublishSteps::partition) step writes, so it mounts only over a
26/// publisher whose [`Publisher::Options`](ruststream::Publisher::Options) are [`KafkaOptions`].
27/// On an `Out` slot that setting arrives carrying whatever the call site chose, and the cycle
28/// steps aside for it; a reply has no call site, so there the cycle always places.
29///
30/// The count is explicit on purpose (cheap and predictable); it must match the destination
31/// topic's partition count, or the tail partitions simply receive nothing (a smaller count)
32/// or publishes fail (a larger one).
33///
34/// # Examples
35///
36/// The transform is a step of the mount site's chain, right after the publish policy:
37///
38/// ```
39/// # #[cfg(feature = "json")]
40/// # mod demo {
41/// use ruststream_rdkafka::prelude::*;
42/// # #[derive(serde::Deserialize)]
43/// # struct Order { id: u64 }
44/// # #[derive(serde::Serialize, Outgoing)]
45/// # #[outgoing(name = "work-items")]
46/// # struct WorkItem { order_id: u64 }
47///
48/// #[subscriber("orders", reply)]
49/// async fn plan(order: &Order) -> WorkItem {
50///     WorkItem { order_id: order.id }
51/// }
52///
53/// fn app() -> RustStream {
54///     RustStream::new(AppInfo::new("planner", "0.1.0"))
55///         .with_broker(KafkaBroker::new(["localhost:9092"]), |b| {
56///             b.include(plan)
57///                 .out_reply(Publish::default())
58///                 .transform(RoundRobin::partitions(8));
59///         })
60/// }
61/// # }
62/// ```
63#[derive(Debug)]
64pub struct RoundRobin {
65    count: u64,
66    next: AtomicU64,
67}
68
69impl RoundRobin {
70    /// A round-robin cycle over partitions `0..count`.
71    ///
72    /// # Panics
73    ///
74    /// Panics when `count` is zero: a cycle over no partitions cannot place anything.
75    #[must_use]
76    pub fn partitions(count: i32) -> Self {
77        assert!(
78            count > 0,
79            "a round-robin cycle needs at least one partition"
80        );
81        Self {
82            #[allow(clippy::cast_sign_loss)] // asserted positive above
83            count: count as u64,
84            next: AtomicU64::new(0),
85        }
86    }
87}
88
89impl<K: ContextKind> PublishTransform<K, KafkaOptions> for RoundRobin {
90    // It reads nothing from the position, so it mounts on a reply and on a slot alike, and the
91    // destination is not its business.
92    type Destination = Reads;
93
94    fn apply(&self, out: &mut Outgoing<'_>, options: &mut Option<KafkaOptions>, _cx: &K::View<'_>) {
95        if out.headers().get(PARTITION_KEY_HEADER).is_some()
96            || options.is_some_and(|options| options.partition_setting().is_some())
97        {
98            return;
99        }
100        // The cycle position is below `count`, which came from a positive `i32`, so it is one.
101        #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
102        let partition = (self.next.fetch_add(1, Ordering::Relaxed) % self.count) as i32;
103        let placed = options.unwrap_or_default().partition(partition);
104        *options = Some(placed);
105    }
106}
ruststream_rdkafka/distribution.rs

ruststream_rdkafka/
distribution.rs
