laminar_core/state/vnode/
mod.rs

1//! [`VnodeRegistry`] — runtime-configurable virtual node topology.
2//!
3//! A registry owns:
4//!
5//! - the runtime-selected vnode count,
6//! - the node-per-vnode assignment (for distributed modes),
7//! - a monotonically increasing `assignment_version` used to fence stale
8//!   assignment publications.
9//!
10//! Vnode assignment is derived from the row's primary key via
11//! [`key_hash`] (xxh3) and modulo `vnode_count`. Connectors that
12//! need a vnode ID for an event call [`VnodeRegistry::vnode_for_key`].
13
14use std::collections::BTreeMap;
15use std::fmt;
16use std::num::NonZeroU16;
17use std::sync::atomic::{AtomicU64, Ordering};
18use std::sync::Arc;
19
20use parking_lot::{RwLock, RwLockReadGuard};
21use serde::{Deserialize, Serialize};
22
23/// Durable encoding and hashing contract used to map keys to key groups.
24///
25/// Any change to Arrow row encoding, key hashing, or modulo mapping requires a
26/// version bump; recovery rejects checkpoints with a different identity. Placement
27/// is not part of this ABI: durable assignment publications are the ownership authority.
28pub const PARTITIONING_ABI_VERSION: u16 = 2;
29
30/// Validated number of stable key groups in a pipeline.
31///
32/// Zero and values above [`u16::MAX`] are rejected. The compact upper bound keeps
33/// ownership metadata bounded while providing substantially more rescale slots
34/// than a production pipeline should need.
35#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
36#[serde(transparent)]
37pub struct KeyGroupCount(NonZeroU16);
38
39/// Default stable key-group count for every deployment tier.
40pub const DEFAULT_KEY_GROUP_COUNT: KeyGroupCount = match NonZeroU16::new(256) {
41    Some(value) => KeyGroupCount(value),
42    None => unreachable!(),
43};
44
45/// Largest key-group count representable by the persisted checkpoint ABI.
46pub const MAX_KEY_GROUP_COUNT: u32 = u16::MAX as u32;
47
48impl KeyGroupCount {
49    /// Build from a nonzero value.
50    #[must_use]
51    pub const fn new(value: NonZeroU16) -> Self {
52        Self(value)
53    }
54
55    /// Exact underlying count.
56    #[must_use]
57    pub const fn get(self) -> u16 {
58        self.0.get()
59    }
60
61    /// Exact nonzero representation.
62    #[must_use]
63    pub const fn into_non_zero(self) -> NonZeroU16 {
64        self.0
65    }
66}
67
68impl fmt::Display for KeyGroupCount {
69    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
70        self.get().fmt(f)
71    }
72}
73
74/// A key-group count outside the supported `1..=65535` range.
75#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
76#[error("key-group count {value} is outside 1..={}", u16::MAX)]
77pub struct InvalidKeyGroupCount {
78    value: u32,
79}
80
81impl InvalidKeyGroupCount {
82    /// Rejected input value.
83    #[must_use]
84    pub const fn value(self) -> u32 {
85        self.value
86    }
87}
88
89impl TryFrom<u16> for KeyGroupCount {
90    type Error = InvalidKeyGroupCount;
91
92    fn try_from(value: u16) -> Result<Self, Self::Error> {
93        NonZeroU16::new(value)
94            .map(Self)
95            .ok_or(InvalidKeyGroupCount {
96                value: u32::from(value),
97            })
98    }
99}
100
101impl TryFrom<u32> for KeyGroupCount {
102    type Error = InvalidKeyGroupCount;
103
104    fn try_from(value: u32) -> Result<Self, Self::Error> {
105        let narrowed = u16::try_from(value).map_err(|_| InvalidKeyGroupCount { value })?;
106        Self::try_from(narrowed).map_err(|_| InvalidKeyGroupCount { value })
107    }
108}
109
110impl From<NonZeroU16> for KeyGroupCount {
111    fn from(value: NonZeroU16) -> Self {
112        Self::new(value)
113    }
114}
115
116impl From<KeyGroupCount> for NonZeroU16 {
117    fn from(value: KeyGroupCount) -> Self {
118        value.into_non_zero()
119    }
120}
121
122impl From<KeyGroupCount> for u16 {
123    fn from(value: KeyGroupCount) -> Self {
124        value.get()
125    }
126}
127
128impl From<KeyGroupCount> for u32 {
129    fn from(value: KeyGroupCount) -> Self {
130        u32::from(value.get())
131    }
132}
133
134impl From<KeyGroupCount> for usize {
135    fn from(value: KeyGroupCount) -> Self {
136        usize::from(value.get())
137    }
138}
139
140/// Unique identifier for a node. Also the owner id for vnodes; cluster
141/// membership and vnode ownership identify the same thing.
142#[derive(
143    Debug,
144    Clone,
145    Copy,
146    PartialEq,
147    Eq,
148    PartialOrd,
149    Ord,
150    Hash,
151    Serialize,
152    Deserialize,
153    rkyv::Archive,
154    rkyv::Serialize,
155    rkyv::Deserialize,
156)]
157pub struct NodeId(pub u64);
158
159/// Stable owner identity for embedded and standalone runtimes.
160pub const LOCAL_NODE_ID: NodeId = NodeId(1);
161
162impl NodeId {
163    /// Sentinel meaning "unassigned".
164    pub const UNASSIGNED: Self = Self(0);
165
166    /// True if this is the unassigned sentinel.
167    #[must_use]
168    pub const fn is_unassigned(&self) -> bool {
169        self.0 == 0
170    }
171}
172
173impl fmt::Display for NodeId {
174    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
175        write!(f, "node-{}", self.0)
176    }
177}
178
179/// One immutable publication of vnode ownership and its assignment version.
180#[derive(Clone)]
181pub struct VnodeAssignmentSnapshot {
182    version: u64,
183    owners: Arc<[NodeId]>,
184    owner_changed_versions: Arc<[u64]>,
185}
186
187impl VnodeAssignmentSnapshot {
188    /// Monotonic assignment version for this publication.
189    #[must_use]
190    pub const fn version(&self) -> u64 {
191        self.version
192    }
193
194    /// Immutable vnode-owner vector for this publication.
195    #[must_use]
196    pub fn owners(&self) -> &[NodeId] {
197        &self.owners
198    }
199
200    /// Version of the last ownership change for `vnode`.
201    #[must_use]
202    pub fn owner_changed_version(&self, vnode: u32) -> Option<u64> {
203        self.owner_changed_versions.get(vnode as usize).copied()
204    }
205}
206
207/// Read guard that pins one assignment publication while a source captures a
208/// data batch or checkpoint cursor.
209pub struct VnodeAssignmentReadGuard<'a>(RwLockReadGuard<'a, VnodeAssignmentSnapshot>);
210
211impl std::ops::Deref for VnodeAssignmentReadGuard<'_> {
212    type Target = VnodeAssignmentSnapshot;
213
214    fn deref(&self) -> &Self::Target {
215        &self.0
216    }
217}
218
219fn changed_owner_versions(
220    current: &VnodeAssignmentSnapshot,
221    next_owners: &[NodeId],
222    next_version: u64,
223) -> Arc<[u64]> {
224    let skipped_generation = next_version > current.version.saturating_add(1);
225    current
226        .owners
227        .iter()
228        .zip(next_owners)
229        .enumerate()
230        .map(|(vnode, (current_owner, next_owner))| {
231            if current_owner == next_owner && !skipped_generation {
232                current.owner_changed_versions[vnode]
233            } else {
234                next_version
235            }
236        })
237        .collect::<Vec<_>>()
238        .into()
239}
240
241/// Runtime registry of vnode topology and assignment.
242pub struct VnodeRegistry {
243    vnode_count: u32,
244    assignment: RwLock<VnodeAssignmentSnapshot>,
245    /// Lock-free observation fence for hot runtime readiness checks. The
246    /// authoritative version also lives inside `assignment`, where it is bound
247    /// to owners.
248    assignment_version: AtomicU64,
249}
250
251impl std::fmt::Debug for VnodeRegistry {
252    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
253        f.debug_struct("VnodeRegistry")
254            .field("vnode_count", &self.vnode_count)
255            .field(
256                "assignment_version",
257                &self.assignment_version.load(Ordering::Relaxed),
258            )
259            .finish_non_exhaustive()
260    }
261}
262
263impl VnodeRegistry {
264    /// Create a registry sized for `vnode_count` vnodes, all marked
265    /// as [`NodeId::UNASSIGNED`]. The assignment version starts at 1.
266    ///
267    /// # Panics
268    /// Panics if `vnode_count == 0`.
269    #[must_use]
270    pub fn new(vnode_count: u32) -> Self {
271        assert!(vnode_count > 0, "vnode_count must be > 0");
272        let assignment: Arc<[NodeId]> =
273            std::iter::repeat_n(NodeId::UNASSIGNED, vnode_count as usize)
274                .collect::<Vec<_>>()
275                .into();
276        Self {
277            vnode_count,
278            assignment: RwLock::new(VnodeAssignmentSnapshot {
279                version: 1,
280                owners: assignment,
281                owner_changed_versions: std::iter::repeat_n(1, vnode_count as usize)
282                    .collect::<Vec<_>>()
283                    .into(),
284            }),
285            assignment_version: AtomicU64::new(1),
286        }
287    }
288
289    /// Create a registry with every vnode unassigned at version 0, so the first stored
290    /// assignment snapshot (version >= 1) adopts through the standard rotation path.
291    ///
292    /// # Panics
293    /// Panics if `vnode_count == 0`.
294    #[must_use]
295    pub fn new_unassigned(vnode_count: u32) -> Self {
296        let registry = Self::new(vnode_count);
297        {
298            let mut initial = registry.assignment.write();
299            initial.version = 0;
300            initial.owner_changed_versions = std::iter::repeat_n(0, vnode_count as usize)
301                .collect::<Vec<_>>()
302                .into();
303        }
304        registry.assignment_version.store(0, Ordering::Release);
305        registry
306    }
307
308    /// Create a registry where every vnode is owned by the same node.
309    ///
310    /// Used by single-instance / embedded deployments.
311    ///
312    /// # Panics
313    /// Panics if `vnode_count == 0` or `owner` is [`NodeId::UNASSIGNED`].
314    #[must_use]
315    pub fn single_owner(vnode_count: u32, owner: NodeId) -> Self {
316        assert!(vnode_count > 0, "vnode_count must be > 0");
317        assert_ne!(owner, NodeId::UNASSIGNED, "owner must be assigned");
318        let assignment: Arc<[NodeId]> = std::iter::repeat_n(owner, vnode_count as usize)
319            .collect::<Vec<_>>()
320            .into();
321        Self {
322            vnode_count,
323            assignment: RwLock::new(VnodeAssignmentSnapshot {
324                version: 1,
325                owners: assignment,
326                owner_changed_versions: std::iter::repeat_n(1, vnode_count as usize)
327                    .collect::<Vec<_>>()
328                    .into(),
329            }),
330            assignment_version: AtomicU64::new(1),
331        }
332    }
333
334    /// Number of vnodes.
335    #[must_use]
336    pub fn vnode_count(&self) -> u32 {
337        self.vnode_count
338    }
339
340    /// Current monotonic assignment version.
341    #[must_use]
342    pub fn assignment_version(&self) -> u64 {
343        self.assignment_version.load(Ordering::Acquire)
344    }
345
346    /// Owner of a given vnode. Returns [`NodeId::UNASSIGNED`] if the
347    /// vnode is out of range or unassigned.
348    #[must_use]
349    pub fn owner(&self, vnode: u32) -> NodeId {
350        if vnode >= self.vnode_count {
351            return NodeId::UNASSIGNED;
352        }
353        self.assignment.read().owners[vnode as usize]
354    }
355
356    /// Snapshot the current assignment vector. Cheap — internally an
357    /// `Arc::clone`.
358    #[must_use]
359    pub fn snapshot(&self) -> Arc<[NodeId]> {
360        Arc::clone(&self.assignment.read().owners)
361    }
362
363    /// Consistent ownership, version, and source-handoff publication.
364    #[must_use]
365    pub fn versioned_snapshot(&self) -> VnodeAssignmentSnapshot {
366        self.assignment.read().clone()
367    }
368
369    /// Pin the current publication for a short, non-blocking source capture.
370    #[must_use]
371    pub fn read_assignment(&self) -> VnodeAssignmentReadGuard<'_> {
372        VnodeAssignmentReadGuard(self.assignment.read())
373    }
374
375    /// Replace the full assignment and bump the version.
376    ///
377    /// # Panics
378    /// Panics if `new_assignment.len() != self.vnode_count`.
379    pub fn set_assignment(&self, new_assignment: Arc<[NodeId]>) {
380        assert_eq!(
381            new_assignment.len(),
382            self.vnode_count as usize,
383            "assignment length mismatch: got {}, expected {}",
384            new_assignment.len(),
385            self.vnode_count,
386        );
387        let mut current = self.assignment.write();
388        let version = current
389            .version
390            .checked_add(1)
391            .expect("assignment version overflow");
392        let owner_changed_versions = changed_owner_versions(&current, &new_assignment, version);
393        *current = VnodeAssignmentSnapshot {
394            version,
395            owners: new_assignment,
396            owner_changed_versions,
397        };
398        self.assignment_version
399            .store(current.version, Ordering::Release);
400    }
401
402    /// Replace the full assignment and set the version to `version`
403    /// atomically. For recovery paths that must restore the registry to
404    /// a persisted fence generation, not a fresh bump.
405    ///
406    /// # Panics
407    /// Panics on length mismatch, or if `version` is less than the
408    /// current one (assignment versions are monotonic).
409    pub fn set_assignment_and_version(&self, new_assignment: Arc<[NodeId]>, version: u64) {
410        self.publish_assignment(new_assignment, version);
411    }
412
413    fn publish_assignment(&self, new_assignment: Arc<[NodeId]>, version: u64) {
414        assert_eq!(
415            new_assignment.len(),
416            self.vnode_count as usize,
417            "assignment length mismatch: got {}, expected {}",
418            new_assignment.len(),
419            self.vnode_count,
420        );
421        let mut guard = self.assignment.write();
422        let current = guard.version;
423        assert!(
424            version > current,
425            "assignment version must advance: got {version}, current {current}",
426        );
427        let owner_changed_versions = changed_owner_versions(&guard, &new_assignment, version);
428        *guard = VnodeAssignmentSnapshot {
429            version,
430            owners: new_assignment,
431            owner_changed_versions,
432        };
433        self.assignment_version.store(version, Ordering::Release);
434    }
435
436    /// Map a primary key to a vnode.
437    #[must_use]
438    pub fn vnode_for_key(&self, key: &[u8]) -> u32 {
439        #[allow(clippy::cast_possible_truncation)]
440        let h = (key_hash(key) % u64::from(self.vnode_count)) as u32;
441        h
442    }
443}
444
445/// Hash a key to a 64-bit value. Used to derive vnode IDs and for any
446/// other keyed-partitioning decisions.
447///
448/// Fixed to xxh3 so all pipeline stages produce the same vnode for the
449/// same key without needing to share a hasher instance.
450#[must_use]
451pub fn key_hash(key: &[u8]) -> u64 {
452    xxhash_rust::xxh3::xxh3_64(key)
453}
454
455/// Build a vnode-to-owner assignment using Rendezvous Hashing (Highest Random Weight).
456///
457/// Deterministic for a given `(vnode_count, peers)` input. Minimizes partition
458/// reshuffling on membership changes (node joins/leaves).
459///
460/// # Panics
461/// Panics if `peers` is empty.
462#[must_use]
463pub fn rendezvous_assignment(vnode_count: u32, peers: &[NodeId]) -> Arc<[NodeId]> {
464    assert!(
465        !peers.is_empty(),
466        "rendezvous_assignment needs at least one peer"
467    );
468    let mut sorted_peers = peers.to_vec();
469    sorted_peers.sort_by_key(|n| n.0);
470
471    let mut assignment = Vec::with_capacity(vnode_count as usize);
472    for v in 0..vnode_count {
473        let mut max_weight = 0;
474        let mut selected_node = sorted_peers[0];
475
476        for &node in &sorted_peers {
477            // Hash the combination of vnode ID and node ID
478            let mut buf = [0u8; 16];
479            buf[0..8].copy_from_slice(&u64::from(v).to_le_bytes());
480            buf[8..16].copy_from_slice(&node.0.to_le_bytes());
481            let weight = xxhash_rust::xxh3::xxh3_64(&buf);
482
483            // Highest weight wins, tie-break by NodeId
484            if weight > max_weight || (weight == max_weight && node.0 > selected_node.0) {
485                max_weight = weight;
486                selected_node = node;
487            }
488        }
489        assignment.push(selected_node);
490    }
491    assignment.into()
492}
493
494/// A node's failure-domain locality: ordered tier values, coarsest first
495/// (e.g. `["us-east-1", "us-east-1a", "r17"]`). Parsed from the node's
496/// `failure_domain` gossip string.
497#[derive(Debug, Clone, Default, PartialEq, Eq)]
498pub struct Locality {
499    tiers: Vec<String>,
500}
501
502impl Locality {
503    /// Build from ordered tier values, coarsest first.
504    #[must_use]
505    pub fn new(tiers: Vec<String>) -> Self {
506        Self { tiers }
507    }
508
509    /// Parse `"region=us-east-1;zone=us-east-1a;rack=r17"` (or a bare label
510    /// `"rack17"`) into its tier values. Tier names are ignored.
511    #[must_use]
512    pub fn parse(s: &str) -> Self {
513        let tiers = s
514            .split(';')
515            .map(str::trim)
516            .filter(|seg| !seg.is_empty())
517            .map(|seg| {
518                seg.split_once('=')
519                    .map_or(seg, |(_, v)| v.trim())
520                    .to_string()
521            })
522            .collect();
523        Self { tiers }
524    }
525
526    /// Failure-domain key at `tier`: the `;`-joined value prefix `0..=tier`.
527    /// Two nodes share a domain iff equal; an unlabeled node yields the empty key.
528    #[must_use]
529    pub fn domain_at(&self, tier: usize) -> String {
530        if self.tiers.is_empty() {
531            return String::new();
532        }
533        let end = tier.min(self.tiers.len() - 1);
534        self.tiers[..=end].join(";")
535    }
536}
537
538/// Resolve each node's failure-domain key at `isolation_tier`.
539fn resolve_domains(nodes: &[(NodeId, Locality)], isolation_tier: usize) -> Vec<(NodeId, String)> {
540    nodes
541        .iter()
542        .map(|(id, loc)| (*id, loc.domain_at(isolation_tier)))
543        .collect()
544}
545
546/// Owner counts per failure domain at `isolation_tier`. The largest value over
547/// `vnode_count` is the blast radius — the share of state affected if that one
548/// domain fails at once.
549#[must_use]
550pub fn owners_per_domain(
551    owners: &[NodeId],
552    nodes: &[(NodeId, Locality)],
553    isolation_tier: usize,
554) -> BTreeMap<String, u32> {
555    let dom: BTreeMap<NodeId, String> =
556        resolve_domains(nodes, isolation_tier).into_iter().collect();
557    let mut counts = BTreeMap::new();
558    for &o in owners {
559        *counts
560            .entry(dom.get(&o).cloned().unwrap_or_default())
561            .or_default() += 1;
562    }
563    counts
564}
565
566/// Vnodes currently assigned to `owner`.
567///
568/// Used by the checkpoint coordinator to decide which vnodes' durability
569/// markers it is responsible for writing each epoch, and by the leader's
570/// `seal_checkpoint` gate to know the full set to check.
571#[must_use]
572pub fn owned_vnodes(registry: &VnodeRegistry, owner: NodeId) -> Vec<u32> {
573    let assignment = registry.snapshot();
574    assignment
575        .iter()
576        .enumerate()
577        .filter(|(_, assigned)| **assigned == owner)
578        .filter_map(|(vnode, _)| u32::try_from(vnode).ok())
579        .collect()
580}
581
582/// Distinct assigned nodes other than `self_id`, sorted by id — the peer set a
583/// node fans checkpoint barriers and shuffle data out to.
584#[must_use]
585pub fn peer_owners(registry: &VnodeRegistry, self_id: NodeId) -> Vec<NodeId> {
586    let assignment = registry.snapshot();
587    let mut peers: Vec<NodeId> = assignment
588        .iter()
589        .copied()
590        .filter(|o| !o.is_unassigned() && *o != self_id)
591        .collect();
592    peers.sort_unstable();
593    peers.dedup();
594    peers
595}
596
597#[cfg(test)]
598mod tests;
laminar_core/state/vnode/mod.rs

laminar_core/state/vnode/
mod.rs
