shardcache_client_rs/
routing.rs

1use std::io::Write;
2use std::net::{SocketAddr, ToSocketAddrs};
3
4use crate::error::{Result, ShardCacheClientError};
5
6/// Router for shard-owned SCNP direct ports.
7///
8/// The base address is the first shard-owned port. Shard `n` is expected at
9/// `base_port + n`; when the server also exposes a fanout port, this base is
10/// usually `SHARDCACHE_DIRECT_SHARD_BASE_PORT` or the fanout port + 1.
11#[derive(Debug, Clone, Copy)]
12pub struct ShardCacheDirectRouter {
13    base_addr: SocketAddr,
14    shard_count: usize,
15    shift: u32,
16    route_mode: ShardCacheRouteMode,
17}
18
19impl ShardCacheDirectRouter {
20    /// Creates a direct router for `shard_count` server shards.
21    pub fn new(addr: impl ToSocketAddrs, shard_count: usize) -> Result<Self> {
22        if shard_count == 0 || !shard_count.is_power_of_two() {
23            return Err(ShardCacheClientError::Config(format!(
24                "SCNP direct shard count must be a non-zero power of two: {shard_count}"
25            )));
26        }
27        let base_addr = resolve_one(addr)?;
28        Ok(Self {
29            base_addr,
30            shard_count,
31            shift: shift_for(shard_count),
32            route_mode: ShardCacheRouteMode::FullKey,
33        })
34    }
35
36    /// Sets how direct shard routing chooses the owning shard.
37    ///
38    /// `FullKey` is the normal point-key mode. `SessionPrefix` routes keys of
39    /// the form `s:<session>:c:<chunk>` by `s:<session>` while preserving the
40    /// full-key hash used for lookup within that shard.
41    pub fn with_route_mode(mut self, route_mode: ShardCacheRouteMode) -> Self {
42        self.route_mode = route_mode;
43        self
44    }
45
46    /// Returns the number of direct shard ports.
47    pub fn shard_count(&self) -> usize {
48        self.shard_count
49    }
50
51    /// Computes the routed SCNP metadata for `key`.
52    pub fn route_key(&self, key: &[u8]) -> ShardCacheRoute {
53        let key_hash = hash_key(key);
54        let shard_id = match self.route_mode {
55            ShardCacheRouteMode::OverflowSlot => overflow_slot_shard(key, self.shard_count)
56                .unwrap_or_else(|| stripe_index(key_hash, self.shift)),
57            mode => stripe_index(mode.route_hash(key, key_hash), self.shift),
58        };
59        ShardCacheRoute {
60            key_hash,
61            key_tag: hash_key_tag_from_hash(key_hash),
62            shard_id,
63        }
64    }
65
66    /// Returns the socket address for `shard_id`.
67    pub fn shard_addr(&self, shard_id: usize) -> Result<SocketAddr> {
68        if shard_id >= self.shard_count {
69            return Err(ShardCacheClientError::Config(format!(
70                "SCNP direct shard {shard_id} out of range for {} shards",
71                self.shard_count
72            )));
73        }
74        let mut addr = self.base_addr;
75        let offset = u16::try_from(shard_id).map_err(|_| {
76            ShardCacheClientError::Config(format!("SCNP direct shard id exceeds u16: {shard_id}"))
77        })?;
78        let port = self.base_addr.port().checked_add(offset).ok_or_else(|| {
79            ShardCacheClientError::Config(format!(
80                "SCNP direct shard port overflows for shard {shard_id}"
81            ))
82        })?;
83        addr.set_port(port);
84        Ok(addr)
85    }
86}
87
88/// Shard-routing mode for direct SCNP clients.
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
90pub enum ShardCacheRouteMode {
91    /// Route every key by its full key hash.
92    FullKey,
93    /// Route `s:<session>:c:<chunk>` keys by the session prefix.
94    SessionPrefix,
95    /// Route a versioned internal overflow key by its encoded shard field.
96    OverflowSlot,
97}
98
99impl ShardCacheRouteMode {
100    /// Parses the route mode used by benchmark and deployment knobs.
101    pub fn parse(value: &str) -> Result<Self> {
102        match value {
103            "full_key" | "full-key" | "point" => Ok(Self::FullKey),
104            "session_prefix" | "session-prefix" | "session" => Ok(Self::SessionPrefix),
105            "overflow_slot" | "overflow-slot" | "overflow" => Ok(Self::OverflowSlot),
106            other => Err(ShardCacheClientError::Config(format!(
107                "unknown SCNP direct route mode `{other}`; use full_key, session_prefix, or overflow_slot"
108            ))),
109        }
110    }
111
112    fn route_hash(self, key: &[u8], key_hash: u64) -> u64 {
113        match self {
114            Self::FullKey => key_hash,
115            Self::SessionPrefix => hash_key(session_route_prefix(key)),
116            Self::OverflowSlot => key_hash,
117        }
118    }
119}
120
121const OVERFLOW_SLOT_KEY_MAGIC: &[u8; 8] = b"SCKVKEY1";
122
123fn overflow_slot_shard(key: &[u8], shard_count: usize) -> Option<usize> {
124    if key.len() < 18 || &key[..8] != OVERFLOW_SLOT_KEY_MAGIC {
125        return None;
126    }
127    let shard_id = u32::from_le_bytes(key[8..12].try_into().ok()?) as usize;
128    (shard_id < shard_count).then_some(shard_id)
129}
130
131/// Precomputed routed SCNP metadata for a key.
132#[derive(Debug, Clone, Copy, PartialEq, Eq)]
133pub struct ShardCacheRoute {
134    /// Primary key hash.
135    pub key_hash: u64,
136    /// Secondary key tag used by direct full-key lookups.
137    pub key_tag: u64,
138    /// Owning shard id.
139    pub shard_id: usize,
140}
141
142impl ShardCacheRoute {
143    pub(crate) fn write_to<W: Write>(&self, w: &mut W) -> Result<()> {
144        w.write_all(&self.key_hash.to_le_bytes())?;
145        w.write_all(&(self.shard_id as u32).to_le_bytes())?;
146        w.write_all(&self.key_tag.to_le_bytes())?;
147        Ok(())
148    }
149}
150
151/// Computes shardcache's primary XXH3 key hash.
152pub fn hash_key(key: &[u8]) -> u64 {
153    xxhash_rust::xxh3::xxh3_64(key)
154}
155
156/// Computes shardcache's secondary key fingerprint.
157pub fn hash_key_tag(key: &[u8]) -> u64 {
158    hash_key_tag_from_hash(hash_key(key))
159}
160
161/// Computes the secondary key fingerprint from an already-computed primary hash.
162pub fn hash_key_tag_from_hash(hash: u64) -> u64 {
163    hash >> 56
164}
165
166fn session_route_prefix(key: &[u8]) -> &[u8] {
167    if !key.starts_with(b"s:") {
168        return key;
169    }
170
171    if let Some(index) = session_chunk_separator(key) {
172        return &key[..index];
173    }
174
175    key
176}
177
178#[inline(always)]
179fn session_chunk_separator(key: &[u8]) -> Option<usize> {
180    if key.len() < 3 {
181        return None;
182    }
183
184    let mut index = key.len() - 3;
185    loop {
186        if key[index] == b':' && key[index + 1] == b'c' && key[index + 2] == b':' {
187            return Some(index);
188        }
189        if index == 0 {
190            return None;
191        }
192        index -= 1;
193    }
194}
195
196/// Computes the shard index for `hash` and `shard_count`.
197pub fn shard_index(hash: u64, shard_count: usize) -> Result<usize> {
198    if shard_count == 0 || !shard_count.is_power_of_two() {
199        return Err(ShardCacheClientError::Config(format!(
200            "shard count must be a non-zero power of two: {shard_count}"
201        )));
202    }
203    Ok(stripe_index(hash, shift_for(shard_count)))
204}
205
206fn stripe_index(hash: u64, shift: u32) -> usize {
207    if shift == usize::BITS {
208        0
209    } else {
210        ((hash as usize) << 7) >> shift
211    }
212}
213
214fn shift_for(shard_count: usize) -> u32 {
215    debug_assert!(shard_count > 0 && shard_count.is_power_of_two());
216    usize::BITS - shard_count.trailing_zeros()
217}
218
219fn resolve_one(addr: impl ToSocketAddrs) -> Result<SocketAddr> {
220    addr.to_socket_addrs()?.next().ok_or_else(|| {
221        ShardCacheClientError::Config("SCNP address resolved to no socket addresses".into())
222    })
223}
shardcache_client_rs/routing.rs

shardcache_client_rs/
routing.rs
