kvbm_logical/manager/
builder.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Builder and configuration types for [`BlockManager`](super::BlockManager).
5
6use std::num::NonZeroUsize;
7use std::sync::Arc;
8
9use crate::metrics::{BlockPoolMetrics, MetricsAggregator, short_type_name};
10use crate::tinylfu::TinyLFUTracker;
11
12use crate::{
13    blocks::BlockMetadata,
14    pools::{
15        BlockDuplicationPolicy, BlockStore, InactiveIndex,
16        backends::{
17            FifoReusePolicy, HashMapBackend, LeafPolicy, LineageBackend, LruBackend,
18            MultiLruBackend,
19        },
20    },
21    registry::BlockRegistry,
22};
23
24use super::BlockManager;
25
26/// Capacity settings for the TinyLFU frequency tracker used by
27/// [`BlockRegistry`] and the multi-level LRU backend.
28#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
29pub enum FrequencyTrackingCapacity {
30    /// Small capacity: 2^18 (262,144) entries
31    Small,
32    /// Medium capacity: 2^21 (2,097,152) entries - default
33    #[default]
34    Medium,
35    /// Large capacity: 2^24 (16,777,216) entries
36    Large,
37}
38
39impl FrequencyTrackingCapacity {
40    /// Get the size in number of entries.
41    pub fn size(&self) -> usize {
42        match self {
43            Self::Small => 1 << 18,
44            Self::Medium => 1 << 21,
45            Self::Large => 1 << 24,
46        }
47    }
48
49    /// Create a new [`TinyLFUTracker`] with this capacity.
50    pub fn create_tracker(&self) -> Arc<TinyLFUTracker<u128>> {
51        Arc::new(TinyLFUTracker::new(self.size()))
52    }
53}
54
55/// Configuration for the inactive pool backend.
56pub enum InactiveBackendConfig {
57    /// HashMap with FIFO reuse order.
58    HashMap,
59    /// Simple LRU — capacity automatically set to block_count.
60    Lru,
61    /// Multi-level LRU with 4 fixed levels — capacity automatically set to block_count.
62    MultiLru {
63        /// Frequency thresholds: [cold->warm, warm->hot, hot->very_hot].
64        /// Default: [3, 8, 15].
65        frequency_thresholds: [u8; 3],
66    },
67    /// Lineage backend with a selectable leaf-eviction policy.
68    Lineage {
69        /// Leaf-eviction ordering. Default: [`LineageEviction::Tick`].
70        eviction: LineageEviction,
71    },
72}
73
74/// Leaf-eviction ordering for the [`Lineage`](InactiveBackendConfig::Lineage)
75/// inactive backend.
76#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
77pub enum LineageEviction {
78    /// `BTreeMap` ordered by a per-node insertion tick — a node that
79    /// re-becomes a leaf returns to its original position. Historical
80    /// behavior; the default. O(log n) per hook, with B-tree node churn.
81    #[default]
82    Tick,
83    /// Intrusive FIFO over leaves — O(1) and allocation-free, but a node
84    /// that re-becomes a leaf is appended at the tail.
85    Fifo,
86}
87
88/// Build the runtime [`LeafPolicy`] for a [`LineageEviction`] selection.
89fn lineage_leaf_policy(eviction: LineageEviction, capacity: usize) -> LeafPolicy {
90    match eviction {
91        LineageEviction::Tick => LeafPolicy::tick(capacity),
92        LineageEviction::Fifo => LeafPolicy::fifo(capacity),
93    }
94}
95
96/// Error types for [`BlockManager`] builder validation.
97#[derive(Debug, thiserror::Error)]
98pub enum BlockManagerBuilderError {
99    #[error("Block count must be greater than 0")]
100    InvalidBlockCount,
101    #[error("Block size mismatch: expected {expected} tokens, got {actual}")]
102    BlockSizeMismatch { expected: usize, actual: usize },
103    #[error("Invalid backend configuration: {0}")]
104    InvalidBackend(String),
105    #[error("Builder validation failed: {0}")]
106    ValidationError(String),
107}
108
109/// Error types for [`BlockManager::reset_inactive_pool`].
110#[derive(Debug, thiserror::Error)]
111pub enum BlockManagerResetError {
112    #[error("Reset pool count mismatch: expected {expected}, got {actual}")]
113    BlockCountMismatch { expected: usize, actual: usize },
114}
115
116/// Builder for [`BlockManager`] configuration.
117///
118/// Construct via [`BlockManager::builder()`] and finish with [`build()`](Self::build).
119pub struct BlockManagerConfigBuilder<T: BlockMetadata> {
120    /// Number of blocks in the pool
121    block_count: Option<usize>,
122
123    /// Size of each block in tokens (must be power of 2, 1-1024)
124    /// Default: 16
125    block_size: Option<usize>,
126
127    /// Block registry for tracking blocks and frequency
128    registry: Option<BlockRegistry>,
129
130    /// Inactive pool backend configuration
131    inactive_backend: Option<InactiveBackendConfig>,
132
133    /// Policy for handling duplicate sequence hashes
134    duplication_policy: Option<BlockDuplicationPolicy>,
135
136    /// Optional metrics aggregator for prometheus export
137    aggregator: Option<MetricsAggregator>,
138
139    /// Default value of the per-block `reset_on_release` flag. When
140    /// `Some(true)`, every `ImmutableBlock` constructed by this manager
141    /// starts with the flag set, so its last drop bypasses the inactive
142    /// pool and resets the slot directly. Individual blocks can still
143    /// override via [`ImmutableBlock::set_evict_on_reset`].
144    default_reset_on_release: Option<bool>,
145
146    /// Phantom data for type parameter
147    _phantom: std::marker::PhantomData<T>,
148}
149
150impl<T: BlockMetadata> Default for BlockManagerConfigBuilder<T> {
151    fn default() -> Self {
152        Self {
153            block_count: None,
154            block_size: Some(16), // Default to 16 tokens per block
155            registry: None,
156            inactive_backend: None,
157            duplication_policy: None,
158            aggregator: None,
159            default_reset_on_release: None,
160            _phantom: std::marker::PhantomData,
161        }
162    }
163}
164
165impl<T: BlockMetadata> BlockManagerConfigBuilder<T> {
166    /// Create a new builder.
167    pub fn new() -> Self {
168        Self::default()
169    }
170
171    /// Set the number of blocks in the pool.
172    pub fn block_count(mut self, count: usize) -> Self {
173        self.block_count = Some(count);
174        self
175    }
176
177    /// Set the block size (number of tokens per block).
178    ///
179    /// # Requirements
180    /// - Must be >= 1 and <= 1024
181    /// - Must be a power of 2
182    ///
183    /// # Panics
184    /// Panics if the block size doesn't meet requirements.
185    pub fn block_size(mut self, size: usize) -> Self {
186        assert!(
187            (1..=1024).contains(&size),
188            "block_size must be between 1 and 1024, got {}",
189            size
190        );
191        assert!(
192            size.is_power_of_two(),
193            "block_size must be a power of 2, got {}",
194            size
195        );
196        self.block_size = Some(size);
197        self
198    }
199
200    /// Set the duplication policy.
201    pub fn duplication_policy(mut self, policy: BlockDuplicationPolicy) -> Self {
202        self.duplication_policy = Some(policy);
203        self
204    }
205
206    /// Set the block registry.
207    pub fn registry(mut self, registry: BlockRegistry) -> Self {
208        self.registry = Some(registry);
209        self
210    }
211
212    /// Use simple LRU backend (capacity automatically set to block_count).
213    pub fn with_lru_backend(mut self) -> Self {
214        self.inactive_backend = Some(InactiveBackendConfig::Lru);
215        self
216    }
217
218    /// Use multi-level LRU backend with 4 fixed priority levels.
219    ///
220    /// Default thresholds: `[3, 8, 15]` for transitions between:
221    /// Cold (0-2 hits) -> Warm (3-7) -> Hot (8-14) -> Very Hot (15+).
222    pub fn with_multi_lru_backend(mut self) -> Self {
223        self.inactive_backend = Some(InactiveBackendConfig::MultiLru {
224            frequency_thresholds: [3, 8, 15],
225        });
226        self
227    }
228
229    /// Use multi-level LRU with custom frequency thresholds.
230    ///
231    /// # Requirements
232    /// - Thresholds must be in ascending order: cold_to_warm < warm_to_hot < hot_to_very_hot
233    /// - hot_to_very_hot must be <= 15 (4-bit counter maximum)
234    /// - cold_to_warm must be >= 1 (to distinguish from never-accessed blocks)
235    ///
236    /// # Arguments
237    /// * `cold_to_warm` - Minimum frequency to move from Cold to Warm level
238    /// * `warm_to_hot` - Minimum frequency to move from Warm to Hot level
239    /// * `hot_to_very_hot` - Minimum frequency to move from Hot to Very Hot level
240    ///
241    /// # Panics
242    /// Panics if thresholds don't meet the requirements above.
243    pub fn with_multi_lru_backend_custom_thresholds(
244        mut self,
245        cold_to_warm: u8,
246        warm_to_hot: u8,
247        hot_to_very_hot: u8,
248    ) -> Self {
249        // Validate ascending order
250        assert!(
251            cold_to_warm < warm_to_hot && warm_to_hot < hot_to_very_hot,
252            "Thresholds must be in ascending order: {} < {} < {} failed",
253            cold_to_warm,
254            warm_to_hot,
255            hot_to_very_hot
256        );
257
258        // Validate maximum value (4-bit counter limit)
259        assert!(
260            hot_to_very_hot <= 15,
261            "hot_to_very_hot threshold ({}) must be <= 15 (4-bit counter maximum)",
262            hot_to_very_hot
263        );
264
265        // Additional validation: ensure reasonable gaps between levels
266        assert!(
267            cold_to_warm >= 1,
268            "cold_to_warm threshold must be >= 1 to distinguish from zero-access blocks"
269        );
270
271        self.inactive_backend = Some(InactiveBackendConfig::MultiLru {
272            frequency_thresholds: [cold_to_warm, warm_to_hot, hot_to_very_hot],
273        });
274        self
275    }
276
277    /// Use HashMap backend with FIFO reuse order.
278    pub fn with_hashmap_backend(mut self) -> Self {
279        self.inactive_backend = Some(InactiveBackendConfig::HashMap);
280        self
281    }
282
283    /// Use the lineage backend with the default ([`Tick`](LineageEviction::Tick))
284    /// leaf-eviction policy.
285    pub fn with_lineage_backend(mut self) -> Self {
286        self.inactive_backend = Some(InactiveBackendConfig::Lineage {
287            eviction: LineageEviction::default(),
288        });
289        self
290    }
291
292    /// Use the lineage backend with an explicit leaf-eviction policy.
293    pub fn with_lineage_backend_eviction(mut self, eviction: LineageEviction) -> Self {
294        self.inactive_backend = Some(InactiveBackendConfig::Lineage { eviction });
295        self
296    }
297
298    /// Set a metrics aggregator for prometheus export.
299    ///
300    /// The aggregator will automatically receive this manager's metrics source.
301    pub fn aggregator(mut self, aggregator: MetricsAggregator) -> Self {
302        self.aggregator = Some(aggregator);
303        self
304    }
305
306    /// Set the default value of the per-block `reset_on_release` flag.
307    ///
308    /// When `true`, every `ImmutableBlock` constructed by this manager
309    /// starts with the flag set. On its last drop, the slot bypasses
310    /// the inactive pool and is reset back to the free list directly
311    /// (matching `release_duplicate` semantics for primary releases).
312    ///
313    /// Individual blocks can still override via
314    /// [`crate::blocks::ImmutableBlock::set_evict_on_reset`].
315    ///
316    /// Default: `false`.
317    pub fn with_default_reset_on_release(mut self, value: bool) -> Self {
318        self.default_reset_on_release = Some(value);
319        self
320    }
321
322    /// Validate the configuration.
323    fn validate(&self) -> Result<(), String> {
324        let registry = self.registry.as_ref().ok_or("registry is required")?;
325
326        let block_count = self.block_count.ok_or("block_count is required")?;
327
328        if block_count == 0 {
329            return Err("block_count must be greater than 0".to_string());
330        }
331
332        // Validate block_size
333        let block_size = self.block_size.unwrap_or(16);
334        if !block_size.is_power_of_two() || !(1..=1024).contains(&block_size) {
335            return Err(format!(
336                "Invalid block_size {}: must be a power of 2 between 1 and 1024",
337                block_size
338            ));
339        }
340
341        // Additional validation for MultiLRU thresholds at build time
342        if let Some(InactiveBackendConfig::MultiLru {
343            frequency_thresholds,
344        }) = &self.inactive_backend
345        {
346            let [t1, t2, t3] = frequency_thresholds;
347            if !(*t1 < *t2 && *t2 < *t3) {
348                return Err(format!(
349                    "Invalid thresholds [{}, {}, {}]: must be in ascending order",
350                    t1, t2, t3
351                ));
352            }
353            if *t3 > 15 {
354                return Err(format!(
355                    "Invalid threshold {}: maximum frequency is 15 (4-bit counter)",
356                    t3
357                ));
358            }
359
360            // Validate MultiLRU requires frequency tracking
361            if !registry.has_frequency_tracking() {
362                return Err(
363                    "MultiLRU backend requires a registry with frequency tracking".to_string(),
364                );
365            }
366        }
367
368        Ok(())
369    }
370
371    /// Build the [`BlockManager`].
372    ///
373    /// Validates configuration and constructs all pools, the upgrade closure,
374    /// and the metrics source. Returns an error if validation fails or
375    /// backend construction fails.
376    pub fn build(mut self) -> Result<BlockManager<T>, BlockManagerBuilderError> {
377        // First validate the configuration
378        self.validate()
379            .map_err(BlockManagerBuilderError::ValidationError)?;
380
381        let block_count = self.block_count.unwrap();
382        let block_size = self.block_size.unwrap_or(16);
383
384        // Use provided registry
385        let registry = self.registry.unwrap();
386
387        // Create metrics
388        let metrics = Arc::new(BlockPoolMetrics::new(short_type_name::<T>()));
389
390        metrics.set_reset_pool_size(block_count as i64);
391
392        // Create backend based on configuration
393        let backend: Box<dyn InactiveIndex> = match self.inactive_backend.take() {
394            Some(InactiveBackendConfig::HashMap) => {
395                tracing::info!("Using HashMap for inactive pool");
396                Box::new(HashMapBackend::new(Box::new(FifoReusePolicy::new())))
397            }
398            Some(InactiveBackendConfig::Lru) => {
399                // Capacity automatically set to block_count
400                let capacity = NonZeroUsize::new(block_count).expect("block_count must be > 0");
401                tracing::info!("Using LRU for inactive pool");
402                Box::new(LruBackend::new(capacity))
403            }
404            Some(InactiveBackendConfig::MultiLru {
405                frequency_thresholds,
406            }) => {
407                // Require frequency tracker for MultiLRU
408                let frequency_tracker = registry.frequency_tracker().ok_or_else(|| {
409                    BlockManagerBuilderError::InvalidBackend(
410                        "MultiLRU backend requires a registry with frequency tracking".to_string(),
411                    )
412                })?;
413
414                // Each level needs capacity for all blocks since the frequency
415                // distribution is unpredictable — all blocks could land in one level.
416                let level_capacity =
417                    NonZeroUsize::new(block_count).expect("block_count must be > 0");
418
419                tracing::info!(
420                    "Using MultiLRU inactive backend with thresholds: {:?}",
421                    frequency_thresholds
422                );
423                Box::new(
424                    MultiLruBackend::new_with_thresholds(
425                        level_capacity,
426                        &frequency_thresholds,
427                        frequency_tracker,
428                    )
429                    .map_err(|e| BlockManagerBuilderError::InvalidBackend(e.to_string()))?,
430                )
431            }
432            Some(InactiveBackendConfig::Lineage { eviction }) => {
433                tracing::info!("Using Lineage inactive backend ({eviction:?})");
434                Box::new(LineageBackend::with_policy(
435                    block_count,
436                    lineage_leaf_policy(eviction, block_count),
437                ))
438            }
439            None => {
440                let eviction = LineageEviction::default();
441                tracing::info!("Using default inactive backend: Lineage ({eviction:?})");
442                Box::new(LineageBackend::with_policy(
443                    block_count,
444                    lineage_leaf_policy(eviction, block_count),
445                ))
446            }
447        };
448
449        // Construct unified store
450        let store = BlockStore::new(
451            block_count,
452            block_size,
453            backend,
454            metrics.clone(),
455            self.default_reset_on_release.unwrap_or(false),
456        );
457
458        // Register with aggregator if provided
459        if let Some(ref aggregator) = self.aggregator {
460            aggregator.register_source(metrics.clone());
461        }
462
463        Ok(BlockManager {
464            store,
465            block_registry: registry,
466            duplication_policy: self
467                .duplication_policy
468                .unwrap_or(BlockDuplicationPolicy::Allow),
469            total_blocks: block_count,
470            block_size,
471            metrics,
472        })
473    }
474}
kvbm_logical/manager/builder.rs

kvbm_logical/manager/
builder.rs
