kvbm_logical/blocks/
complete.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 for a block in the **Staged** state.
5//!
6//! A [`CompleteBlock`] holds a slot that has been assigned a
7//! [`SequenceHash`] but has not yet been registered. It is produced by
8//! [`MutableBlock::stage`](super::MutableBlock::stage) or
9//! [`MutableBlock::complete`](super::MutableBlock::complete) and can be
10//! either registered via
11//! [`BlockManager::register_block`](crate::manager::BlockManager::register_block)
12//! or rolled back to a [`MutableBlock`] with [`reset`](CompleteBlock::reset).
13
14use std::sync::Arc;
15
16use crate::blocks::{BlockId, BlockMetadata, MutableBlock, SequenceHash};
17use crate::pools::BlockStore;
18
19/// RAII guard for a block in the **Staged** state.
20pub struct CompleteBlock<T: BlockMetadata> {
21    store: Arc<BlockStore<T>>,
22    block_id: BlockId,
23    block_size: usize,
24    seq_hash: SequenceHash,
25    /// `false` once the guard has been consumed by `reset` or by
26    /// `BlockManager::register_block`; Drop becomes a no-op.
27    armed: bool,
28}
29
30impl<T: BlockMetadata + Sync> CompleteBlock<T> {
31    pub(crate) fn from_store(
32        store: Arc<BlockStore<T>>,
33        block_id: BlockId,
34        block_size: usize,
35        seq_hash: SequenceHash,
36    ) -> Self {
37        Self {
38            store,
39            block_id,
40            block_size,
41            seq_hash,
42            armed: true,
43        }
44    }
45
46    /// Returns the [`BlockId`] assigned to this block.
47    pub fn block_id(&self) -> BlockId {
48        self.block_id
49    }
50
51    /// Returns the fixed block size of this block in tokens.
52    pub fn block_size(&self) -> usize {
53        self.block_size
54    }
55
56    /// Returns the [`SequenceHash`] assigned during staging.
57    pub fn sequence_hash(&self) -> SequenceHash {
58        self.seq_hash
59    }
60
61    /// Roll back to a [`MutableBlock`].
62    pub fn reset(mut self) -> MutableBlock<T> {
63        self.store.transition_back_to_mutable(self.block_id);
64        self.armed = false;
65        MutableBlock::from_store(self.store.clone(), self.block_id, self.block_size)
66    }
67
68    /// Disarm the guard so Drop is a no-op. Used by registration when
69    /// the slot will be transitioned `Staged → Primary` / `Duplicate`
70    /// under the store mutex; the caller takes responsibility for the
71    /// slot's onward state.
72    pub(crate) fn disarm(&mut self) {
73        self.armed = false;
74    }
75
76    /// Re-arm a previously disarmed guard so Drop releases the slot
77    /// `Staged → Reset`. Used by the registration `Reject` path after
78    /// the lookup commits to dropping the staged block.
79    pub(crate) fn rearm(&mut self) {
80        self.armed = true;
81    }
82}
83
84impl<T: BlockMetadata> Drop for CompleteBlock<T> {
85    #[inline]
86    fn drop(&mut self) {
87        if self.armed {
88            self.store.release_staged(self.block_id);
89        }
90    }
91}
92
93impl<T: BlockMetadata> std::fmt::Debug for CompleteBlock<T> {
94    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95        f.debug_struct("CompleteBlock")
96            .field("block_id", &self.block_id)
97            .field("seq_hash", &self.seq_hash)
98            .finish()
99    }
100}
kvbm_logical/blocks/complete.rs

kvbm_logical/blocks/
complete.rs
