kvbm_logical/blocks/
pin.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Type-erased lifecycle pin for cross-policy block references.
5//!
6//! [`LifecyclePinRef`] is a thin wrapper around `Arc<dyn LifecyclePin>`
7//! obtained from [`ImmutableBlock::pin`](super::ImmutableBlock::pin). It
8//! serves two purposes for downstream callers (e.g., a KV offload
9//! connector that must stash a heterogeneous list of in-flight transfers):
10//!
11//! 1. **Keepalive.** As long as the pin is alive, the underlying slot
12//!    will not transition from `Active` to `Inactive` and will not be
13//!    recycled to a different sequence hash. Cloning the pin is one
14//!    `Arc::clone` — no extra heap allocation beyond the original
15//!    `ImmutableBlockInner`.
16//! 2. **Logical address.** The pin exposes
17//!    `(manager_id, block_id, sequence_hash)`, which uniquely identifies
18//!    the slot at runtime even after the originating
19//!    `ImmutableBlock<T>`'s metadata parameter `T` has been type-erased
20//!    away. Two pins with the same `manager_id` come from the same
21//!    physical pool.
22//!
23//! kvbm-logical does **not** own bytes — the pin keeps the *logical*
24//! lifecycle alive. The executor / runtime owning the physical buffer
25//! is responsible for honoring the `(manager_id, block_id)` address as
26//! the lookup key into its own per-pool storage.
27
28use std::sync::Arc;
29
30use crate::ManagerId;
31use crate::blocks::{BlockId, BlockRegistrationHandle, SequenceHash};
32
33/// Type-erased view of a registered slot's lifecycle. Implemented on the
34/// crate-private `ImmutableBlockInner<T>` so the policy parameter `T`
35/// drops out at the trait-object boundary.
36pub trait LifecyclePin: Send + Sync {
37    fn block_id(&self) -> BlockId;
38    fn sequence_hash(&self) -> SequenceHash;
39    fn manager_id(&self) -> ManagerId;
40    fn registration_handle(&self) -> BlockRegistrationHandle;
41}
42
43/// Cheap-to-clone strong reference to a registered slot's lifecycle.
44///
45/// Each `Clone` is a single `Arc::clone`. Drop releases the keepalive;
46/// when the last pin (and the last `ImmutableBlock` clone) are dropped,
47/// the underlying slot transitions to `Inactive` and may be evicted.
48#[derive(Clone)]
49pub struct LifecyclePinRef(Arc<dyn LifecyclePin>);
50
51impl LifecyclePinRef {
52    pub(crate) fn new(inner: Arc<dyn LifecyclePin>) -> Self {
53        Self(inner)
54    }
55
56    pub fn block_id(&self) -> BlockId {
57        self.0.block_id()
58    }
59
60    pub fn sequence_hash(&self) -> SequenceHash {
61        self.0.sequence_hash()
62    }
63
64    pub fn manager_id(&self) -> ManagerId {
65        self.0.manager_id()
66    }
67
68    pub fn registration_handle(&self) -> BlockRegistrationHandle {
69        self.0.registration_handle()
70    }
71
72    /// Strong-count of the underlying `Arc<dyn LifecyclePin>`. Useful for
73    /// metrics and tests; do not branch on this in production code.
74    pub fn use_count(&self) -> usize {
75        Arc::strong_count(&self.0)
76    }
77}
78
79impl std::fmt::Debug for LifecyclePinRef {
80    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
81        f.debug_struct("LifecyclePinRef")
82            .field("manager_id", &self.0.manager_id())
83            .field("block_id", &self.0.block_id())
84            .field("sequence_hash", &self.0.sequence_hash())
85            .finish()
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use crate::ManagerId;
92    use crate::testing::{TestMeta, create_iota_token_block, create_test_manager};
93
94    /// Helper: register a single block in `manager` and return both the
95    /// resulting `ImmutableBlock` and the `ManagerId` derived from its pin.
96    fn one_block_and_manager_id(
97        manager: &crate::BlockManager<TestMeta>,
98        salt: u32,
99    ) -> (crate::ImmutableBlock<TestMeta>, ManagerId) {
100        let token = create_iota_token_block(salt, 4);
101        let mb = manager
102            .allocate_blocks(1)
103            .unwrap()
104            .into_iter()
105            .next()
106            .unwrap();
107        let cb = mb.complete(&token).unwrap();
108        let block = manager
109            .register_blocks(vec![cb])
110            .into_iter()
111            .next()
112            .unwrap();
113        let id = block.pin().manager_id();
114        (block, id)
115    }
116
117    #[test]
118    fn manager_id_unique_across_managers() {
119        let mgr_a = create_test_manager::<TestMeta>(4);
120        let mgr_b = create_test_manager::<TestMeta>(4);
121        let (_a, id_a) = one_block_and_manager_id(&mgr_a, 1);
122        let (_b, id_b) = one_block_and_manager_id(&mgr_b, 2);
123        assert_ne!(id_a, id_b);
124        assert_ne!(id_a, ManagerId::NULL);
125        assert_ne!(id_b, ManagerId::NULL);
126    }
127
128    #[test]
129    fn manager_id_stable_across_clones() {
130        let manager = create_test_manager::<TestMeta>(4);
131        let (block, manager_id) = one_block_and_manager_id(&manager, 0);
132
133        let pin1 = block.pin();
134        let pin2 = block.clone().pin();
135        let pin3 = pin1.clone();
136
137        assert_eq!(pin1.manager_id(), manager_id);
138        assert_eq!(pin2.manager_id(), manager_id);
139        assert_eq!(pin3.manager_id(), manager_id);
140        assert_eq!(pin1.block_id(), block.block_id());
141        assert_eq!(pin1.sequence_hash(), block.sequence_hash());
142    }
143
144    #[test]
145    fn pin_keeps_slot_active_after_block_drops() {
146        let manager = create_test_manager::<TestMeta>(4);
147        let initial = manager.available_blocks();
148        let token = create_iota_token_block(100, 4);
149        let mb = manager
150            .allocate_blocks(1)
151            .unwrap()
152            .into_iter()
153            .next()
154            .unwrap();
155        let cb = mb.complete(&token).unwrap();
156        let block = manager
157            .register_blocks(vec![cb])
158            .into_iter()
159            .next()
160            .unwrap();
161
162        // Slot now Active: not in reset+inactive count.
163        assert_eq!(manager.available_blocks(), initial - 1);
164
165        let pin = block.pin();
166        // Dropping the block while pin is held: slot must stay Active.
167        drop(block);
168        assert_eq!(
169            manager.available_blocks(),
170            initial - 1,
171            "pin should prevent the slot from transitioning to Inactive"
172        );
173
174        // Dropping the pin releases the keepalive: slot transitions to
175        // Inactive (which counts as available).
176        drop(pin);
177        assert_eq!(manager.available_blocks(), initial);
178    }
179
180    #[test]
181    fn pin_clone_is_arc_bump() {
182        let manager = create_test_manager::<TestMeta>(4);
183        let token = create_iota_token_block(200, 4);
184        let mb = manager
185            .allocate_blocks(1)
186            .unwrap()
187            .into_iter()
188            .next()
189            .unwrap();
190        let cb = mb.complete(&token).unwrap();
191        let block = manager
192            .register_blocks(vec![cb])
193            .into_iter()
194            .next()
195            .unwrap();
196
197        let pin = block.pin();
198        let count_before = pin.use_count();
199        let _clone = pin.clone();
200        assert_eq!(pin.use_count(), count_before + 1);
201    }
202
203    #[test]
204    fn pins_from_different_managers_have_different_addresses() {
205        let mgr_a = create_test_manager::<TestMeta>(4);
206        let mgr_b = create_test_manager::<TestMeta>(4);
207
208        let token_a = create_iota_token_block(300, 4);
209        let mb_a = mgr_a
210            .allocate_blocks(1)
211            .unwrap()
212            .into_iter()
213            .next()
214            .unwrap();
215        let block_a = mgr_a
216            .register_blocks(vec![mb_a.complete(&token_a).unwrap()])
217            .into_iter()
218            .next()
219            .unwrap();
220
221        let token_b = create_iota_token_block(300, 4);
222        let mb_b = mgr_b
223            .allocate_blocks(1)
224            .unwrap()
225            .into_iter()
226            .next()
227            .unwrap();
228        let block_b = mgr_b
229            .register_blocks(vec![mb_b.complete(&token_b).unwrap()])
230            .into_iter()
231            .next()
232            .unwrap();
233
234        let pin_a = block_a.pin();
235        let pin_b = block_b.pin();
236
237        // Same content (same iota tokens) ⇒ same SequenceHash.
238        assert_eq!(pin_a.sequence_hash(), pin_b.sequence_hash());
239        // ...but different physical pools.
240        assert_ne!(pin_a.manager_id(), pin_b.manager_id());
241        // The (manager_id, block_id) pair is the disambiguating address.
242    }
243}
kvbm_logical/blocks/pin.rs

kvbm_logical/blocks/
pin.rs
