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

krabka_client_consumer/
builder.rs
