kvbm_logical/manager/
mod.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Block lifecycle orchestration over the unified [`BlockStore`].
5//!
6//! [`BlockManager`] owns a single [`BlockStore`] and the [`BlockRegistry`].
7//! All pool transitions go through the store's single mutex; the manager
8//! adds the registry coordination, allocation eviction policy, and metrics.
9
10mod builder;
11
12#[cfg(test)]
13mod tests;
14
15pub use builder::{
16    BlockManagerBuilderError, BlockManagerConfigBuilder, BlockManagerResetError,
17    FrequencyTrackingCapacity, InactiveBackendConfig, LineageEviction,
18};
19
20use std::collections::HashMap;
21use std::sync::Arc;
22
23use crate::blocks::{BlockMetadata, CompleteBlock, ImmutableBlock, MutableBlock};
24use crate::metrics::BlockPoolMetrics;
25use crate::pools::{BlockDuplicationPolicy, BlockStore, SequenceHash};
26use crate::registry::BlockRegistry;
27
28/// Manages the full block lifecycle over the unified [`BlockStore`].
29///
30/// Construct via [`BlockManager::builder()`].
31pub struct BlockManager<T: BlockMetadata> {
32    pub(crate) store: Arc<BlockStore<T>>,
33    pub(crate) block_registry: BlockRegistry,
34    pub(crate) duplication_policy: BlockDuplicationPolicy,
35    pub(crate) total_blocks: usize,
36    pub(crate) block_size: usize,
37    pub(crate) metrics: Arc<BlockPoolMetrics>,
38}
39
40impl<T: BlockMetadata + Sync> BlockManager<T> {
41    /// Create a new builder for `BlockManager`.
42    pub fn builder() -> BlockManagerConfigBuilder<T> {
43        BlockManagerConfigBuilder::default()
44    }
45
46    /// Stable, process-unique identifier for this manager's underlying
47    /// [`BlockStore`](crate::pools::BlockStore). See [`crate::ManagerId`].
48    /// Cheap (one field load via the store).
49    ///
50    /// Together with a [`BlockId`](crate::BlockId) this names a specific
51    /// physical pool slot — the disambiguating runtime address that
52    /// downstream consumers need after the policy parameter `T` has been
53    /// type-erased through [`crate::LifecyclePinRef`].
54    pub fn id(&self) -> crate::ManagerId {
55        self.store.id()
56    }
57
58    /// Allocate `count` mutable blocks, drawing first from the reset pool
59    /// and then evicting from the inactive pool if needed.
60    ///
61    /// Returns `None` if fewer than `count` blocks are available across both pools.
62    pub fn allocate_blocks(&self, count: usize) -> Option<Vec<MutableBlock<T>>> {
63        self.allocate_blocks_with_evictions(count)
64            .map(|(blocks, _evicted)| blocks)
65    }
66
67    /// Like [`allocate_blocks`](Self::allocate_blocks) but also reports the
68    /// [`SequenceHash`] of each block evicted from the inactive pool.
69    pub fn allocate_blocks_with_evictions(
70        &self,
71        count: usize,
72    ) -> Option<(Vec<MutableBlock<T>>, Vec<SequenceHash>)> {
73        self.store.allocate_atomic(count)
74    }
75
76    /// Drain the inactive pool, returning all blocks to the reset pool.
77    pub fn reset_inactive_pool(&self) -> Result<(), BlockManagerResetError> {
78        let blocks = self.store.drain_inactive_to_mutable();
79        drop(blocks);
80
81        let reset_count = self.store.reset_len();
82        if reset_count != self.total_blocks {
83            return Err(BlockManagerResetError::BlockCountMismatch {
84                expected: self.total_blocks,
85                actual: reset_count,
86            });
87        }
88
89        Ok(())
90    }
91
92    /// Register a batch of completed blocks.
93    pub fn register_blocks(&self, blocks: Vec<CompleteBlock<T>>) -> Vec<ImmutableBlock<T>> {
94        if blocks.is_empty() {
95            return Vec::new();
96        }
97
98        let handles = self
99            .block_registry
100            .register_sequence_hashes(blocks.iter().map(CompleteBlock::sequence_hash));
101        let batch_size = blocks.len();
102        let registered = self.store.register_completed_blocks(
103            blocks.into_iter().zip(handles).collect(),
104            self.duplication_policy,
105        );
106        // The offline settlement bridge observes this counter as a
107        // publication watermark, so publish only after every store transition
108        // and presence marker in the batch is complete.
109        self.metrics
110            .inc_registrations_by(u64::try_from(batch_size).unwrap_or(u64::MAX));
111        registered
112            .into_iter()
113            .map(ImmutableBlock::from_inner)
114            .collect()
115    }
116
117    /// Register a single completed block and return an immutable handle.
118    pub fn register_block(&self, block: CompleteBlock<T>) -> ImmutableBlock<T> {
119        let handle = self
120            .block_registry
121            .register_sequence_hash(block.sequence_hash());
122        let inner = handle.register_block(block, self.duplication_policy, &self.store);
123        self.metrics.inc_registrations();
124        ImmutableBlock::from_inner(inner)
125    }
126
127    /// Linear prefix match: walks `seq_hash` left-to-right, stopping on
128    /// the first hash that hits neither the active nor the inactive pool.
129    ///
130    /// The whole active-or-inactive prefix is resolved under a **single**
131    /// store-mutex acquisition via [`BlockStore::match_prefix_locked_batch`]
132    /// — no per-hash registry radix-tree lookup, no per-hash store lock.
133    /// Frequency-tracker touches are batched and applied *after* the store
134    /// lock is released: every returned block is touched exactly once
135    /// (including inactive resurrections).
136    pub fn match_blocks(&self, seq_hash: &[SequenceHash]) -> Vec<ImmutableBlock<T>> {
137        self.metrics
138            .inc_match_hashes_requested(seq_hash.len() as u64);
139
140        if seq_hash.is_empty() {
141            self.metrics.inc_match_blocks_returned(0);
142            return Vec::new();
143        }
144
145        // ONE store-lock acquisition for the whole active+inactive prefix.
146        let inners = self.store.match_prefix_locked_batch(seq_hash);
147
148        // Frequency-tracker touches, batched, AFTER the store lock is
149        // released. Touches every returned hit exactly once — including
150        // inactive resurrections, which the old `find_inactive_primaries`
151        // path never touched.
152        if self.block_registry.has_frequency_tracking() {
153            for inner in &inners {
154                self.block_registry.touch(inner.sequence_hash());
155            }
156        }
157
158        let matched: Vec<ImmutableBlock<T>> =
159            inners.into_iter().map(ImmutableBlock::from_inner).collect();
160
161        self.metrics.inc_match_blocks_returned(matched.len() as u64);
162        tracing::debug!(
163            num_hashes = seq_hash.len(),
164            total_matched = matched.len(),
165            "match_blocks result"
166        );
167        tracing::trace!(matched = ?matched, "matched blocks");
168        matched
169    }
170
171    /// Scatter-gather scan: finds all blocks matching any hash, without
172    /// stopping on misses.
173    pub fn scan_matches(
174        &self,
175        seq_hashes: &[SequenceHash],
176        touch: bool,
177    ) -> HashMap<SequenceHash, ImmutableBlock<T>> {
178        self.metrics
179            .inc_scan_hashes_requested(seq_hashes.len() as u64);
180
181        let mut result = HashMap::new();
182
183        let active_found = self.scan_active_matches(seq_hashes, touch);
184        for (hash, inner) in active_found {
185            result.insert(hash, ImmutableBlock::from_inner(inner));
186        }
187
188        let remaining: Vec<SequenceHash> = seq_hashes
189            .iter()
190            .filter(|h| !result.contains_key(h))
191            .copied()
192            .collect();
193
194        if !remaining.is_empty() {
195            let inactive_found = self.store.scan_inactive_primaries(&remaining, touch);
196            for (hash, inner) in inactive_found {
197                result.insert(hash, ImmutableBlock::from_inner(inner));
198            }
199        }
200
201        self.metrics.inc_scan_blocks_returned(result.len() as u64);
202
203        result
204    }
205
206    /// Scan-style active lookup by sequence hash via the registry's
207    /// stored Weak references — does not stop on miss.
208    fn scan_active_matches(
209        &self,
210        hashes: &[SequenceHash],
211        touch: bool,
212    ) -> Vec<(SequenceHash, Arc<crate::blocks::ImmutableBlockInner<T>>)> {
213        hashes
214            .iter()
215            .filter_map(|hash| {
216                self.block_registry
217                    .match_sequence_hash(*hash, touch)
218                    .and_then(|handle| {
219                        handle
220                            .try_get_inner::<T>(&self.store, touch)
221                            .map(|inner| (*hash, inner))
222                    })
223            })
224            .collect()
225    }
226
227    /// Total number of blocks managed (constant after construction).
228    pub fn total_blocks(&self) -> usize {
229        self.total_blocks
230    }
231
232    /// Blocks available for allocation (reset + inactive pools).
233    ///
234    /// Reads both pool sizes under a single store-lock acquisition so the
235    /// returned value is a coherent snapshot, never an over- or under-count
236    /// produced by a concurrent reset↔inactive transition.
237    pub fn available_blocks(&self) -> usize {
238        self.store.available_len()
239    }
240
241    /// Tokens per block (constant after construction).
242    pub fn block_size(&self) -> usize {
243        self.block_size
244    }
245
246    /// Current duplication policy.
247    pub fn duplication_policy(&self) -> &BlockDuplicationPolicy {
248        &self.duplication_policy
249    }
250
251    /// Reference to the shared block registry.
252    pub fn block_registry(&self) -> &BlockRegistry {
253        &self.block_registry
254    }
255
256    /// Reference to the block pool metrics.
257    pub fn metrics(&self) -> &Arc<BlockPoolMetrics> {
258        &self.metrics
259    }
260
261    /// Test-only accessor for the underlying [`BlockStore`]. Used to
262    /// reach test hooks like `BlockStore::pause_release_primary` from
263    /// race-window tests.
264    #[cfg(test)]
265    pub(crate) fn store_for_test(&self) -> &Arc<BlockStore<T>> {
266        &self.store
267    }
268}
kvbm_logical/manager/mod.rs

kvbm_logical/manager/
mod.rs
