kvbm_logical/blocks/
mutable.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 **Reset** state.
5//!
6//! A [`MutableBlock`] is the entry point of the block lifecycle. It is
7//! obtained from
8//! [`BlockManager::allocate_blocks`](crate::manager::BlockManager::allocate_blocks)
9//! or by calling [`CompleteBlock::reset`](super::CompleteBlock::reset), and
10//! can be advanced to a [`CompleteBlock`](super::CompleteBlock) via
11//! [`stage`](MutableBlock::stage) or [`complete`](MutableBlock::complete).
12
13use std::sync::Arc;
14
15use dynamo_tokens::TokenBlock;
16
17use crate::KvbmSequenceHashProvider;
18use crate::blocks::{BlockError, BlockId, BlockMetadata, CompleteBlock, SequenceHash};
19use crate::pools::BlockStore;
20
21/// RAII guard for a block in the **Reset** state.
22///
23/// Holds an `Arc<BlockStore<T>>` and a `BlockId`; the slot at that id is in
24/// `SlotState::Mutable` while this guard exists. Drop returns the slot to
25/// the reset pool.
26///
27/// # Drop behaviour
28///
29/// Dropping a `MutableBlock` returns the slot to the reset pool. The store's
30/// `release_mutable` updates the `inflight_mutable` gauge.
31pub struct MutableBlock<T: BlockMetadata> {
32    store: Arc<BlockStore<T>>,
33    block_id: BlockId,
34    block_size: usize,
35    /// `false` once the guard has been consumed by a state-transition
36    /// method (`stage` / `complete`); Drop becomes a no-op.
37    armed: bool,
38}
39
40impl<T: BlockMetadata + Sync> MutableBlock<T> {
41    /// Build a new `MutableBlock` for a slot the store has just transitioned
42    /// to `Mutable`. The store has already incremented the `inflight_mutable`
43    /// gauge as part of the transition.
44    pub(crate) fn from_store(
45        store: Arc<BlockStore<T>>,
46        block_id: BlockId,
47        block_size: usize,
48    ) -> Self {
49        Self {
50            store,
51            block_id,
52            block_size,
53            armed: true,
54        }
55    }
56
57    /// Returns the [`BlockId`] assigned to this block.
58    pub fn block_id(&self) -> BlockId {
59        self.block_id
60    }
61
62    /// Returns the fixed block size of this block in tokens.
63    pub fn block_size(&self) -> usize {
64        self.block_size
65    }
66
67    /// Transition from **Reset** to **Staged** with a pre-computed
68    /// [`SequenceHash`] and an explicit `block_size` check.
69    ///
70    /// On size mismatch returns `Err(`[`BlockError::BlockSizeMismatch`]`)`
71    /// containing this `MutableBlock` so the caller can recover it.
72    pub fn stage(
73        mut self,
74        seq_hash: SequenceHash,
75        block_size: usize,
76    ) -> Result<CompleteBlock<T>, BlockError<MutableBlock<T>>> {
77        if block_size != self.block_size {
78            return Err(BlockError::BlockSizeMismatch {
79                expected: self.block_size,
80                actual: block_size,
81                block: self,
82            });
83        }
84        self.store.transition_to_staged(self.block_id, seq_hash);
85        let id = self.block_id;
86        let bsize = self.block_size;
87        let store = self.store.clone();
88        // Disarm so Drop is a no-op; the slot is now in Staged state.
89        self.armed = false;
90        drop(self);
91        Ok(CompleteBlock::from_store(store, id, bsize, seq_hash))
92    }
93
94    /// Transition from **Reset** to **Staged** by extracting the
95    /// [`SequenceHash`] from a [`TokenBlock`].
96    pub fn complete(
97        self,
98        token_block: &TokenBlock,
99    ) -> Result<CompleteBlock<T>, BlockError<MutableBlock<T>>> {
100        let actual = token_block.block_size();
101        if actual != self.block_size {
102            return Err(BlockError::BlockSizeMismatch {
103                expected: self.block_size,
104                actual,
105                block: self,
106            });
107        }
108        let seq_hash = token_block.kvbm_sequence_hash();
109        self.stage(seq_hash, actual)
110    }
111}
112
113impl<T: BlockMetadata> Drop for MutableBlock<T> {
114    #[inline]
115    fn drop(&mut self) {
116        if self.armed {
117            self.store.release_mutable(self.block_id);
118        }
119    }
120}
121
122impl<T: BlockMetadata> std::fmt::Debug for MutableBlock<T> {
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        f.debug_struct("MutableBlock")
125            .field("block_id", &self.block_id)
126            .finish()
127    }
128}
kvbm_logical/blocks/mutable.rs

kvbm_logical/blocks/
mutable.rs
