kvbm_logical/blocks/
mod.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! RAII guard types that enforce the block lifecycle state machine.
5//!
6//! Every logical block in KVBM progresses through a fixed sequence of states.
7//! The guard types in this module make those transitions explicit in the type
8//! system: each state is represented by a distinct Rust type, and moving to the
9//! next state consumes the current guard and produces the next one. Dropping a
10//! guard at any point automatically returns the underlying block to the
11//! appropriate pool, so blocks are never leaked. This ensures proper behavior
12//! on early exits or exception handling.
13//!
14//! # State machine
15//!
16//! ```text
17//!                  stage / complete          register_block
18//!   Reset ──────────────────────► Staged ─────────────────► Registered
19//!     ▲                              │                          │
20//!     │          reset               │                          │
21//!     ├──────────────────────────────┘                          │
22//!     │                          drop (reset + return)          │
23//!     └─────────────────────────────────────────────────────────┘
24//! ```
25//!
26//! # Public guard types
27//!
28//! - [`MutableBlock`] -- guards a block in the **Reset** state.
29//!   - Created by [`crate::BlockManager::allocate_blocks`] or [`CompleteBlock::reset`].
30//!   - Transitions to [`CompleteBlock`] via [`stage`](MutableBlock::stage)
31//!     or [`complete`](MutableBlock::complete).
32//!
33//! - [`CompleteBlock`] -- guards a block in the **Staged** state.
34//!   - Created by [`MutableBlock::stage`] or [`MutableBlock::complete`].
35//!   - Transitions to [`MutableBlock`] via [`reset`](CompleteBlock::reset),
36//!     or to [`ImmutableBlock`] via [`crate::BlockManager::register_block`].
37//!
38//! - [`ImmutableBlock`] -- guards a block in the **Registered** state.
39//!   - Created by [`crate::BlockManager::register_block`],
40//!     [`crate::BlockManager::match_blocks`], [`crate::BlockManager::scan_matches`],
41//!     or [`WeakBlock::upgrade`].
42//!   - Transitions to [`WeakBlock`] via [`downgrade`](ImmutableBlock::downgrade).
43//!
44//! - [`WeakBlock`] -- non-owning reference to a registered block.
45//!   - Created by [`ImmutableBlock::downgrade`].
46//!   - Transitions to [`ImmutableBlock`] via [`upgrade`](WeakBlock::upgrade).
47//!
48//! # Supporting types
49//!
50//! - [`BlockMetadata`] -- trait bound satisfied by any `Clone + Send + Sync + 'static` type.
51//! - [`BlockError`] -- error type that returns the originating block to prevent leaks.
52//! - [`BlockDuplicationPolicy`] -- controls whether duplicate sequence hashes are allowed
53//!   at registration time.
54//!
55//! # Internal types (crate-visible only)
56//!
57//! `PrimaryBlock` and `DuplicateBlock` are RAII guards for the Registered state,
58//! used internally by the [`BlockRegistry`] to
59//! distinguish between the canonical holder of a sequence hash and any additional
60//! logical copies.
61
62mod complete;
63mod immutable;
64mod mutable;
65mod registered;
66
67pub use complete::CompleteBlock;
68pub use immutable::{ImmutableBlock, WeakBlock};
69pub use mutable::MutableBlock;
70
71pub(crate) mod state;
72pub(crate) use registered::{DuplicateBlock, PrimaryBlock, WeakBlockEntry};
73
74// Re-export from the new registry module location for backward compatibility
75pub use crate::registry::BlockRegistrationHandle;
76pub use crate::registry::BlockRegistry;
77
78/// Marker trait for types that can serve as block-level metadata.
79///
80/// A blanket implementation covers every type that is
81/// `Clone + Send + Sync + 'static`, so callers rarely need to think about
82/// this trait directly -- any ordinary data type already satisfies it.
83pub trait BlockMetadata: Clone + Send + Sync + 'static {}
84impl<T: Clone + Send + Sync + 'static> BlockMetadata for T {}
85
86/// Logical Block Identifier
87pub use crate::{BlockId, SequenceHash};
88use dynamo_tokens::TokenBlock;
89use std::sync::Arc;
90
91/// Return function for blocks transitioning back to Reset state.
92/// Used by MutableBlock, DuplicateBlock, and CompleteBlock drop paths.
93pub(crate) type ResetReturnFn<T> = Arc<dyn Fn(Block<T, state::Reset>) + Send + Sync>;
94
95/// Return function for registered blocks returning to the inactive pool.
96/// Used by PrimaryBlock drop and pool management.
97pub(crate) type RegisteredReturnFn<T> = Arc<dyn Fn(Arc<Block<T, state::Registered>>) + Send + Sync>;
98
99/// Upgrade function for finding/promoting blocks by sequence hash.
100/// Used by ImmutableBlock and WeakBlock upgrade paths.
101pub(crate) type UpgradeFn<T> =
102    Arc<dyn Fn(SequenceHash) -> Option<Arc<dyn RegisteredBlock<T>>> + Send + Sync>;
103
104/// Error returned by block state transitions.
105///
106/// Every variant carries the originating block back to the caller so that
107/// the block is never silently leaked on failure. The caller can inspect the
108/// error, recover the block from the variant, and retry or drop it as needed.
109#[derive(Debug, thiserror::Error)]
110pub enum BlockError<B> {
111    /// The number of tokens in the provided data did not match the block's
112    /// fixed size. The block is returned in the `block` field.
113    #[error("Block size mismatch: expected {expected} tokens, got {actual}")]
114    BlockSizeMismatch {
115        expected: usize,
116        actual: usize,
117        block: B,
118    },
119}
120
121/// Controls whether the [`BlockRegistry`] accepts
122/// multiple physical blocks that share the same [`SequenceHash`].
123///
124/// The policy is set once when a
125/// [`BlockManager`](crate::manager::BlockManager) is built and applies to
126/// every subsequent registration.
127#[derive(Debug, Clone, Copy, PartialEq, Eq)]
128pub enum BlockDuplicationPolicy {
129    /// Multiple physical blocks may hold the same data for a single logical
130    /// block / sequence hash.
131    ///
132    /// This is the typical choice for G1 / device memory in LLM inference
133    /// frameworks.
134    Allow,
135    /// Only one physical block is retained per sequence hash. Attempting to
136    /// register a duplicate returns the existing primary block instead.
137    ///
138    /// Used internally by KVBM for G2+ storage layers (host, disk,
139    /// distributed, object store) where deduplication saves capacity.
140    Reject,
141}
142
143/// The raw block value parameterised by metadata `T` and a type-state marker.
144///
145/// External code never sees this type directly; it is always wrapped in one of
146/// the public RAII guards ([`MutableBlock`], [`CompleteBlock`], [`ImmutableBlock`]).
147#[derive(Debug)]
148pub(crate) struct Block<T, State> {
149    block_id: BlockId,
150    block_size: usize,
151    state: State,
152    marker: std::marker::PhantomData<T>,
153}
154
155/// Common interface for blocks in the Registered state.
156///
157/// Both [`PrimaryBlock`] and [`DuplicateBlock`] implement this trait so that
158/// [`ImmutableBlock`] can hold either variant behind a `dyn RegisteredBlock<T>`.
159pub(crate) trait RegisteredBlock<T>: Send + Sync {
160    fn block_id(&self) -> BlockId;
161    fn sequence_hash(&self) -> SequenceHash;
162    fn registration_handle(&self) -> &BlockRegistrationHandle;
163}
kvbm_logical/blocks/mod.rs

kvbm_logical/blocks/
mod.rs
