ruststream_rdkafka/testing/
subscriber.rs

1//! The in-process subscriber and its delivery type.
2
3use std::fmt;
4use std::sync::Arc;
5
6use futures::Stream;
7use ruststream::testing::Coordinator;
8use ruststream::{AckError, Headers, IncomingMessage, Partitioned, Subscriber};
9
10use super::broker::TestBrokerState;
11use super::router::{DeliveryReceiver, DeliverySender, SubscriptionId, TestDelivery};
12use crate::error::KafkaError;
13
14/// In-process subscriber on one topic name.
15///
16/// Yielded messages settle like the routing contract expects: ack finalizes, `nack(true)`
17/// re-enqueues to this same subscription, `nack(false)` drops. The real transport's
18/// committed-position semantics (holes, watermarks, redelivery on rebalance) are deliberately
19/// not simulated.
20pub struct KafkaTestSubscriber {
21    state: Arc<TestBrokerState>,
22    id: SubscriptionId,
23    topic: String,
24    sender: DeliverySender,
25    receiver: DeliveryReceiver,
26    coordinator: Option<Coordinator>,
27}
28
29impl KafkaTestSubscriber {
30    pub(crate) fn open(state: &Arc<TestBrokerState>, topic: String) -> Self {
31        let (id, sender, receiver) = state.router.subscribe(topic.clone());
32        let coordinator = state.coordinator();
33        Self {
34            state: Arc::clone(state),
35            id,
36            topic,
37            sender,
38            receiver,
39            coordinator,
40        }
41    }
42
43    /// The topic this subscriber consumes.
44    #[must_use]
45    pub fn topic(&self) -> &str {
46        &self.topic
47    }
48}
49
50impl Drop for KafkaTestSubscriber {
51    fn drop(&mut self) {
52        self.state.router.unsubscribe(self.id);
53    }
54}
55
56impl fmt::Debug for KafkaTestSubscriber {
57    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
58        f.debug_struct("KafkaTestSubscriber")
59            .field("topic", &self.topic)
60            .finish_non_exhaustive()
61    }
62}
63
64impl Subscriber for KafkaTestSubscriber {
65    type Message = KafkaTestMessage;
66    type Error = KafkaError;
67
68    /// Streams injected deliveries; never yields an error.
69    ///
70    /// # Cancel safety
71    ///
72    /// Cancel safe and re-enterable: the receiver is polled in place, so dropping the returned
73    /// stream loses nothing and `stream` can be called again.
74    fn stream(&mut self) -> impl Stream<Item = Result<Self::Message, Self::Error>> + Send + '_ {
75        let Self {
76            receiver,
77            sender,
78            coordinator,
79            ..
80        } = self;
81        futures::stream::poll_fn(move |cx| {
82            receiver.poll_recv(cx).map(|delivery| {
83                delivery.map(|delivery| {
84                    Ok(KafkaTestMessage {
85                        delivery: Some(delivery),
86                        sender: sender.clone(),
87                        coordinator: coordinator.clone(),
88                    })
89                })
90            })
91        })
92    }
93}
94
95/// One in-process delivery.
96pub struct KafkaTestMessage {
97    delivery: Option<TestDelivery>,
98    sender: DeliverySender,
99    coordinator: Option<Coordinator>,
100}
101
102impl KafkaTestMessage {
103    fn take(&mut self) -> TestDelivery {
104        // The settle methods consume `self`, so a second settle cannot compile; reaching this
105        // twice is an internal invariant violation.
106        self.delivery
107            .take()
108            .expect("KafkaTestMessage settled twice")
109    }
110}
111
112impl Drop for KafkaTestMessage {
113    fn drop(&mut self) {
114        // Balance the router's `enqueued` exactly once per delivery, whatever the dispatch
115        // path did (ack, nack, panic, or plain drop).
116        if let Some(coordinator) = self.coordinator.take() {
117            coordinator.consumed();
118        }
119    }
120}
121
122impl fmt::Debug for KafkaTestMessage {
123    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
124        f.debug_struct("KafkaTestMessage")
125            .field("delivery", &self.delivery)
126            .finish_non_exhaustive()
127    }
128}
129
130impl IncomingMessage for KafkaTestMessage {
131    fn payload(&self) -> &[u8] {
132        &self
133            .delivery
134            .as_ref()
135            .expect("message accessed after settlement")
136            .payload
137    }
138
139    fn headers(&self) -> &Headers {
140        &self
141            .delivery
142            .as_ref()
143            .expect("message accessed after settlement")
144            .headers
145    }
146
147    /// The partition key from the `PARTITION_KEY_HEADER`, mirroring the real message so keyed
148    /// worker lanes behave the same in-process.
149    fn partition_key(&self) -> Option<&[u8]> {
150        self.headers().get(crate::PARTITION_KEY_HEADER)
151    }
152
153    /// Finalizes the delivery.
154    ///
155    /// # Errors
156    ///
157    /// Never fails; the in-process transport has no position to store.
158    async fn ack(mut self) -> Result<(), AckError> {
159        drop(self.take());
160        Ok(())
161    }
162
163    /// Re-enqueues to the same subscription (`requeue = true`) or drops (`requeue = false`).
164    ///
165    /// # Errors
166    ///
167    /// Never fails; the in-process transport has no position to store.
168    async fn nack(mut self, requeue: bool) -> Result<(), AckError> {
169        let delivery = self.take();
170        if requeue && self.sender.send(delivery).is_ok() {
171            // This bypasses the router fanout, so account for the new in-flight delivery here.
172            if let Some(coordinator) = &self.coordinator {
173                coordinator.enqueued();
174            }
175        }
176        Ok(())
177    }
178}
179
180impl Partitioned for KafkaTestMessage {
181    /// The partition key from the `PARTITION_KEY_HEADER`, mirroring the real message.
182    fn partition_key(&self) -> Option<&[u8]> {
183        self.headers().get(crate::PARTITION_KEY_HEADER)
184    }
185}
ruststream_rdkafka/testing/subscriber.rs

ruststream_rdkafka/testing/
subscriber.rs
