ruststream_pulsar/
message.rs1use bytes::Bytes;
7use pulsar::proto::MessageIdData;
8use ruststream::{AckError, Headers, IncomingMessage, OutgoingMessage, Partitioned, Positioned};
9use tokio::sync::{mpsc, oneshot};
10
11use crate::error::PulsarError;
12
13pub const PARTITION_KEY_HEADER: &str = "partition-key";
19
20#[derive(Debug)]
22pub(crate) enum SettleKind {
23 Ack,
25 Nack,
27}
28
29#[derive(Debug, Clone, PartialEq, Eq)]
51pub enum PulsarPosition {
52 Earliest,
59 Latest,
61 MessageId(MessageIdData),
63 Timestamp(u64),
65}
66
67impl PulsarPosition {
68 #[must_use]
71 pub fn earliest() -> Self {
72 Self::Earliest
73 }
74
75 #[must_use]
77 pub fn latest() -> Self {
78 Self::Latest
79 }
80
81 #[must_use]
84 pub fn timestamp(millis: u64) -> Self {
85 Self::Timestamp(millis)
86 }
87}
88
89#[derive(Debug)]
91pub(crate) struct SeekCmd {
92 pub(crate) position: PulsarPosition,
93 pub(crate) done: oneshot::Sender<Result<(), PulsarError>>,
94}
95
96#[derive(Debug)]
99pub(crate) struct SettleCmd {
100 pub(crate) topic: String,
101 pub(crate) id: MessageIdData,
102 pub(crate) kind: SettleKind,
103 pub(crate) done: oneshot::Sender<Result<(), AckError>>,
104}
105
106#[derive(Debug)]
108pub(crate) enum DriverCmd {
109 Settle(SettleCmd),
110 Seek(SeekCmd),
111}
112
113pub(crate) type SettleSender = mpsc::UnboundedSender<DriverCmd>;
114
115pub struct PulsarMessage {
123 payload: Bytes,
124 headers: Headers,
125 topic: String,
126 id: MessageIdData,
127 settle: SettleSender,
128}
129
130impl std::fmt::Debug for PulsarMessage {
131 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
132 f.debug_struct("PulsarMessage")
133 .field("topic", &self.topic)
134 .field("payload_len", &self.payload.len())
135 .finish_non_exhaustive()
136 }
137}
138
139impl PulsarMessage {
140 pub(crate) fn new(message: &pulsar::consumer::Message<Vec<u8>>, settle: SettleSender) -> Self {
141 let metadata = message.metadata();
142 let mut headers = Headers::with_capacity(metadata.properties.len() + 1);
143 for kv in &metadata.properties {
144 headers.insert(kv.key.clone(), kv.value.clone());
145 }
146 if let Some(key) = &metadata.partition_key {
147 headers.insert(PARTITION_KEY_HEADER, key.clone());
148 }
149 Self {
150 payload: Bytes::copy_from_slice(&message.payload.data),
151 headers,
152 topic: message.topic.clone(),
153 id: message.message_id().clone(),
154 settle,
155 }
156 }
157
158 #[must_use]
161 pub fn topic(&self) -> &str {
162 &self.topic
163 }
164
165 async fn send_settle(self, kind: SettleKind) -> Result<(), AckError> {
166 let (done, wait) = oneshot::channel();
167 self.settle
168 .send(DriverCmd::Settle(SettleCmd {
169 topic: self.topic,
170 id: self.id,
171 kind,
172 done,
173 }))
174 .map_err(|_| {
175 AckError::Broker(Box::from("the subscription's driver task has shut down"))
176 })?;
177 wait.await.map_err(|_| {
178 AckError::Broker(Box::from("the subscription's driver task has shut down"))
179 })?
180 }
181}
182
183impl Positioned for PulsarMessage {
184 type Position = PulsarPosition;
185
186 fn position(&self) -> PulsarPosition {
187 PulsarPosition::MessageId(self.id.clone())
188 }
189}
190
191impl Partitioned for PulsarMessage {
192 fn partition_key(&self) -> Option<&[u8]> {
193 self.headers.get(PARTITION_KEY_HEADER)
194 }
195}
196
197impl IncomingMessage for PulsarMessage {
198 fn payload(&self) -> &[u8] {
199 &self.payload
200 }
201
202 fn headers(&self) -> &Headers {
203 &self.headers
204 }
205
206 async fn ack(self) -> Result<(), AckError> {
207 self.send_settle(SettleKind::Ack).await
208 }
209
210 async fn nack(self, requeue: bool) -> Result<(), AckError> {
211 if requeue {
212 self.send_settle(SettleKind::Nack).await
213 } else {
214 self.send_settle(SettleKind::Ack).await
217 }
218 }
219
220 fn partition_key(&self) -> Option<&[u8]> {
221 Partitioned::partition_key(self)
222 }
223}
224
225pub(crate) fn to_pulsar_message(msg: &OutgoingMessage<'_>) -> pulsar::producer::Message {
227 let headers = msg.headers();
228 let mut properties = std::collections::HashMap::with_capacity(headers.len());
229 let mut partition_key = None;
230 for (name, value) in headers.iter() {
231 let text = String::from_utf8_lossy(value).into_owned();
232 if name == PARTITION_KEY_HEADER {
233 partition_key = Some(text);
234 } else {
235 properties.insert(name.to_owned(), text);
236 }
237 }
238 pulsar::producer::Message {
239 payload: msg.payload().to_vec(),
240 properties,
241 partition_key,
242 ..Default::default()
243 }
244}
245
246#[cfg(test)]
247mod tests {
248 use super::*;
249
250 #[test]
251 fn partition_key_header_becomes_the_partition_key() {
252 let mut headers = Headers::new();
253 headers.insert(PARTITION_KEY_HEADER, "user-42");
254 headers.insert("x-tenant", "acme");
255 let outgoing = OutgoingMessage::new("orders", b"{}".as_slice()).with_headers(headers);
256
257 let message = to_pulsar_message(&outgoing);
258 assert_eq!(message.partition_key.as_deref(), Some("user-42"));
259 assert_eq!(
260 message.properties.get("x-tenant").map(String::as_str),
261 Some("acme")
262 );
263 assert!(!message.properties.contains_key(PARTITION_KEY_HEADER));
264 }
265}