Skip to main content

krabka_remote_storage_topic/
partitioning.rs

1//! Deterministic `TopicIdPartition` → metadata-topic-partition mapping.
2//!
3//! Every metadata event for one user partition must land in the same
4//! `__remote_log_metadata` partition, so its lifecycle ordering survives the
5//! round trip through Kafka. The hash inputs are `topic_id` and `partition`,
6//! which match the identity of `TopicIdPartition`. A topic rename therefore
7//! does not move the bucket.
8
9use {
10    crate::error::MetadataLogError, krabka_remote_storage::TopicIdPartition,
11    krabka_verified::remote_metadata_partition,
12};
13
14/// Pick the metadata-topic partition for `tp`, given the
15/// `partition_count` of `__remote_log_metadata`.
16///
17/// # Errors
18///
19/// Returns [`MetadataLogError::NonPositivePartitionCount`] when
20/// `partition_count <= 0`.
21pub fn metadata_partition_for(
22    tp: &TopicIdPartition,
23    partition_count: i32,
24) -> Result<i32, MetadataLogError> {
25    remote_metadata_partition(tp.topic_id.as_u128(), tp.partition, partition_count).ok_or(
26        MetadataLogError::NonPositivePartitionCount {
27            count: partition_count,
28        },
29    )
30}
31
32/// Deduped, sorted set of `__remote_log_metadata` partitions that carry
33/// metadata for the given user-topic-partitions, for the metadata topic's
34/// `partition_count`.
35///
36/// This is the set a broker must consume to serve remote
37/// reads for the partitions it leads or follows.
38///
39/// # Errors
40///
41/// Returns [`MetadataLogError::NonPositivePartitionCount`] when
42/// `partition_count <= 0`, even for an empty `tps`.
43pub fn metadata_partitions_for<'a, I>(
44    tps: I,
45    partition_count: i32,
46) -> Result<Vec<i32>, MetadataLogError>
47where
48    I: IntoIterator<Item = &'a TopicIdPartition>,
49{
50    if partition_count <= 0 {
51        return Err(MetadataLogError::NonPositivePartitionCount {
52            count: partition_count,
53        });
54    }
55    let mut set = tps
56        .into_iter()
57        .map(|tp| metadata_partition_for(tp, partition_count))
58        .collect::<Result<Vec<i32>, _>>()?;
59    set.sort_unstable();
60    set.dedup();
61    Ok(set)
62}
63
64#[cfg(test)]
65mod tests {
66    use {super::*, assert2::assert, std::collections::HashSet, uuid::Uuid};
67
68    fn tp(name: &str, partition: i32) -> TopicIdPartition {
69        TopicIdPartition::new(
70            Uuid::from_u128(0x1234_5678_9ABC_DEF0_0011_2233_4455_6677),
71            name,
72            partition,
73        )
74    }
75
76    #[test]
77    fn is_in_range() {
78        for p in 0..10 {
79            let bucket = metadata_partition_for(&tp("orders", p), 50).unwrap();
80            assert!((0..50).contains(&bucket), "bucket {bucket} out of [0,50)");
81        }
82    }
83
84    #[yare::parameterized(
85        first = { 0, 6 },
86        second = { 1, 36 },
87        fourth = { 3, 2 },
88        eighth = { 7, 26 },
89    )]
90    fn matches_jvm_remote_metadata_partitioner_goldens(partition: i32, expected: i32) {
91        assert2::assert!(metadata_partition_for(&tp("orders", partition), 50).unwrap() == expected);
92    }
93
94    #[test]
95    fn is_deterministic_across_calls() {
96        let a = metadata_partition_for(&tp("orders", 7), 50).unwrap();
97        let b = metadata_partition_for(&tp("orders", 7), 50).unwrap();
98        assert!(a == b);
99    }
100
101    #[test]
102    fn ignores_topic_name() {
103        // Identity is (topic_id, partition); name is informational.
104        let a = metadata_partition_for(&tp("orders", 3), 50).unwrap();
105        let b = metadata_partition_for(&tp("renamed", 3), 50).unwrap();
106        assert!(a == b, "renaming a topic must not re-bucket its metadata");
107    }
108
109    #[test]
110    fn different_partitions_distribute() {
111        let mut seen = HashSet::new();
112        for p in 0..200 {
113            seen.insert(metadata_partition_for(&tp("orders", p), 50).unwrap());
114        }
115        // Sanity: a hash that always returned 0 would land only one bucket.
116        assert!(
117            seen.len() > 5,
118            "hash should spread across buckets, got {}",
119            seen.len()
120        );
121    }
122
123    #[test]
124    fn single_partition_count_collapses_to_zero() {
125        for p in 0..20 {
126            assert!(metadata_partition_for(&tp("t", p), 1).unwrap() == 0);
127        }
128    }
129
130    #[test]
131    fn rejects_zero_partition_count() {
132        assert!(matches!(
133            metadata_partition_for(&tp("t", 0), 0),
134            Err(MetadataLogError::NonPositivePartitionCount { count: 0 })
135        ));
136    }
137
138    #[test]
139    fn metadata_partitions_for_dedupes_and_sorts() {
140        // Two user-partitions that hash to the same metadata partition must
141        // collapse to one entry; the result is sorted ascending.
142        let a = tp("orders", 0);
143        let b = tp("orders", 1);
144        let pa = metadata_partition_for(&a, 50).unwrap();
145        let pb = metadata_partition_for(&b, 50).unwrap();
146        let got = metadata_partitions_for([a.clone(), b, a].iter(), 50).unwrap();
147        let mut expected: Vec<i32> = vec![pa, pb];
148        expected.sort_unstable();
149        expected.dedup();
150        assert!(got == expected);
151        assert!(got.windows(2).all(|w| w[0] < w[1]), "sorted, deduped");
152    }
153
154    #[test]
155    fn metadata_partitions_for_empty_is_empty() {
156        let none: [TopicIdPartition; 0] = [];
157        assert!(metadata_partitions_for(none.iter(), 50).unwrap().is_empty());
158    }
159
160    #[yare::parameterized(
161        zero = { 0 },
162        negative = { -1 },
163    )]
164    fn metadata_partitions_for_empty_rejects_nonpositive_count(count: i32) {
165        let none: [TopicIdPartition; 0] = [];
166        assert2::assert!(
167            matches!(
168                metadata_partitions_for(none.iter(), count),
169                Err(MetadataLogError::NonPositivePartitionCount { count: c }) if c == count
170            ),
171            "count {count} must be rejected before iteration"
172        );
173    }
174}