ruststream_sqs_sns/
message.rs

1//! [`SqsMessage`] and the mapping between `RustStream` headers and SQS message attributes.
2//!
3//! Message attributes carry headers directly (String for UTF-8 values, Binary otherwise) - no
4//! envelope format is invented. The one transport constraint is the body: SQS bodies are text,
5//! so a payload that is not valid UTF-8 travels base64-encoded with a marker attribute, and is
6//! decoded transparently on receive.
7
8use std::time::Duration;
9
10use aws_sdk_sqs::Client;
11use aws_sdk_sqs::primitives::Blob;
12use aws_sdk_sqs::types::{
13    Message as AwsMessage, MessageAttributeValue, MessageSystemAttributeName,
14};
15use base64::Engine as _;
16use base64::engine::general_purpose::STANDARD as BASE64;
17use bytes::Bytes;
18use ruststream::{AckError, Headers, IncomingMessage, Partitioned};
19use tokio::task::JoinHandle;
20
21use crate::error::sdk_err;
22
23/// Header carrying the partition key, mapped onto the FIFO message group id.
24///
25/// Mirrors the in-memory broker's convention, so services can switch brokers without changing
26/// their headers.
27pub const PARTITION_KEY_HEADER: &str = "partition-key";
28
29/// Header exposing the approximate receive count on received messages.
30pub const RECEIVE_COUNT_HEADER: &str = "sqs-receive-count";
31
32/// Marker attribute set when the payload travels base64-encoded (SQS bodies are text; binary
33/// payloads have no other faithful form).
34pub(crate) const ENCODING_ATTRIBUTE: &str = "ruststream-payload-encoding";
35
36/// A message delivered by an [`SqsSubscriber`](crate::SqsSubscriber).
37///
38/// `ack` deletes the message; `nack(requeue = true)` zeroes its visibility so it redelivers
39/// immediately; `nack_after(delay)` sets the visibility to the delay, so deferred retry is
40/// native. `nack(requeue = false)` deletes: SQS has no drop verb short of deletion - poison
41/// routing belongs to the queue's redrive policy, driven by repeated requeues.
42///
43/// While the handle is alive, a background task keeps extending the message's visibility, so a
44/// handler outliving the visibility timeout does not cause a concurrent redelivery.
45pub struct SqsMessage {
46    payload: Bytes,
47    headers: Headers,
48    client: Client,
49    queue_url: String,
50    receipt: String,
51    extender: JoinHandle<()>,
52}
53
54impl std::fmt::Debug for SqsMessage {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        f.debug_struct("SqsMessage")
57            .field("payload_len", &self.payload.len())
58            .field("queue_url", &self.queue_url)
59            .finish_non_exhaustive()
60    }
61}
62
63impl Drop for SqsMessage {
64    fn drop(&mut self) {
65        // An unsettled drop stops the extension; the message redelivers when its current
66        // visibility lapses, which is the at-least-once contract.
67        self.extender.abort();
68    }
69}
70
71impl SqsMessage {
72    pub(crate) fn new(
73        message: &AwsMessage,
74        client: Client,
75        queue_url: String,
76        receipt: String,
77        visibility: Duration,
78    ) -> Self {
79        let (payload, headers) = decode_message(message);
80        // Why a per-message watchdog: SQS has no lease API - a handler outliving the
81        // visibility timeout would get a concurrent redelivery, so the crate extends the
82        // visibility for as long as the handle is held (the issue's one piece of real
83        // machinery). Aborted on settle or drop.
84        let extender = tokio::spawn(extend_visibility(
85            client.clone(),
86            queue_url.clone(),
87            receipt.clone(),
88            visibility,
89        ));
90        Self {
91            payload,
92            headers,
93            client,
94            queue_url,
95            receipt,
96            extender,
97        }
98    }
99
100    async fn delete(&self) -> Result<(), AckError> {
101        self.client
102            .delete_message()
103            .queue_url(&self.queue_url)
104            .receipt_handle(&self.receipt)
105            .send()
106            .await
107            .map(|_| ())
108            .map_err(|e| AckError::Broker(sdk_err(&e)))
109    }
110
111    async fn set_visibility(&self, seconds: i32) -> Result<(), AckError> {
112        self.client
113            .change_message_visibility()
114            .queue_url(&self.queue_url)
115            .receipt_handle(&self.receipt)
116            .visibility_timeout(seconds)
117            .send()
118            .await
119            .map(|_| ())
120            .map_err(|e| AckError::Broker(sdk_err(&e)))
121    }
122}
123
124impl Partitioned for SqsMessage {
125    fn partition_key(&self) -> Option<&[u8]> {
126        self.headers.get(PARTITION_KEY_HEADER)
127    }
128}
129
130impl IncomingMessage for SqsMessage {
131    fn payload(&self) -> &[u8] {
132        &self.payload
133    }
134
135    fn headers(&self) -> &Headers {
136        &self.headers
137    }
138
139    async fn ack(self) -> Result<(), AckError> {
140        self.extender.abort();
141        self.delete().await
142    }
143
144    async fn nack(self, requeue: bool) -> Result<(), AckError> {
145        self.extender.abort();
146        if requeue {
147            self.set_visibility(0).await
148        } else {
149            // Deleting IS the drop: SQS cannot discard without deleting, and the redrive
150            // policy owns poison-message routing.
151            self.delete().await
152        }
153    }
154
155    /// Every SQS delivery honors a delayed redelivery: the visibility timeout is the delay, so
156    /// the runtime must take `nack_after` here instead of its broker-agnostic deferred
157    /// re-publish, which would re-publish a copy and reset the receive count.
158    fn supports_nack_after(&self) -> bool {
159        true
160    }
161
162    async fn nack_after(self, delay: Duration) -> Result<(), AckError> {
163        self.extender.abort();
164        // Setting the visibility to the delay is the native deferred retry (capped at the
165        // protocol's 12 hours).
166        let seconds = i32::try_from(delay.as_secs().min(43_200)).unwrap_or(43_200);
167        self.set_visibility(seconds).await
168    }
169
170    fn partition_key(&self) -> Option<&[u8]> {
171        Partitioned::partition_key(self)
172    }
173}
174
175/// Keeps a message invisible while its handle is alive: re-arms the visibility to `visibility`
176/// every half period. Aborted on settle/drop; a failed extension is logged and retried on the
177/// next tick (the message may redeliver, which at-least-once permits).
178async fn extend_visibility(
179    client: Client,
180    queue_url: String,
181    receipt: String,
182    visibility: Duration,
183) {
184    let period = (visibility / 2).max(Duration::from_secs(1));
185    let seconds = i32::try_from(visibility.as_secs().min(43_200)).unwrap_or(43_200);
186    loop {
187        tokio::time::sleep(period).await;
188        let outcome = client
189            .change_message_visibility()
190            .queue_url(&queue_url)
191            .receipt_handle(&receipt)
192            .visibility_timeout(seconds)
193            .send()
194            .await;
195        if let Err(err) = outcome {
196            tracing::debug!(
197                queue_url = %queue_url,
198                error = %aws_sdk_sqs::error::DisplayErrorContext(&err),
199                "sqs visibility extension failed"
200            );
201        }
202    }
203}
204
205fn decode_message(message: &AwsMessage) -> (Bytes, Headers) {
206    let mut headers = Headers::new();
207    let mut base64_payload = false;
208    if let Some(attributes) = message.message_attributes() {
209        for (name, value) in attributes {
210            if name == ENCODING_ATTRIBUTE {
211                base64_payload = value.string_value() == Some("base64");
212                continue;
213            }
214            if let Some(text) = value.string_value() {
215                headers.insert(name.clone(), text.to_owned());
216            } else if let Some(blob) = value.binary_value() {
217                headers.insert(name.clone(), Bytes::copy_from_slice(blob.as_ref()));
218            }
219        }
220    }
221    if let Some(system) = message.attributes() {
222        if let Some(group) = system.get(&MessageSystemAttributeName::MessageGroupId) {
223            headers.insert(PARTITION_KEY_HEADER, group.clone());
224        }
225        if let Some(count) = system.get(&MessageSystemAttributeName::ApproximateReceiveCount) {
226            headers.insert(RECEIVE_COUNT_HEADER, count.clone());
227        }
228    }
229
230    let body = message.body().unwrap_or_default();
231    let payload = if base64_payload {
232        BASE64
233            .decode(body)
234            .map_or_else(|_| Bytes::copy_from_slice(body.as_bytes()), Bytes::from)
235    } else {
236        Bytes::copy_from_slice(body.as_bytes())
237    };
238    (payload, headers)
239}
240
241/// Encodes a payload into an SQS body: UTF-8 passes through, anything else travels base64 with
242/// the marker attribute. Returns the body and whether the marker must be set.
243pub(crate) fn encode_body(payload: &[u8]) -> (String, bool) {
244    std::str::from_utf8(payload).map_or_else(
245        |_| (BASE64.encode(payload), true),
246        |text| (text.to_owned(), false),
247    )
248}
249
250/// Converts headers into SQS message attributes (String for UTF-8 values, Binary otherwise),
251/// pulling the partition key out for the FIFO group id.
252pub(crate) fn encode_attributes(
253    headers: &Headers,
254    base64_marker: bool,
255) -> (
256    std::collections::HashMap<String, MessageAttributeValue>,
257    Option<String>,
258) {
259    let mut attributes = std::collections::HashMap::new();
260    let mut group = None;
261    for (name, value) in headers.iter() {
262        if name == PARTITION_KEY_HEADER {
263            group = Some(String::from_utf8_lossy(value).into_owned());
264            continue;
265        }
266        let attribute = std::str::from_utf8(value).map_or_else(
267            |_| {
268                MessageAttributeValue::builder()
269                    .data_type("Binary")
270                    .binary_value(Blob::new(value))
271                    .build()
272            },
273            |text| {
274                MessageAttributeValue::builder()
275                    .data_type("String")
276                    .string_value(text)
277                    .build()
278            },
279        );
280        if let Ok(attribute) = attribute {
281            attributes.insert(name.to_owned(), attribute);
282        }
283    }
284    if base64_marker
285        && let Ok(marker) = MessageAttributeValue::builder()
286            .data_type("String")
287            .string_value("base64")
288            .build()
289    {
290        attributes.insert(ENCODING_ATTRIBUTE.to_owned(), marker);
291    }
292    (attributes, group)
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298
299    #[test]
300    fn utf8_payloads_pass_through() {
301        let (body, marker) = encode_body(b"{\"id\":1}");
302        assert_eq!(body, "{\"id\":1}");
303        assert!(!marker);
304    }
305
306    #[test]
307    fn binary_payloads_travel_base64_with_marker() {
308        let raw = [0u8, 159, 146, 150];
309        let (body, marker) = encode_body(&raw);
310        assert!(marker);
311        assert_eq!(BASE64.decode(body).expect("valid base64"), raw);
312    }
313
314    #[test]
315    fn partition_key_header_becomes_the_group_id() {
316        let mut headers = Headers::new();
317        headers.insert(PARTITION_KEY_HEADER, "user-42");
318        headers.insert("x-tenant", "acme");
319        let (attributes, group) = encode_attributes(&headers, false);
320        assert_eq!(group.as_deref(), Some("user-42"));
321        assert!(attributes.contains_key("x-tenant"));
322        assert!(!attributes.contains_key(PARTITION_KEY_HEADER));
323    }
324
325    /// A client built from a bare config: no network happens until an operation is sent, and
326    /// this test never sends one.
327    fn offline_client() -> Client {
328        let config = aws_config::SdkConfig::builder()
329            .behavior_version(aws_config::BehaviorVersion::latest())
330            .region(aws_config::Region::new("us-east-1"))
331            .build();
332        Client::new(&config)
333    }
334
335    /// The runtime picks the native path off this flag, so a delivery that can change its own
336    /// visibility has to report it; without it `retry_after` silently falls back to the
337    /// deferred re-publish.
338    #[tokio::test]
339    async fn deliveries_advertise_native_delayed_redelivery() {
340        let raw = AwsMessage::builder()
341            .body("{}")
342            .receipt_handle("receipt")
343            .build();
344        let message = SqsMessage::new(
345            &raw,
346            offline_client(),
347            "http://localhost:4566/000000000000/queue".to_owned(),
348            "receipt".to_owned(),
349            Duration::from_secs(30),
350        );
351        assert!(message.supports_nack_after());
352    }
353}
ruststream_sqs_sns/message.rs

ruststream_sqs_sns/
message.rs
