rocketmq_protocol/protocol/body/
batch_ack_builder.rs

1// Copyright 2023 The RocketMQ Rust Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::BTreeMap;
16
17use bitvec::prelude::BitVec;
18use bitvec::prelude::Lsb0;
19use cheetah_string::CheetahString;
20use rocketmq_error::Error;
21
22use crate::protocol::body::batch_ack::BatchAck;
23use crate::protocol::body::batch_ack::SerializableBitVec;
24use crate::protocol::body::batch_ack_message_request_body::BatchAckMessageRequestBody;
25use crate::protocol::header::extra_info_util::AckExtraInfo;
26use crate::protocol::header::extra_info_util::ExtraInfoUtil;
27
28pub const DEFAULT_MAX_BATCH_ACK_ENTRIES: usize = 4096;
29pub const DEFAULT_MAX_BATCH_ACK_OFFSET_SPAN: usize = 1_048_576;
30
31#[derive(Debug, Clone, Copy)]
32pub struct BatchAckBuildLimits {
33    pub max_entries: usize,
34    pub max_offset_span: usize,
35}
36
37impl Default for BatchAckBuildLimits {
38    fn default() -> Self {
39        Self {
40            max_entries: DEFAULT_MAX_BATCH_ACK_ENTRIES,
41            max_offset_span: DEFAULT_MAX_BATCH_ACK_OFFSET_SPAN,
42        }
43    }
44}
45
46#[derive(Debug, Clone, Copy)]
47pub struct BatchAckInput<'a> {
48    pub entry_index: usize,
49    pub consumer_group: &'a str,
50    pub topic: &'a str,
51    pub receipt_handle: &'a str,
52}
53
54#[derive(Debug)]
55pub struct BatchAckBuildFailure {
56    pub entry_index: usize,
57    pub error: Error,
58}
59
60#[derive(Debug)]
61pub struct BatchAckBrokerRequest {
62    pub broker_name: CheetahString,
63    pub body: BatchAckMessageRequestBody,
64    pub entry_indexes: Vec<usize>,
65}
66
67#[derive(Debug, Default)]
68pub struct BatchAckBuildResult {
69    pub requests: Vec<BatchAckBrokerRequest>,
70    pub failures: Vec<BatchAckBuildFailure>,
71}
72
73#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
74struct CheckpointKey {
75    consumer_group: String,
76    topic: String,
77    retry: String,
78    start_offset: i64,
79    queue_id: i32,
80    revive_queue_id: i32,
81    pop_time: i64,
82    invisible_time: i64,
83}
84
85struct PendingAck {
86    ack: BatchAck,
87    entry_indexes: Vec<usize>,
88}
89
90pub fn build_batch_ack_requests(inputs: &[BatchAckInput<'_>]) -> BatchAckBuildResult {
91    build_batch_ack_requests_with_limits(inputs, BatchAckBuildLimits::default())
92}
93
94pub fn build_batch_ack_requests_with_limits(
95    inputs: &[BatchAckInput<'_>],
96    limits: BatchAckBuildLimits,
97) -> BatchAckBuildResult {
98    let mut result = BatchAckBuildResult::default();
99    let mut brokers: BTreeMap<String, BTreeMap<CheckpointKey, PendingAck>> = BTreeMap::new();
100
101    for (position, input) in inputs.iter().enumerate() {
102        if position >= limits.max_entries {
103            result.failures.push(BatchAckBuildFailure {
104                entry_index: input.entry_index,
105                error: crate::error::invalid_argument(format!(
106                    "batch ACK input exceeds the configured {} entry limit",
107                    limits.max_entries
108                )),
109            });
110            continue;
111        }
112
113        match parse_input(*input, limits.max_offset_span) {
114            Ok((extra, bit_index)) => {
115                let key = CheckpointKey {
116                    consumer_group: input.consumer_group.to_owned(),
117                    topic: input.topic.to_owned(),
118                    retry: extra.retry.to_owned(),
119                    start_offset: extra.ck_queue_offset,
120                    queue_id: extra.queue_id,
121                    revive_queue_id: extra.revive_queue_id,
122                    pop_time: extra.pop_time,
123                    invisible_time: extra.invisible_time,
124                };
125                let pending = brokers
126                    .entry(extra.broker_name.to_owned())
127                    .or_default()
128                    .entry(key)
129                    .or_insert_with(|| PendingAck {
130                        ack: BatchAck {
131                            consumer_group: CheetahString::from(input.consumer_group),
132                            topic: CheetahString::from(input.topic),
133                            retry: CheetahString::from(extra.retry),
134                            start_offset: extra.ck_queue_offset,
135                            queue_id: extra.queue_id,
136                            revive_queue_id: extra.revive_queue_id,
137                            pop_time: extra.pop_time,
138                            invisible_time: extra.invisible_time,
139                            bit_set: SerializableBitVec(BitVec::<u64, Lsb0>::new()),
140                        },
141                        entry_indexes: Vec::new(),
142                    });
143                if pending.ack.bit_set.0.len() <= bit_index {
144                    pending.ack.bit_set.0.resize(bit_index + 1, false);
145                }
146                pending.ack.bit_set.0.set(bit_index, true);
147                pending.entry_indexes.push(input.entry_index);
148            }
149            Err(error) => result.failures.push(BatchAckBuildFailure {
150                entry_index: input.entry_index,
151                error,
152            }),
153        }
154    }
155
156    result.requests = brokers
157        .into_iter()
158        .map(|(broker_name, checkpoints)| {
159            let mut acks = Vec::with_capacity(checkpoints.len());
160            let mut entry_indexes = Vec::new();
161            for pending in checkpoints.into_values() {
162                acks.push(pending.ack);
163                entry_indexes.extend(pending.entry_indexes);
164            }
165            let broker_name = CheetahString::from_string(broker_name);
166            BatchAckBrokerRequest {
167                body: BatchAckMessageRequestBody {
168                    broker_name: broker_name.clone(),
169                    acks,
170                },
171                broker_name,
172                entry_indexes,
173            }
174        })
175        .collect();
176    result
177}
178
179fn parse_input<'a>(input: BatchAckInput<'a>, max_offset_span: usize) -> Result<(AckExtraInfo<'a>, usize), Error> {
180    if input.consumer_group.trim().is_empty() {
181        return Err(crate::error::invalid_argument("batch ACK consumer group is blank"));
182    }
183    if input.topic.trim().is_empty() {
184        return Err(crate::error::invalid_argument("batch ACK topic is blank"));
185    }
186
187    let extra = ExtraInfoUtil::parse_ack_extra_info(input.receipt_handle)?;
188    let delta = extra
189        .queue_offset
190        .checked_sub(extra.ck_queue_offset)
191        .ok_or_else(|| crate::error::invalid_argument("batch ACK queue offset overflow"))?;
192    let bit_index = usize::try_from(delta)
193        .map_err(|_| crate::error::invalid_argument("batch ACK queue offset precedes checkpoint offset"))?;
194    if bit_index >= max_offset_span {
195        return Err(crate::error::invalid_argument(format!(
196            "batch ACK offset span {bit_index} exceeds the configured {max_offset_span} bit limit"
197        )));
198    }
199    Ok((extra, bit_index))
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205    use crate::protocol::header::extra_info_util::ExtraInfoUtil;
206
207    fn receipt(start: i64, broker: &str, queue_id: i32, offset: i64) -> String {
208        ExtraInfoUtil::build_extra_info_with_offset(start, 100, 15_000, 0, "topic", broker, queue_id, offset)
209    }
210
211    #[test]
212    fn blank_groups_and_topics_preserve_failure_indexes() {
213        let handle = receipt(10, "broker", 1, 10);
214        let inputs = [("", "topic"), (" \t", "topic"), ("group", ""), ("group", "\n ")]
215            .into_iter()
216            .enumerate()
217            .map(|(index, (consumer_group, topic))| BatchAckInput {
218                entry_index: index + 10,
219                consumer_group,
220                topic,
221                receipt_handle: &handle,
222            })
223            .collect::<Vec<_>>();
224        let result = build_batch_ack_requests(&inputs);
225        assert!(result.requests.is_empty());
226        assert_eq!(
227            result
228                .failures
229                .iter()
230                .map(|failure| failure.entry_index)
231                .collect::<Vec<_>>(),
232            vec![10, 11, 12, 13]
233        );
234        for failure in result.failures {
235            assert_eq!(failure.error.code(), rocketmq_error::CORE_ARGUMENT_INVALID.code());
236        }
237    }
238
239    #[test]
240    fn entry_limit_keeps_the_first_request_and_rejects_every_later_input() {
241        let handle = receipt(10, "broker", 1, 10);
242        let inputs = [handle.as_str(), handle.as_str(), "malformed"]
243            .into_iter()
244            .enumerate()
245            .map(|(index, receipt_handle)| BatchAckInput {
246                entry_index: index + 20,
247                consumer_group: "group",
248                topic: "topic",
249                receipt_handle,
250            })
251            .collect::<Vec<_>>();
252        let result = build_batch_ack_requests_with_limits(
253            &inputs,
254            BatchAckBuildLimits {
255                max_entries: 1,
256                ..Default::default()
257            },
258        );
259        assert_eq!(result.requests.len(), 1);
260        assert_eq!(result.requests[0].entry_indexes, vec![20]);
261        assert_eq!(result.requests[0].body.acks[0].bit_set.0.count_ones(), 1);
262        assert_eq!(
263            result
264                .failures
265                .iter()
266                .map(|failure| failure.entry_index)
267                .collect::<Vec<_>>(),
268            vec![21, 22]
269        );
270        for failure in result.failures {
271            assert_eq!(failure.error.code(), rocketmq_error::CORE_ARGUMENT_INVALID.code());
272        }
273    }
274
275    #[test]
276    fn invalid_offset_arithmetic_does_not_discard_valid_requests() {
277        let handles = [
278            receipt(10, "broker", 1, 9),
279            receipt(10, "broker", 1, 11),
280            receipt(i64::MIN, "broker", 1, i64::MAX),
281        ];
282        let inputs = handles
283            .iter()
284            .zip([31, 7, 12])
285            .map(|(receipt_handle, entry_index)| BatchAckInput {
286                entry_index,
287                consumer_group: "group",
288                topic: "topic",
289                receipt_handle,
290            })
291            .collect::<Vec<_>>();
292        let result = build_batch_ack_requests(&inputs);
293        assert_eq!(result.requests.len(), 1);
294        assert_eq!(result.requests[0].broker_name, "broker");
295        assert_eq!(result.requests[0].entry_indexes, vec![7]);
296        assert_eq!(
297            result.requests[0].body.acks[0]
298                .bit_set
299                .0
300                .iter()
301                .by_vals()
302                .collect::<Vec<_>>(),
303            vec![false, true]
304        );
305        assert_eq!(
306            result
307                .failures
308                .iter()
309                .map(|failure| failure.entry_index)
310                .collect::<Vec<_>>(),
311            vec![31, 12]
312        );
313        for failure in result.failures {
314            assert_eq!(failure.error.code(), rocketmq_error::CORE_ARGUMENT_INVALID.code());
315        }
316    }
317
318    #[test]
319    fn thirty_two_offsets_share_one_broker_request_and_checkpoint_bitmap() {
320        let receipts = (0..32)
321            .map(|offset| receipt(10, "broker-a", 1, 10 + offset))
322            .collect::<Vec<_>>();
323        let inputs = receipts
324            .iter()
325            .enumerate()
326            .map(|(entry_index, receipt_handle)| BatchAckInput {
327                entry_index,
328                consumer_group: "group",
329                topic: "topic",
330                receipt_handle,
331            })
332            .collect::<Vec<_>>();
333
334        let built = build_batch_ack_requests(&inputs);
335        assert!(built.failures.is_empty());
336        assert_eq!(built.requests.len(), 1);
337        assert_eq!(built.requests[0].body.acks.len(), 1);
338        assert_eq!(built.requests[0].entry_indexes.len(), 32);
339        assert!(built.requests[0].body.acks[0].bit_set.0.iter().by_vals().all(|bit| bit));
340    }
341
342    #[test]
343    fn broker_and_checkpoint_identity_are_grouped_without_mixing() {
344        let receipts = [
345            receipt(10, "broker-a", 1, 10),
346            receipt(20, "broker-a", 1, 20),
347            receipt(10, "broker-b", 1, 11),
348        ];
349        let inputs = receipts
350            .iter()
351            .enumerate()
352            .map(|(entry_index, receipt_handle)| BatchAckInput {
353                entry_index,
354                consumer_group: "group",
355                topic: "topic",
356                receipt_handle,
357            })
358            .collect::<Vec<_>>();
359
360        let built = build_batch_ack_requests(&inputs);
361        assert!(built.failures.is_empty());
362        assert_eq!(built.requests.len(), 2);
363        assert_eq!(built.requests[0].body.acks.len(), 2);
364        assert_eq!(built.requests[1].body.acks.len(), 1);
365    }
366
367    #[test]
368    fn malformed_and_huge_spans_are_failures_without_large_bitmaps() {
369        let huge = receipt(0, "broker-a", 1, i64::MAX);
370        let inputs = [
371            BatchAckInput {
372                entry_index: 3,
373                consumer_group: "group",
374                topic: "topic",
375                receipt_handle: "bad",
376            },
377            BatchAckInput {
378                entry_index: 7,
379                consumer_group: "group",
380                topic: "topic",
381                receipt_handle: &huge,
382            },
383        ];
384
385        let built = build_batch_ack_requests_with_limits(
386            &inputs,
387            BatchAckBuildLimits {
388                max_entries: 32,
389                max_offset_span: 64,
390            },
391        );
392        assert!(built.requests.is_empty());
393        assert_eq!(
394            built
395                .failures
396                .iter()
397                .map(|failure| failure.entry_index)
398                .collect::<Vec<_>>(),
399            vec![3, 7]
400        );
401    }
402
403    #[test]
404    fn duplicate_offsets_preserve_all_entry_indexes() {
405        let receipt = receipt(10, "broker-a", 1, 10);
406        let inputs = [
407            BatchAckInput {
408                entry_index: 1,
409                consumer_group: "group",
410                topic: "topic",
411                receipt_handle: &receipt,
412            },
413            BatchAckInput {
414                entry_index: 2,
415                consumer_group: "group",
416                topic: "topic",
417                receipt_handle: &receipt,
418            },
419        ];
420        let built = build_batch_ack_requests(&inputs);
421        assert_eq!(built.requests[0].entry_indexes, vec![1, 2]);
422        assert_eq!(built.requests[0].body.acks[0].bit_set.0.count_ones(), 1);
423    }
424}
rocketmq_protocol/protocol/body/batch_ack_builder.rs

rocketmq_protocol/protocol/body/
batch_ack_builder.rs
