krabka_client_consumer/
builder.rs

1//! Codec helpers for the `ConsumerProtocol` subscription and assignment
2//! payloads, plus the [`AutoOffsetReset`] and [`IsolationLevel`] enums that
3//! [`Consumer::builder`] uses.
4
5use std::{collections::BTreeMap, str::FromStr, time::Duration};
6
7use bytes::{Bytes, BytesMut};
8
9/// What to do when a partition has no committed offset. Kafka's
10/// `auto.offset.reset`.
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum AutoOffsetReset {
13    /// Start from the log-start offset. `Consumer::poll` resolves it lazily
14    /// with `ListOffsets(timestamp=-2)`, so a partition never reports offset 0
15    /// before the broker's real log start is known.
16    Earliest,
17    /// Start from the log-end offset. `Consumer::poll` resolves it lazily with
18    /// `ListOffsets(timestamp=-1)`.
19    Latest,
20    /// Do not reset automatically. On a missing committed offset, `poll` and
21    /// `position` return `ConsumerError::NoOffsetForPartition` naming the
22    /// partitions, instead of synthesizing a reset. On a detected truncation,
23    /// `poll` returns `ConsumerError::LogTruncation`.
24    None,
25    /// Start from the first offset whose timestamp is at or after now minus
26    /// this duration (KIP-1106, `by_duration:<ISO-8601 duration>`).
27    /// `Consumer::poll` resolves it with `ListOffsets(timestamp)`.
28    ByDuration(Duration),
29}
30
31impl FromStr for AutoOffsetReset {
32    type Err = String;
33
34    /// Parse Kafka's `auto.offset.reset` values: `earliest`, `latest`, `none`
35    /// and `by_duration:<ISO-8601 duration>`, as
36    /// `AutoOffsetResetStrategy.fromString` does.
37    fn from_str(value: &str) -> Result<Self, Self::Err> {
38        match value {
39            "earliest" => Ok(Self::Earliest),
40            "latest" => Ok(Self::Latest),
41            "none" => Ok(Self::None),
42            "by_duration" => {
43                Err("<:duration> part is missing in by_duration auto offset reset strategy.".into())
44            }
45            _ => value.strip_prefix("by_duration:").map_or_else(
46                || Err(format!("Unknown auto offset reset strategy: {value}")),
47                |duration| parse_iso_duration(duration).map(Self::ByDuration),
48            ),
49        }
50    }
51}
52
53/// Split the text before `unit` from the front of `part`.
54fn duration_component(part: &str, unit: char) -> (Option<&str>, &str) {
55    part.split_once(unit)
56        .map_or((None, part), |(number, rest)| (Some(number), rest))
57}
58
59/// A decimal integer with an optional sign that fits `i64`, as Java's
60/// `Long.parseLong` in `Duration.parse` requires.
61fn duration_integer(number: &str) -> Option<i128> {
62    let digits = number.strip_prefix(['-', '+']).unwrap_or(number);
63    if digits.is_empty() || !digits.bytes().all(|byte| byte.is_ascii_digit()) {
64        return None;
65    }
66    number.parse::<i64>().ok().map(i128::from)
67}
68
69/// `value * unit_seconds` seconds in nanoseconds, added to `nanos`, or `None`
70/// on overflow.
71fn add_duration_part(nanos: i128, value: i128, unit_seconds: i128) -> Option<i128> {
72    const NANOS_PER_SECOND: i128 = 1_000_000_000;
73    let part = value
74        .checked_mul(unit_seconds)?
75        .checked_mul(NANOS_PER_SECOND)?;
76    nanos.checked_add(part)
77}
78
79/// Parse an ISO-8601 duration as Java's `Duration.parse` does:
80/// `[-+]P[nD][T[nH][nM][n[.fffffffff]S]]`, case-insensitive, each number with
81/// an optional sign. `by_duration` rejects a negative result.
82fn parse_iso_duration(text: &str) -> Result<Duration, String> {
83    const NANOS_PER_SECOND: i128 = 1_000_000_000;
84    let invalid =
85        || "Unable to parse duration string in by_duration offset reset strategy.".to_owned();
86    let upper = text.to_ascii_uppercase();
87    let (negative, rest) = upper.strip_prefix('-').map_or_else(
88        || (false, upper.strip_prefix('+').unwrap_or(upper.as_str())),
89        |rest| (true, rest),
90    );
91    let rest = rest.strip_prefix('P').ok_or_else(invalid)?;
92    let (date, time) = match rest.split_once('T') {
93        Some((date, time)) => (date, Some(time)),
94        None => (rest, None),
95    };
96    let mut nanos: i128 = 0;
97    let mut components = 0;
98    let (days, date_rest) = duration_component(date, 'D');
99    if !date_rest.is_empty() {
100        return Err(invalid());
101    }
102    if let Some(days) = days {
103        nanos = add_duration_part(nanos, duration_integer(days).ok_or_else(invalid)?, 86_400)
104            .ok_or_else(invalid)?;
105        components += 1;
106    }
107    if let Some(time) = time {
108        let (hours, time) = duration_component(time, 'H');
109        let (minutes, time) = duration_component(time, 'M');
110        let (seconds, time) = duration_component(time, 'S');
111        if !time.is_empty() || (hours.is_none() && minutes.is_none() && seconds.is_none()) {
112            return Err(invalid());
113        }
114        if let Some(hours) = hours {
115            nanos = add_duration_part(nanos, duration_integer(hours).ok_or_else(invalid)?, 3_600)
116                .ok_or_else(invalid)?;
117        }
118        if let Some(minutes) = minutes {
119            nanos = add_duration_part(nanos, duration_integer(minutes).ok_or_else(invalid)?, 60)
120                .ok_or_else(invalid)?;
121        }
122        if let Some(seconds) = seconds {
123            let (whole, fraction) = match seconds.split_once(['.', ',']) {
124                Some((whole, fraction)) => (whole, fraction),
125                None => (seconds, ""),
126            };
127            if fraction.len() > 9 || !fraction.bytes().all(|byte| byte.is_ascii_digit()) {
128                return Err(invalid());
129            }
130            let whole_value = duration_integer(whole).ok_or_else(invalid)?;
131            let fraction_value = if fraction.is_empty() {
132                0
133            } else {
134                format!("{fraction:0<9}")
135                    .parse::<i128>()
136                    .map_err(|_| invalid())?
137            };
138            let sign = if whole.starts_with('-') { -1 } else { 1 };
139            nanos = add_duration_part(nanos, whole_value, 1)
140                .and_then(|nanos| nanos.checked_add(sign * fraction_value))
141                .ok_or_else(invalid)?;
142        }
143        components += 1;
144    }
145    if components == 0 {
146        return Err(invalid());
147    }
148    if negative {
149        nanos = -nanos;
150    }
151    if nanos < 0 {
152        return Err(
153            "Negative duration is not supported in by_duration offset reset strategy.".into(),
154        );
155    }
156    // Java's `Duration` holds the seconds in a `long`.
157    let seconds = i64::try_from(nanos / NANOS_PER_SECOND).map_err(|_| invalid())?;
158    let seconds = u64::try_from(seconds).map_err(|_| invalid())?;
159    let subsec = u32::try_from(nanos % NANOS_PER_SECOND).map_err(|_| invalid())?;
160    Ok(Duration::new(seconds, subsec))
161}
162
163/// Controls which records are visible to this consumer.
164///
165/// This maps to Kafka's `isolation.level` configuration and to the
166/// `isolation_level` field in the `Fetch` request, whose wire value is an `i8`.
167///
168/// The default is [`ReadUncommitted`](IsolationLevel::ReadUncommitted) for
169/// backward compatibility.
170#[derive(Debug, Clone, Copy, PartialEq, Eq)]
171pub enum IsolationLevel {
172    /// All records are visible, including those from open or aborted
173    /// transactions. Equivalent to `isolation.level=read_uncommitted`.
174    ReadUncommitted,
175    /// Only records from committed transactions, plus non-transactional
176    /// records, are visible. Equivalent to `isolation.level=read_committed`.
177    ReadCommitted,
178}
179
180impl FromStr for IsolationLevel {
181    type Err = String;
182
183    fn from_str(value: &str) -> Result<Self, Self::Err> {
184        // Kafka's `ConsumerConfig` accepts exactly these `isolation.level`
185        // values.
186        match value {
187            "read_uncommitted" => Ok(Self::ReadUncommitted),
188            "read_committed" => Ok(Self::ReadCommitted),
189            _ => Err(format!("invalid isolation level: {value}")),
190        }
191    }
192}
193
194impl IsolationLevel {
195    /// Returns the `i8` wire encoding that the `Fetch` request uses.
196    pub(crate) fn wire(self) -> i8 {
197        match self {
198            Self::ReadUncommitted => 0,
199            Self::ReadCommitted => 1,
200        }
201    }
202}
203
204// ── subscription / assignment codec (ConsumerProtocol v3) ─────────────────
205
206use krabka_protocol::{
207    Decode, Encode, ProtocolError, UnknownTaggedFields,
208    owned::{
209        consumer_protocol_assignment::{
210            ConsumerProtocolAssignment, TopicPartition as AssignTopicPartition,
211        },
212        consumer_protocol_subscription::{
213            ConsumerProtocolSubscription, TopicPartition as SubTopicPartition,
214        },
215    },
216};
217
218const SUBSCRIPTION_WIRE_VERSION: i16 = 3;
219const ASSIGNMENT_WIRE_VERSION: i16 = 3;
220
221#[derive(Debug, PartialEq, Eq)]
222pub struct DecodedSubscription {
223    pub topics: Vec<String>,
224    pub owned: Vec<(String, i32)>,
225    pub generation_id: i32,
226    /// Part of the `ConsumerProtocolSubscription` v3 wire surface. No built-in
227    /// assignor reads it yet: rack-aware assignment needs replica racks.
228    pub rack_id: Option<String>,
229    /// The assignor-specific `userData`.
230    pub user_data: Option<Bytes>,
231}
232
233fn group_by_topic(pairs: &[(String, i32)]) -> BTreeMap<&str, Vec<i32>> {
234    let mut by_topic: BTreeMap<&str, Vec<i32>> = BTreeMap::new();
235    for (t, p) in pairs {
236        by_topic.entry(t.as_str()).or_default().push(*p);
237    }
238    by_topic
239}
240
241fn peek_version(bytes: &[u8]) -> i16 {
242    bytes
243        .first_chunk()
244        .map_or(0, |version| i16::from_be_bytes(*version))
245}
246
247fn empty_subscription() -> DecodedSubscription {
248    DecodedSubscription {
249        topics: Vec::new(),
250        owned: Vec::new(),
251        generation_id: -1,
252        rack_id: None,
253        user_data: None,
254    }
255}
256
257/// Encode a `ConsumerProtocolSubscription` v3 as Kafka's
258/// `ConsumerProtocol.serializeSubscription` does: sorted topics, and owned
259/// partitions sorted by topic and partition.
260///
261/// # Errors
262/// Returns an error when the subscription does not encode.
263pub fn encode_subscription(
264    topics: &[String],
265    owned: &[(String, i32)],
266    generation_id: i32,
267    rack_id: Option<&str>,
268    user_data: Option<Bytes>,
269) -> Result<Bytes, ProtocolError> {
270    use bytes::BufMut;
271    let owned_partitions: Vec<SubTopicPartition> = group_by_topic(owned)
272        .into_iter()
273        .map(|(topic, mut partitions)| {
274            partitions.sort_unstable();
275            SubTopicPartition {
276                topic: topic.to_owned(),
277                partitions,
278                unknown_tagged_fields: UnknownTaggedFields::default(),
279            }
280        })
281        .collect();
282    let mut sorted_topics = topics.to_vec();
283    sorted_topics.sort();
284    let msg = ConsumerProtocolSubscription {
285        topics: sorted_topics,
286        user_data,
287        owned_partitions,
288        generation_id,
289        rack_id: rack_id.map(str::to_owned),
290        unknown_tagged_fields: UnknownTaggedFields::default(),
291    };
292    let mut buf = BytesMut::with_capacity(2 + msg.encoded_len(SUBSCRIPTION_WIRE_VERSION));
293    buf.put_i16(SUBSCRIPTION_WIRE_VERSION);
294    msg.encode(&mut buf, SUBSCRIPTION_WIRE_VERSION)?;
295    Ok(buf.freeze())
296}
297
298pub fn decode_subscription(bytes: &[u8]) -> DecodedSubscription {
299    let Some(payload) = bytes.get(2..) else {
300        return empty_subscription();
301    };
302    let version = peek_version(bytes).clamp(0, SUBSCRIPTION_WIRE_VERSION);
303    let mut cur = payload;
304    let Ok(msg) = ConsumerProtocolSubscription::decode(&mut cur, version) else {
305        return empty_subscription();
306    };
307    let mut owned = Vec::new();
308    for tp in msg.owned_partitions {
309        for p in tp.partitions {
310            owned.push((tp.topic.clone(), p));
311        }
312    }
313    DecodedSubscription {
314        topics: msg.topics,
315        owned,
316        generation_id: msg.generation_id,
317        rack_id: msg.rack_id,
318        user_data: msg.user_data,
319    }
320}
321
322/// Encode a `ConsumerProtocolAssignment` v3 of `partitions`.
323///
324/// # Errors
325/// Returns an error when the assignment does not encode.
326pub fn encode_assignment(partitions: &[(String, i32)]) -> Result<Bytes, ProtocolError> {
327    use bytes::BufMut;
328    let assigned_partitions: Vec<AssignTopicPartition> = group_by_topic(partitions)
329        .into_iter()
330        .map(|(topic, partitions)| AssignTopicPartition {
331            topic: topic.to_owned(),
332            partitions,
333            unknown_tagged_fields: UnknownTaggedFields::default(),
334        })
335        .collect();
336    let msg = ConsumerProtocolAssignment {
337        assigned_partitions,
338        user_data: None,
339        unknown_tagged_fields: UnknownTaggedFields::default(),
340    };
341    let mut buf = BytesMut::with_capacity(2 + msg.encoded_len(ASSIGNMENT_WIRE_VERSION));
342    buf.put_i16(ASSIGNMENT_WIRE_VERSION);
343    msg.encode(&mut buf, ASSIGNMENT_WIRE_VERSION)?;
344    Ok(buf.freeze())
345}
346
347pub fn decode_assignment(bytes: &[u8]) -> Vec<(String, i32)> {
348    let Some(payload) = bytes.get(2..) else {
349        return Vec::new();
350    };
351    let version = peek_version(bytes).clamp(0, ASSIGNMENT_WIRE_VERSION);
352    let mut cur = payload;
353    let Ok(msg) = ConsumerProtocolAssignment::decode(&mut cur, version) else {
354        return Vec::new();
355    };
356    let mut out = Vec::new();
357    for tp in msg.assigned_partitions {
358        for p in tp.partitions {
359            out.push((tp.topic.clone(), p));
360        }
361    }
362    out
363}
364
365#[cfg(test)]
366mod tests {
367
368    use super::*;
369
370    #[test]
371    fn consumer_behavior_values_parse_kafka_spellings() {
372        use std::time::Duration;
373        let reset = |value: &str| value.parse::<AutoOffsetReset>();
374        let isolation = |value: &str| value.parse::<IsolationLevel>().map_err(|_| ());
375        let actual = (
376            [
377                reset("earliest"),
378                reset("latest"),
379                reset("none"),
380                reset("by_duration:PT1H"),
381                reset("by_duration:P1DT2H3M4.5S"),
382                reset("by_duration:pt0.000000001s"),
383                reset("by_duration:PT1H-30M"),
384                reset("by_duration:+P2D"),
385                reset("by_duration"),
386                reset("by_duration:-PT1H"),
387                reset("by_duration:PT"),
388                reset("by_duration:P"),
389                reset("by_duration:1H"),
390                reset("by_duration:PT1.0000000001S"),
391                reset("by_duration:P9223372036854775807D"),
392                reset("by_duration:PT9223372036854775808H"),
393                reset("by_duration:P106751991167300DT24H"),
394                reset("EARLIEST"),
395                reset("unknown"),
396            ],
397            [
398                isolation("read_committed"),
399                isolation("read_uncommitted"),
400                isolation("read-committed"),
401                isolation("READ_COMMITTED"),
402            ],
403        );
404        let unparsable =
405            Err("Unable to parse duration string in by_duration offset reset strategy.".to_owned());
406        let expected = (
407            [
408                Ok(AutoOffsetReset::Earliest),
409                Ok(AutoOffsetReset::Latest),
410                Ok(AutoOffsetReset::None),
411                Ok(AutoOffsetReset::ByDuration(Duration::from_hours(1))),
412                Ok(AutoOffsetReset::ByDuration(Duration::from_millis(
413                    ((24 + 2) * 3600 + 3 * 60 + 4) * 1000 + 500,
414                ))),
415                Ok(AutoOffsetReset::ByDuration(Duration::from_nanos(1))),
416                Ok(AutoOffsetReset::ByDuration(Duration::from_mins(30))),
417                Ok(AutoOffsetReset::ByDuration(Duration::from_hours(48))),
418                Err(
419                    "<:duration> part is missing in by_duration auto offset reset strategy."
420                        .to_owned(),
421                ),
422                Err(
423                    "Negative duration is not supported in by_duration offset reset strategy."
424                        .to_owned(),
425                ),
426                unparsable.clone(),
427                unparsable.clone(),
428                unparsable.clone(),
429                unparsable.clone(),
430                unparsable.clone(),
431                unparsable.clone(),
432                unparsable,
433                Err("Unknown auto offset reset strategy: EARLIEST".to_owned()),
434                Err("Unknown auto offset reset strategy: unknown".to_owned()),
435            ],
436            [
437                Ok(IsolationLevel::ReadCommitted),
438                Ok(IsolationLevel::ReadUncommitted),
439                Err(()),
440                Err(()),
441            ],
442        );
443        assert2::assert!(actual == expected);
444    }
445
446    #[test]
447    fn isolation_level_wire_values_match_fetch_request_encoding() {
448        for (_name, isolation, expected) in [
449            ("read uncommitted", IsolationLevel::ReadUncommitted, 0),
450            ("read committed", IsolationLevel::ReadCommitted, 1),
451        ] {
452            assert2::assert!(isolation.wire() == expected);
453        }
454    }
455
456    #[test]
457    fn peek_version_requires_two_bytes() {
458        for (_name, bytes, want) in [
459            ("empty", &[][..], 0),
460            ("truncated", &[0x7f][..], 0),
461            ("two byte version", &[0, 3][..], 3),
462        ] {
463            assert2::assert!(peek_version(bytes) == want);
464        }
465    }
466
467    #[test]
468    fn decode_subscription_short_or_malformed_payload_uses_empty_fallback() {
469        for (_name, payload) in [
470            ("empty", &[][..]),
471            ("truncated version", &[0x7f][..]),
472            ("version only", &[0, 3][..]),
473        ] {
474            let decoded = decode_subscription(payload);
475            assert2::assert!(
476                decoded
477                    == DecodedSubscription {
478                        topics: Vec::new(),
479                        owned: Vec::new(),
480                        generation_id: -1,
481                        rack_id: None,
482                        user_data: None,
483                    }
484            );
485        }
486    }
487
488    #[test]
489    fn decode_assignment_short_or_malformed_payload_uses_empty_fallback() {
490        for (_name, payload) in [
491            ("empty", &[][..]),
492            ("truncated version", &[0x7f][..]),
493            ("version only", &[0, 3][..]),
494        ] {
495            assert2::assert!(decode_assignment(payload).is_empty());
496        }
497    }
498
499    /// Kafka's `ConsumerProtocol.serializeSubscription` sorts the topics and
500    /// the owned partitions, and carries the user data.
501    #[test]
502    fn subscription_sorts_topics_and_owned_partitions_and_keeps_user_data() {
503        let s = encode_subscription(
504            &["t2".into(), "t1".into()],
505            &[("t2".into(), 1), ("t1".into(), 3), ("t2".into(), 0)],
506            5,
507            None,
508            Some(Bytes::from_static(&[0, 0, 0, 5])),
509        )
510        .unwrap();
511        assert2::assert!(
512            decode_subscription(&s)
513                == DecodedSubscription {
514                    topics: vec!["t1".into(), "t2".into()],
515                    owned: vec![("t1".into(), 3), ("t2".into(), 0), ("t2".into(), 1)],
516                    generation_id: 5,
517                    rack_id: None,
518                    user_data: Some(Bytes::from_static(&[0, 0, 0, 5])),
519                }
520        );
521    }
522
523    #[test]
524    fn subscription_round_trip() {
525        let s = encode_subscription(&["t1".into(), "t2".into()], &[], -1, None, None).unwrap();
526        let decoded = decode_subscription(&s);
527        assert2::assert!(decoded.topics == vec!["t1", "t2"]);
528    }
529
530    #[test]
531    fn subscription_empty_round_trip() {
532        let s = encode_subscription(&[], &[], -1, None, None).unwrap();
533        let decoded = decode_subscription(&s);
534        assert2::assert!(
535            decoded
536                == DecodedSubscription {
537                    topics: Vec::new(),
538                    owned: Vec::new(),
539                    generation_id: -1,
540                    rack_id: None,
541                    user_data: None,
542                }
543        );
544    }
545
546    #[test]
547    fn subscription_v3_owned_partitions_round_trip() {
548        let owned = vec![("t".into(), 0), ("t".into(), 1), ("u".into(), 0)];
549        let s = encode_subscription(&["t".into(), "u".into()], &owned, -1, None, None).unwrap();
550        let decoded = decode_subscription(&s);
551        let mut got = decoded.owned;
552        got.sort();
553        let mut want = owned;
554        want.sort();
555        assert2::assert!(got == want);
556    }
557
558    #[test]
559    fn subscription_v3_generation_and_rack_round_trip() {
560        let s = encode_subscription(&["t".into()], &[], 42, Some("rack-a"), None).unwrap();
561        let decoded = decode_subscription(&s);
562        assert2::assert!(
563            decoded
564                == DecodedSubscription {
565                    topics: vec!["t".into()],
566                    owned: Vec::new(),
567                    generation_id: 42,
568                    rack_id: Some("rack-a".into()),
569                    user_data: None,
570                }
571        );
572    }
573
574    #[test]
575    fn subscription_decodes_v1_payload() {
576        use bytes::BufMut;
577        let mut buf = BytesMut::new();
578        buf.put_i16(1);
579        buf.put_i32(1);
580        let t = "t1";
581        buf.put_i16(i16::try_from(t.len()).unwrap());
582        buf.put_slice(t.as_bytes());
583        buf.put_i32(-1); // user_data null
584        buf.put_i32(0); // owned_partitions empty (v1)
585        let payload = buf.freeze();
586        let decoded = decode_subscription(&payload);
587        assert2::assert!(
588            decoded
589                == DecodedSubscription {
590                    topics: vec!["t1".to_owned()],
591                    owned: Vec::new(),
592                    generation_id: -1,
593                    rack_id: None,
594                    user_data: None,
595                }
596        );
597    }
598
599    #[test]
600    fn assignment_round_trip() {
601        let s = encode_assignment(&[("t".into(), 0), ("t".into(), 1), ("u".into(), 0)]).unwrap();
602        let mut decoded = decode_assignment(&s);
603        decoded.sort();
604        assert2::assert!(
605            decoded
606                == vec![
607                    ("t".to_owned(), 0),
608                    ("t".to_owned(), 1),
609                    ("u".to_owned(), 0)
610                ]
611        );
612    }
613
614    #[test]
615    fn assignment_empty_round_trip() {
616        let s = encode_assignment(&[]).unwrap();
617        let decoded = decode_assignment(&s);
618        assert2::assert!(decoded.is_empty());
619    }
620}
krabka_client_consumer/builder.rs

krabka_client_consumer/
builder.rs
