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`] represents a block that has been assigned a
7//! [`SequenceHash`](super::SequenceHash) but has not yet been registered in
8//! the [`BlockRegistry`](crate::registry::BlockRegistry). It is produced by
9//! [`MutableBlock::stage`](super::MutableBlock::stage) or
10//! [`MutableBlock::complete`](super::MutableBlock::complete) and can be
11//! either registered via
12//! [`BlockManager::register_block`](crate::manager::BlockManager::register_block)
13//! or rolled back to a [`MutableBlock`](super::MutableBlock) with
14//! [`reset`](CompleteBlock::reset).
15
16use super::{
17    Block, BlockId, BlockMetadata, MutableBlock, ResetReturnFn, SequenceHash, state::Staged,
18};
19
20/// RAII guard for a block in the **Staged** state.
21///
22/// Wraps an internal `Block<T, Staged>` -- a block that carries a
23/// [`SequenceHash`](super::SequenceHash) and is ready for registration.
24///
25/// # Obtaining a `CompleteBlock`
26///
27/// - [`MutableBlock::stage`](super::MutableBlock::stage) -- from a
28///   pre-computed [`SequenceHash`](super::SequenceHash).
29/// - [`MutableBlock::complete`](super::MutableBlock::complete) -- by
30///   extracting the hash from a [`TokenBlock`](dynamo_tokens::TokenBlock).
31///
32/// # State transitions
33///
34/// - Pass to [`BlockManager::register_block`](crate::manager::BlockManager::register_block)
35///   to move the block into the **Registered** state and receive an
36///   [`ImmutableBlock`](super::ImmutableBlock).
37/// - Call [`reset`](Self::reset) to undo the staging and get a
38///   [`MutableBlock`](super::MutableBlock) back (metrics are `None` on this
39///   path).
40///
41/// # Drop behaviour
42///
43/// If the `CompleteBlock` is dropped without being consumed, the underlying
44/// block is reset to the **Reset** state and returned to the reset pool.
45pub struct CompleteBlock<T: BlockMetadata> {
46    pub(crate) block: Option<Block<T, Staged>>,
47    pub(crate) return_fn: ResetReturnFn<T>,
48}
49
50impl<T: BlockMetadata> CompleteBlock<T> {
51    /// Create a new CompleteBlock
52    pub(crate) fn new(block: Block<T, Staged>, return_fn: ResetReturnFn<T>) -> Self {
53        Self {
54            block: Some(block),
55            return_fn,
56        }
57    }
58
59    /// Returns the [`BlockId`] assigned to this block.
60    pub fn block_id(&self) -> BlockId {
61        self.block.as_ref().unwrap().block_id()
62    }
63
64    /// Returns the [`SequenceHash`](super::SequenceHash) that was assigned
65    /// during staging.
66    pub fn sequence_hash(&self) -> SequenceHash {
67        self.block.as_ref().unwrap().sequence_hash()
68    }
69
70    /// Undoes the staging transition, returning a [`MutableBlock`] in the
71    /// **Reset** state.
72    ///
73    /// The returned `MutableBlock` does **not** carry metrics (they are set
74    /// to `None`) because this is an undo/rollback path rather than a fresh
75    /// allocation.
76    pub fn reset(mut self) -> MutableBlock<T> {
77        let block = self.block.take().unwrap().reset();
78
79        MutableBlock::new(block, self.return_fn.clone(), None)
80    }
81}
82
83impl<T: BlockMetadata> Drop for CompleteBlock<T> {
84    #[inline]
85    fn drop(&mut self) {
86        if let Some(block) = self.block.take() {
87            (self.return_fn)(block.reset());
88        }
89    }
90}
kvbm_logical/blocks/complete.rs

kvbm_logical/blocks/
complete.rs
