iggy_binary_protocol/responses/messages/
poll_messages.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Decoder for the `PollMessages` response.
19//!
20//! The body after the 16-byte prefix is a stream of canonical batch records
21//! (see [`crate::batch`]) served as stored: each record's header carries the
22//! stamped `base_offset` / `base_timestamp`, and each frame's deltas resolve
23//! against them. A record may be a server-sliced view of a larger stored
24//! batch, so `base_offset + offset_delta` of the first frame is the first
25//! polled offset, not necessarily `base_offset` itself.
26
27use crate::batch::{BatchHeader, BatchIntegrity, BatchIterator, decode_batch_slice_with};
28use crate::codec::{WireDecode, WireEncode, read_u32_le, read_u64_le};
29use crate::error::WireError;
30use bytes::{BufMut, BytesMut};
31
32/// Size of the `PollMessages` response header: `partition_id(4) + current_offset(8) + count(4)`.
33const POLL_RESPONSE_HEADER_SIZE: usize = 16;
34
35/// The 16-byte metadata prefix of a `PollMessages` response.
36///
37/// Layout: `partition_id(4) + current_offset(8) + messages_count(4)`.
38#[derive(Debug, Clone, PartialEq, Eq)]
39pub struct PollMessagesResponseHeader {
40    pub partition_id: u32,
41    pub current_offset: u64,
42    pub messages_count: u32,
43}
44
45impl WireEncode for PollMessagesResponseHeader {
46    fn encoded_size(&self) -> usize {
47        POLL_RESPONSE_HEADER_SIZE
48    }
49
50    fn encode(&self, buf: &mut BytesMut) {
51        buf.put_u32_le(self.partition_id);
52        buf.put_u64_le(self.current_offset);
53        buf.put_u32_le(self.messages_count);
54    }
55}
56
57impl WireDecode for PollMessagesResponseHeader {
58    fn decode(buf: &[u8]) -> Result<(Self, usize), WireError> {
59        let partition_id = read_u32_le(buf, 0)?;
60        let current_offset = read_u64_le(buf, 4)?;
61        let messages_count = read_u32_le(buf, 12)?;
62        Ok((
63            Self {
64                partition_id,
65                current_offset,
66                messages_count,
67            },
68            POLL_RESPONSE_HEADER_SIZE,
69        ))
70    }
71}
72
73/// One polled message with its deltas resolved to absolute values.
74#[derive(Debug, Clone, Copy)]
75pub struct PolledMessageView<'a> {
76    /// Stored per-message checksum, passed through as served.
77    pub checksum: u64,
78    pub id: u128,
79    pub offset: u64,
80    /// Broker append time: the flat batch `base_timestamp` (the per-message
81    /// delta applies to `origin_timestamp` only).
82    pub timestamp: u64,
83    pub origin_timestamp: u64,
84    pub payload: &'a [u8],
85    pub user_headers: &'a [u8],
86}
87
88/// Iterator over every message in a stream of served batch records.
89///
90/// Walks records by `batch_length` and flattens their frames. Each record's
91/// layout is proven up front (frames must tile `message_count` exactly), so
92/// framing errors, including frame-level corruption inside a record, surface
93/// as `Some(Err(_))` and end the iteration. Checksums are passed through
94/// unverified, as served.
95pub struct PolledBatchesIterator<'a> {
96    batches: &'a [u8],
97    position: usize,
98    current_header: Option<BatchHeader>,
99    frames: Option<BatchIterator<'a>>,
100    failed: bool,
101}
102
103impl<'a> PolledBatchesIterator<'a> {
104    #[must_use]
105    pub const fn new(batches: &'a [u8]) -> Self {
106        Self {
107            batches,
108            position: 0,
109            current_header: None,
110            frames: None,
111            failed: false,
112        }
113    }
114
115    fn advance_batch(&mut self) -> Result<bool, WireError> {
116        if self.position >= self.batches.len() {
117            return Ok(false);
118        }
119        // Layout-only decode proves the frames tile `message_count` exactly
120        // before any of them are yielded, so a corrupt record errors instead
121        // of silently truncating. Poll replies pass checksums through as
122        // served, so no integrity verification here.
123        let batch =
124            decode_batch_slice_with(&self.batches[self.position..], BatchIntegrity::LayoutOnly)?;
125        self.position += batch.header.total_size();
126        self.frames = Some(batch.iter());
127        self.current_header = Some(batch.header);
128        Ok(true)
129    }
130}
131
132impl<'a> Iterator for PolledBatchesIterator<'a> {
133    type Item = Result<PolledMessageView<'a>, WireError>;
134
135    fn next(&mut self) -> Option<Self::Item> {
136        if self.failed {
137            return None;
138        }
139        loop {
140            if let Some(frames) = &mut self.frames
141                && let Some(view) = frames.next()
142            {
143                let header = self.current_header.expect("frames imply a current header");
144                return Some(Ok(PolledMessageView {
145                    checksum: view.header.checksum,
146                    id: view.header.id,
147                    offset: header.base_offset + u64::from(view.header.offset_delta),
148                    timestamp: header.base_timestamp,
149                    origin_timestamp: header.origin_timestamp
150                        + u64::from(view.header.timestamp_delta),
151                    payload: view.payload,
152                    user_headers: view.user_headers,
153                }));
154            }
155            match self.advance_batch() {
156                Ok(true) => {}
157                Ok(false) => return None,
158                Err(error) => {
159                    self.failed = true;
160                    return Some(Err(error));
161                }
162            }
163        }
164    }
165}
166
167/// Borrowed `PollMessages` response. Does not own message data.
168///
169/// Does NOT implement `WireDecode` (trait returns owned data, we borrow).
170/// Use [`PollMessagesResponse::decode`] instead.
171pub struct PollMessagesResponse<'a> {
172    pub header: PollMessagesResponseHeader,
173    pub messages: PolledBatchesIterator<'a>,
174}
175
176impl<'a> PollMessagesResponse<'a> {
177    /// Decode from a response payload buffer. Borrows the buffer.
178    ///
179    /// Reads the 16-byte header then creates an iterator over the batch
180    /// records that follow. Records are validated lazily during iteration.
181    ///
182    /// # Errors
183    /// Returns `WireError` if the buffer is too short for the response header.
184    pub fn decode(buf: &'a [u8]) -> Result<Self, WireError> {
185        let (header, _) = PollMessagesResponseHeader::decode(buf)?;
186        let messages = PolledBatchesIterator::new(&buf[POLL_RESPONSE_HEADER_SIZE..]);
187        Ok(Self { header, messages })
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194    use crate::batch::{BATCH_HEADER_SIZE, BATCH_MESSAGE_HEADER_SIZE, calculate_batch_checksum};
195    use twox_hash::XxHash3_64;
196
197    #[allow(clippy::cast_possible_truncation)]
198    fn frame(id: u128, offset_delta: u32, timestamp_delta: u32, payload: &[u8]) -> Vec<u8> {
199        let mut bytes = vec![0u8; BATCH_MESSAGE_HEADER_SIZE];
200        bytes[8..24].copy_from_slice(&id.to_le_bytes());
201        bytes[24..28].copy_from_slice(&offset_delta.to_le_bytes());
202        bytes[28..32].copy_from_slice(&timestamp_delta.to_le_bytes());
203        bytes[36..40].copy_from_slice(&(payload.len() as u32).to_le_bytes());
204        bytes.extend_from_slice(payload);
205        let checksum = XxHash3_64::oneshot(&bytes[8..]);
206        bytes[0..8].copy_from_slice(&checksum.to_le_bytes());
207        bytes
208    }
209
210    #[allow(clippy::cast_possible_truncation)]
211    fn batch_record(
212        base_offset: u64,
213        base_timestamp: u64,
214        origin_timestamp: u64,
215        frames: &[Vec<u8>],
216    ) -> Vec<u8> {
217        let blob: Vec<u8> = frames.concat();
218        let mut header = BatchHeader::new(
219            1,
220            origin_timestamp,
221            (BATCH_HEADER_SIZE + blob.len()) as u64,
222            frames.len() as u32,
223        );
224        header.base_offset = base_offset;
225        header.base_timestamp = base_timestamp;
226        header.batch_checksum = calculate_batch_checksum(&header, &blob);
227        let mut bytes = vec![0u8; BATCH_HEADER_SIZE];
228        header.encode_into(&mut bytes);
229        bytes.extend_from_slice(&blob);
230        bytes
231    }
232
233    fn response_body(current_offset: u64, count: u32, batches: &[Vec<u8>]) -> Vec<u8> {
234        let mut body = Vec::new();
235        body.extend_from_slice(&7u32.to_le_bytes());
236        body.extend_from_slice(&current_offset.to_le_bytes());
237        body.extend_from_slice(&count.to_le_bytes());
238        for batch in batches {
239            body.extend_from_slice(batch);
240        }
241        body
242    }
243
244    #[test]
245    fn response_header_roundtrip() {
246        let header = PollMessagesResponseHeader {
247            partition_id: 3,
248            current_offset: 42,
249            messages_count: 7,
250        };
251        let bytes = header.to_bytes();
252        let (decoded, consumed) = PollMessagesResponseHeader::decode(&bytes).unwrap();
253        assert_eq!(consumed, bytes.len());
254        assert_eq!(decoded, header);
255    }
256
257    #[test]
258    fn response_header_truncation() {
259        let bytes = [0u8; POLL_RESPONSE_HEADER_SIZE - 1];
260        assert!(PollMessagesResponseHeader::decode(&bytes).is_err());
261    }
262
263    #[test]
264    fn decodes_messages_across_batches() {
265        let first = batch_record(
266            100,
267            5_000,
268            1_000,
269            &[frame(11, 0, 0, b"a"), frame(12, 1, 10, b"b")],
270        );
271        let second = batch_record(102, 6_000, 2_000, &[frame(13, 0, 0, b"c")]);
272        let body = response_body(102, 3, &[first, second]);
273
274        let response = PollMessagesResponse::decode(&body).unwrap();
275        assert_eq!(response.header.partition_id, 7);
276        assert_eq!(response.header.current_offset, 102);
277        assert_eq!(response.header.messages_count, 3);
278
279        let messages: Vec<_> = response.messages.collect::<Result<Vec<_>, _>>().unwrap();
280        assert_eq!(messages.len(), 3);
281
282        assert_eq!(messages[0].id, 11);
283        assert_eq!(messages[0].offset, 100);
284        assert_eq!(messages[0].timestamp, 5_000);
285        assert_eq!(messages[0].origin_timestamp, 1_000);
286        assert_eq!(messages[0].payload, b"a");
287
288        assert_eq!(messages[1].id, 12);
289        assert_eq!(messages[1].offset, 101);
290        assert_eq!(messages[1].timestamp, 5_000);
291        assert_eq!(messages[1].origin_timestamp, 1_010);
292        assert_eq!(messages[1].payload, b"b");
293
294        assert_eq!(messages[2].id, 13);
295        assert_eq!(messages[2].offset, 102);
296        assert_eq!(messages[2].timestamp, 6_000);
297        assert_eq!(messages[2].payload, b"c");
298    }
299
300    #[test]
301    fn sliced_record_resolves_leading_delta() {
302        // A server-sliced record keeps the stored base_offset; the first
303        // frame's delta positions it inside the original batch.
304        let record = batch_record(50, 9_000, 0, &[frame(1, 3, 0, b"tail")]);
305        let body = response_body(53, 1, &[record]);
306
307        let response = PollMessagesResponse::decode(&body).unwrap();
308        let messages: Vec<_> = response.messages.collect::<Result<Vec<_>, _>>().unwrap();
309        assert_eq!(messages.len(), 1);
310        assert_eq!(messages[0].offset, 53);
311    }
312
313    #[test]
314    fn empty_body_yields_no_messages() {
315        let body = response_body(0, 0, &[]);
316        let response = PollMessagesResponse::decode(&body).unwrap();
317        assert_eq!(response.messages.count(), 0);
318    }
319
320    #[test]
321    fn truncated_record_surfaces_error() {
322        let record = batch_record(0, 0, 0, &[frame(1, 0, 0, b"payload")]);
323        let mut body = response_body(0, 1, &[record]);
324        body.truncate(body.len() - 1);
325
326        let response = PollMessagesResponse::decode(&body).unwrap();
327        let result: Result<Vec<_>, _> = response.messages.collect();
328        assert!(result.is_err());
329    }
330}
iggy_binary_protocol/responses/messages/poll_messages.rs

iggy_binary_protocol/responses/messages/
poll_messages.rs
