kvbm_logical/blocks/
immutable.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! RAII guards for blocks in the **Registered** state.
5//!
6//! [`ImmutableBlock`] is the strong, cloneable handle that keeps a registered
7//! block alive. [`WeakBlock`] is its non-owning counterpart -- it does not
8//! prevent the block from being evicted, but can be cheaply upgraded back to
9//! an `ImmutableBlock` if the block is still present.
10
11use super::{
12    BlockId, BlockMetadata, BlockRegistrationHandle, RegisteredBlock, SequenceHash, UpgradeFn,
13};
14
15use crate::metrics::BlockPoolMetrics;
16use std::sync::{Arc, Weak};
17
18/// RAII guard for a block in the **Registered** state.
19///
20/// An `ImmutableBlock` is the primary handle through which callers interact
21/// with registered blocks. It is reference-counted (`Clone`-able) and each
22/// clone independently tracks the `inflight_immutable` metric gauge, so the
23/// gauge reflects the total number of outstanding references across the
24/// system.
25///
26/// # Obtaining an `ImmutableBlock`
27///
28/// - [`BlockManager::register_block`](crate::manager::BlockManager::register_block)
29///   -- registers a [`CompleteBlock`](super::CompleteBlock) and returns an
30///   `ImmutableBlock`.
31/// - [`BlockManager::match_blocks`](crate::manager::BlockManager::match_blocks)
32///   / [`BlockManager::scan_matches`](crate::manager::BlockManager::scan_matches)
33///   -- look up already-registered blocks by [`SequenceHash`].
34/// - [`WeakBlock::upgrade`] -- resurrects a weak reference if the block is
35///   still alive.
36///
37/// # State transitions
38///
39/// - [`downgrade`](Self::downgrade) -- creates a [`WeakBlock`] that does not
40///   keep the block alive.
41///
42/// # Clone behaviour
43///
44/// Cloning an `ImmutableBlock` increments `inflight_immutable`; dropping a
45/// clone decrements it. The underlying registered block is shared via
46/// `Arc`, so clones are cheap.
47///
48/// # Drop behaviour
49///
50/// Dropping the last strong reference (including internal pool references)
51/// triggers the block's return to the inactive or reset pool. Every drop
52/// decrements the `inflight_immutable` gauge.
53pub struct ImmutableBlock<T: BlockMetadata> {
54    block: Arc<dyn RegisteredBlock<T>>,
55    upgrade_fn: UpgradeFn<T>,
56    metrics: Option<Arc<BlockPoolMetrics>>,
57}
58
59/// Non-owning reference to a registered block.
60///
61/// A `WeakBlock` does not keep the underlying block alive -- if all
62/// [`ImmutableBlock`] handles (and internal pool references) are dropped,
63/// the block may be evicted and the weak reference will fail to upgrade.
64///
65/// Created via [`ImmutableBlock::downgrade`]. Cloneable and cheap to hold.
66///
67/// Call [`upgrade`](Self::upgrade) to attempt to recover a full
68/// [`ImmutableBlock`].
69#[derive(Clone)]
70pub struct WeakBlock<T: BlockMetadata> {
71    sequence_hash: SequenceHash,
72    block: Weak<dyn RegisteredBlock<T>>,
73    upgrade_fn: UpgradeFn<T>,
74    metrics: Option<Arc<BlockPoolMetrics>>,
75}
76
77impl<T: BlockMetadata> ImmutableBlock<T> {
78    /// Create a new ImmutableBlock with an upgrade function
79    pub(crate) fn new(
80        block: Arc<dyn RegisteredBlock<T>>,
81        upgrade_fn: UpgradeFn<T>,
82        metrics: Option<Arc<BlockPoolMetrics>>,
83    ) -> Self {
84        if let Some(ref m) = metrics {
85            m.inc_inflight_immutable();
86        }
87        Self {
88            block,
89            upgrade_fn,
90            metrics,
91        }
92    }
93
94    /// Creates a [`WeakBlock`] that references the same registered block
95    /// without preventing it from being evicted.
96    pub fn downgrade(&self) -> WeakBlock<T> {
97        WeakBlock {
98            sequence_hash: self.sequence_hash(),
99            block: Arc::downgrade(&self.block),
100            upgrade_fn: self.upgrade_fn.clone(),
101            metrics: self.metrics.clone(),
102        }
103    }
104
105    /// Returns the [`BlockId`] assigned to this block.
106    pub fn block_id(&self) -> BlockId {
107        self.block.block_id()
108    }
109
110    /// Returns the [`SequenceHash`] that identifies this block's content.
111    pub fn sequence_hash(&self) -> SequenceHash {
112        self.block.sequence_hash()
113    }
114
115    /// Returns a clone of the [`BlockRegistrationHandle`] for this block.
116    pub fn registration_handle(&self) -> BlockRegistrationHandle {
117        self.block.registration_handle().clone()
118    }
119
120    /// Returns the number of strong (`Arc`) references to the underlying
121    /// registered block, including internal pool references.
122    pub fn use_count(&self) -> usize {
123        Arc::strong_count(&self.block)
124    }
125}
126
127impl<T: BlockMetadata> Clone for ImmutableBlock<T> {
128    fn clone(&self) -> Self {
129        if let Some(ref m) = self.metrics {
130            m.inc_inflight_immutable();
131        }
132        Self {
133            block: self.block.clone(),
134            upgrade_fn: self.upgrade_fn.clone(),
135            metrics: self.metrics.clone(),
136        }
137    }
138}
139
140impl<T: BlockMetadata> Drop for ImmutableBlock<T> {
141    #[inline]
142    fn drop(&mut self) {
143        if let Some(ref m) = self.metrics {
144            m.dec_inflight_immutable();
145        }
146    }
147}
148
149impl<T: BlockMetadata> std::fmt::Debug for ImmutableBlock<T> {
150    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
151        f.debug_struct("ImmutableBlock")
152            .field("block_id", &self.block_id())
153            .field("sequence_hash", &self.sequence_hash())
154            .finish()
155    }
156}
157
158impl<T: BlockMetadata> WeakBlock<T> {
159    /// Attempts to upgrade this weak reference back to an [`ImmutableBlock`].
160    ///
161    /// Uses a two-phase strategy:
162    /// 1. Tries a direct `Weak::upgrade` on the stored pointer (fast path).
163    /// 2. Falls back to searching the
164    ///    [`BlockRegistry`](crate::registry::BlockRegistry) by
165    ///    [`SequenceHash`] in case the block was moved between pools.
166    ///
167    /// Returns `None` if the block has been fully evicted.
168    pub fn upgrade(&self) -> Option<ImmutableBlock<T>> {
169        // Fast path: direct weak upgrade
170        if let Some(block) = self.block.upgrade() {
171            return Some(ImmutableBlock::new(
172                block,
173                self.upgrade_fn.clone(),
174                self.metrics.clone(),
175            ));
176        }
177
178        // Slow path: search the registry by sequence hash
179        if let Some(block) = (self.upgrade_fn)(self.sequence_hash) {
180            return Some(ImmutableBlock::new(
181                block,
182                self.upgrade_fn.clone(),
183                self.metrics.clone(),
184            ));
185        }
186
187        None
188    }
189
190    /// Returns the [`SequenceHash`] for the block this weak reference
191    /// points to.
192    pub fn sequence_hash(&self) -> SequenceHash {
193        self.sequence_hash
194    }
195}
196
197impl<T: BlockMetadata> std::fmt::Debug for WeakBlock<T> {
198    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
199        f.debug_struct("WeakBlock")
200            .field("sequence_hash", &self.sequence_hash())
201            .finish()
202    }
203}
kvbm_logical/blocks/immutable.rs

kvbm_logical/blocks/
immutable.rs
