nodedb_types/
array_route.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Routing bucket and vShard of an array tile.
4//!
5//! A tile key orders tiles on a space-filling curve, so its top bits carry
6//! little entropy when tile indices are small. Routing therefore takes the
7//! top bits of a bijective mix of the key. Stored keys stay unchanged.
8//! The cluster router, the Data Plane shard filter and the snapshot paths
9//! use these functions, so a tile has one bucket and one vShard everywhere.
10//!
11//! The router uses at most [`VShardId::BITS`] bucket bits. Bucket `b` of `R`
12//! routing bits owns the vShards `[b * stride, (b + 1) * stride)`, with
13//! `stride = VShardId::COUNT >> R`. Its primary vShard is `b * stride`.
14
15use serde::{Deserialize, Serialize};
16
17use crate::id::VShardId;
18
19/// The largest routing granularity an array declares.
20pub const MAX_PREFIX_BITS: u8 = 16;
21
22/// An array routing granularity outside `1..=MAX_PREFIX_BITS`.
23#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
24#[error("array routing: prefix_bits must be 1-{}, got {}", MAX_PREFIX_BITS, .0)]
25pub struct PrefixBitsOutOfRange(pub u8);
26
27/// MurmurHash3 64-bit finalizer. A bijection on `u64`: distinct keys keep
28/// distinct values, and every input bit affects every output bit.
29pub const fn fmix64(key: u64) -> u64 {
30    let mut k = key;
31    k ^= k >> 33;
32    k = k.wrapping_mul(0xff51_afd7_ed55_8ccd);
33    k ^= k >> 33;
34    k = k.wrapping_mul(0xc4ce_b9fe_1a85_ec53);
35    k ^= k >> 33;
36    k
37}
38
39/// The routing bucket of tile key `tile_key`: the top `prefix_bits` bits of
40/// `fmix64(tile_key)`. Returns a value in `[0, 2^prefix_bits)`. With
41/// `prefix_bits == 0` every key maps to bucket 0.
42pub const fn prefix_bucket(tile_key: u64, prefix_bits: u8) -> u32 {
43    let bits = if prefix_bits > 32 { 32 } else { prefix_bits };
44    if bits == 0 {
45        return 0;
46    }
47    (fmix64(tile_key) >> (64 - bits as u32)) as u32
48}
49
50/// The bucket one shard serves. A shard read keeps only the tiles whose
51/// key falls in this bucket, so shards that share one Data Plane never
52/// return the same tile twice.
53#[derive(
54    Debug,
55    Clone,
56    Copy,
57    PartialEq,
58    Eq,
59    Hash,
60    Serialize,
61    Deserialize,
62    zerompk::ToMessagePack,
63    zerompk::FromMessagePack,
64)]
65pub struct ShardBucket {
66    /// Bucket index in `[0, 2^prefix_bits)`.
67    pub bucket: u32,
68    /// Number of mixed key bits the bucket is taken from.
69    pub prefix_bits: u8,
70}
71
72impl ShardBucket {
73    /// Whether tile key `tile_key` routes to this bucket.
74    pub const fn contains(&self, tile_key: u64) -> bool {
75        prefix_bucket(tile_key, self.prefix_bits) == self.bucket
76    }
77}
78
79/// Whether `tile_key` passes the optional shard filter. `None` keeps every
80/// tile.
81pub fn shard_bucket_admits(filter: Option<ShardBucket>, tile_key: u64) -> bool {
82    filter.is_none_or(|b| b.contains(tile_key))
83}
84
85/// The bucket bits the router uses for `prefix_bits`: at most
86/// [`VShardId::BITS`], since a finer bucket has no vShard of its own.
87pub fn routing_bits(prefix_bits: u8) -> Result<u8, PrefixBitsOutOfRange> {
88    if prefix_bits == 0 || prefix_bits > MAX_PREFIX_BITS {
89        return Err(PrefixBitsOutOfRange(prefix_bits));
90    }
91    Ok(prefix_bits.min(VShardId::BITS))
92}
93
94/// The primary vShard of `bucket` at `bits` routing bits. `bits` comes from
95/// [`routing_bits`], so it is at most [`VShardId::BITS`].
96pub const fn bucket_to_vshard(bucket: u32, bits: u8) -> u32 {
97    bucket << (VShardId::BITS - bits) as u32
98}
99
100/// The primary vShard that owns the tile with key `tile_key`.
101///
102/// `prefix_bits` is the array's routing granularity (1–16). Returns a vShard
103/// ID in `[0, VShardId::COUNT)`.
104pub fn array_vshard_for_tile(tile_key: u64, prefix_bits: u8) -> Result<u32, PrefixBitsOutOfRange> {
105    let bits = routing_bits(prefix_bits)?;
106    Ok(bucket_to_vshard(prefix_bucket(tile_key, bits), bits))
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn fmix64_is_injective_on_a_sample() {
115        let mut seen = std::collections::HashSet::new();
116        for k in 0u64..100_000 {
117            assert!(seen.insert(fmix64(k)), "collision at {k}");
118        }
119        assert_eq!(fmix64(0), 0);
120    }
121
122    #[test]
123    fn prefix_bucket_stays_in_range() {
124        for bits in 1u8..=16 {
125            for k in [0u64, 1, 2, 63, 1 << 40, u64::MAX / 3, u64::MAX] {
126                assert!(prefix_bucket(k, bits) < (1u32 << bits), "bits={bits} k={k}");
127            }
128        }
129        assert_eq!(prefix_bucket(u64::MAX, 0), 0);
130    }
131
132    #[test]
133    fn small_keys_spread_over_buckets() {
134        // Keys whose top bits are all zero land in most of the 256 buckets.
135        let buckets: std::collections::HashSet<u32> =
136            (0u64..4096).map(|k| prefix_bucket(k, 8)).collect();
137        assert!(
138            buckets.len() >= 240,
139            "only {} of 256 buckets used",
140            buckets.len()
141        );
142    }
143
144    #[test]
145    fn shard_bucket_partitions_keys() {
146        let bits = 4u8;
147        for k in 0u64..2048 {
148            let owners = (0u32..(1 << bits))
149                .filter(|&b| {
150                    ShardBucket {
151                        bucket: b,
152                        prefix_bits: bits,
153                    }
154                    .contains(k)
155                })
156                .count();
157            assert_eq!(owners, 1, "key {k} has {owners} owners");
158        }
159        assert!(shard_bucket_admits(None, 7));
160    }
161
162    #[test]
163    fn tile_routing_invalid_prefix_bits() {
164        assert_eq!(array_vshard_for_tile(0, 0), Err(PrefixBitsOutOfRange(0)));
165        assert_eq!(array_vshard_for_tile(0, 17), Err(PrefixBitsOutOfRange(17)));
166    }
167
168    #[test]
169    fn tile_routing_uses_the_mixed_bucket() {
170        // prefix_bits=8: stride=4, so the vShard is the mixed bucket times 4.
171        for key in [0u64, 1, 2, 0x0100_0000_0000_0000, u64::MAX] {
172            let vshard = array_vshard_for_tile(key, 8).unwrap();
173            assert_eq!(vshard, prefix_bucket(key, 8) * 4, "key={key:#x}");
174        }
175        // prefix_bits=10: stride=1, so the vShard is the mixed bucket.
176        let key = 0x0040_0000_0000_0000;
177        assert_eq!(
178            array_vshard_for_tile(key, 10).unwrap(),
179            prefix_bucket(key, 10)
180        );
181    }
182
183    #[test]
184    fn tile_routing_output_in_vshard_range() {
185        // Every vShard returned is below VShardId::COUNT, including
186        // prefix_bits above VShardId::BITS.
187        for bits in 1u8..=MAX_PREFIX_BITS {
188            for &key in &[0u64, 1, u64::MAX / 3, u64::MAX / 2, u64::MAX] {
189                let vshard = array_vshard_for_tile(key, bits).unwrap();
190                assert!(
191                    vshard < VShardId::COUNT,
192                    "bits={bits} key={key:#x} → vshard={vshard}"
193                );
194            }
195        }
196    }
197
198    #[test]
199    fn prefix_bits_above_vshard_bits_route_like_vshard_bits() {
200        for key in [3u64, 99, 1 << 50, u64::MAX - 7] {
201            let at_max = array_vshard_for_tile(key, VShardId::BITS).unwrap();
202            for bits in VShardId::BITS..=MAX_PREFIX_BITS {
203                assert_eq!(array_vshard_for_tile(key, bits).unwrap(), at_max);
204            }
205        }
206    }
207
208    #[test]
209    fn small_tile_keys_spread_over_many_vshards() {
210        // Keys whose top bits are all zero reach most of the 256 buckets.
211        let vshards: std::collections::BTreeSet<u32> = (0u64..4096)
212            .map(|key| array_vshard_for_tile(key, 8).unwrap())
213            .collect();
214        assert!(
215            vshards.len() >= 240,
216            "only {} of 256 buckets used",
217            vshards.len()
218        );
219    }
220
221    #[test]
222    fn bucket_to_vshard_strides() {
223        assert_eq!(VShardId::COUNT, 1 << VShardId::BITS);
224        assert_eq!(VShardId::BITS, 10);
225        assert_eq!(bucket_to_vshard(1, 1), 512);
226        assert_eq!(bucket_to_vshard(1, 8), 4);
227        assert_eq!(bucket_to_vshard(1, 10), 1);
228        assert_eq!(routing_bits(16), Ok(10));
229    }
230
231    #[test]
232    fn shard_bucket_roundtrips_through_msgpack() {
233        let b = ShardBucket {
234            bucket: 41,
235            prefix_bits: 8,
236        };
237        let bytes = zerompk::to_msgpack_vec(&b).expect("encode");
238        let back: ShardBucket = zerompk::from_msgpack(&bytes).expect("decode");
239        assert_eq!(b, back);
240    }
241}
nodedb_types/array_route.rs

nodedb_types/
array_route.rs
