Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
OpenDAS
dynamo
Commits
94fad479
Unverified
Commit
94fad479
authored
Feb 13, 2026
by
Ryan Olson
Committed by
GitHub
Feb 14, 2026
Browse files
feat: kvbm-logical (#6033)
Signed-off-by:
Ryan Olson
<
rolson@nvidia.com
>
parent
0e55e821
Changes
54
Hide whitespace changes
Inline
Side-by-side
Showing
20 changed files
with
6941 additions
and
0 deletions
+6941
-0
lib/kvbm-logical/src/manager/tests.rs
lib/kvbm-logical/src/manager/tests.rs
+2153
-0
lib/kvbm-logical/src/metrics/collector.rs
lib/kvbm-logical/src/metrics/collector.rs
+387
-0
lib/kvbm-logical/src/metrics/mod.rs
lib/kvbm-logical/src/metrics/mod.rs
+52
-0
lib/kvbm-logical/src/metrics/pool_metrics.rs
lib/kvbm-logical/src/metrics/pool_metrics.rs
+315
-0
lib/kvbm-logical/src/metrics/stats.rs
lib/kvbm-logical/src/metrics/stats.rs
+321
-0
lib/kvbm-logical/src/pools/active.rs
lib/kvbm-logical/src/pools/active.rs
+224
-0
lib/kvbm-logical/src/pools/block_proptest.rs
lib/kvbm-logical/src/pools/block_proptest.rs
+272
-0
lib/kvbm-logical/src/pools/inactive/backends/fifo.rs
lib/kvbm-logical/src/pools/inactive/backends/fifo.rs
+358
-0
lib/kvbm-logical/src/pools/inactive/backends/hashmap_backend.rs
...bm-logical/src/pools/inactive/backends/hashmap_backend.rs
+106
-0
lib/kvbm-logical/src/pools/inactive/backends/lineage.rs
lib/kvbm-logical/src/pools/inactive/backends/lineage.rs
+751
-0
lib/kvbm-logical/src/pools/inactive/backends/lru_backend.rs
lib/kvbm-logical/src/pools/inactive/backends/lru_backend.rs
+167
-0
lib/kvbm-logical/src/pools/inactive/backends/mod.rs
lib/kvbm-logical/src/pools/inactive/backends/mod.rs
+28
-0
lib/kvbm-logical/src/pools/inactive/backends/multi_lru_backend.rs
...-logical/src/pools/inactive/backends/multi_lru_backend.rs
+359
-0
lib/kvbm-logical/src/pools/inactive/backends/reuse_policy.rs
lib/kvbm-logical/src/pools/inactive/backends/reuse_policy.rs
+46
-0
lib/kvbm-logical/src/pools/inactive/backends/tests.rs
lib/kvbm-logical/src/pools/inactive/backends/tests.rs
+243
-0
lib/kvbm-logical/src/pools/inactive/mod.rs
lib/kvbm-logical/src/pools/inactive/mod.rs
+496
-0
lib/kvbm-logical/src/pools/mod.rs
lib/kvbm-logical/src/pools/mod.rs
+77
-0
lib/kvbm-logical/src/pools/reset.rs
lib/kvbm-logical/src/pools/reset.rs
+194
-0
lib/kvbm-logical/src/pools/tests.rs
lib/kvbm-logical/src/pools/tests.rs
+363
-0
lib/kvbm-logical/src/pubsub.rs
lib/kvbm-logical/src/pubsub.rs
+29
-0
No files found.
lib/kvbm-logical/src/manager/tests.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
use
super
::
*
;
use
crate
::
KvbmSequenceHashProvider
;
use
crate
::
blocks
::
BlockError
;
use
crate
::
testing
::{
self
,
TestMeta
,
create_iota_token_block
,
create_test_token_block
as
testing_create_token_block
,
};
use
rstest
::
rstest
;
// Type alias for backward compatibility
type
TestBlockData
=
TestMeta
;
/// Helper function to create a token block with specific data (local wrapper)
fn
create_token_block
(
tokens
:
&
[
u32
])
->
dynamo_tokens
::
TokenBlock
{
testing_create_token_block
(
tokens
,
tokens
.len
()
as
u32
)
}
/// Helper function to create a token block using fill_iota pattern
fn
create_test_token_block_from_iota
(
start
:
u32
)
->
dynamo_tokens
::
TokenBlock
{
create_iota_token_block
(
start
,
4
)
}
fn
create_test_token_block_8_from_iota
(
start
:
u32
)
->
dynamo_tokens
::
TokenBlock
{
create_iota_token_block
(
start
,
8
)
}
/// Helper function to create a basic manager for testing
fn
create_test_manager
(
block_count
:
usize
)
->
BlockManager
<
TestBlockData
>
{
testing
::
create_test_manager
::
<
TestBlockData
>
(
block_count
)
}
// ============================================================================
// BUILDER PATTERN TESTS
// ============================================================================
mod
builder_tests
{
use
super
::
*
;
#[test]
fn
test_builder_default
()
{
let
registry
=
BlockRegistry
::
new
();
let
manager
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
100
)
.registry
(
registry
)
.build
()
.expect
(
"Should build with defaults"
);
// Verify initial gauge
let
snap
=
manager
.metrics
()
.snapshot
();
assert_eq!
(
snap
.reset_pool_size
,
100
);
assert_eq!
(
snap
.inactive_pool_size
,
0
);
// Verify we can allocate blocks
let
blocks
=
manager
.allocate_blocks
(
5
);
assert
!
(
blocks
.is_some
());
assert_eq!
(
blocks
.unwrap
()
.len
(),
5
);
}
#[test]
fn
test_builder_with_lru_backend
()
{
let
registry
=
BlockRegistry
::
new
();
let
manager
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
100
)
.registry
(
registry
)
.with_lru_backend
()
.build
()
.expect
(
"Should build with LRU backend"
);
// Verify we can allocate blocks
let
blocks
=
manager
.allocate_blocks
(
10
);
assert
!
(
blocks
.is_some
());
assert_eq!
(
blocks
.unwrap
()
.len
(),
10
);
}
#[test]
fn
test_builder_with_multi_lru_backend
()
{
let
registry
=
BlockRegistry
::
builder
()
.frequency_tracker
(
FrequencyTrackingCapacity
::
Small
.create_tracker
())
.build
();
let
manager
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
100
)
.registry
(
registry
)
.with_multi_lru_backend
()
.build
()
.expect
(
"Should build with MultiLRU backend"
);
// Verify we can allocate blocks
let
blocks
=
manager
.allocate_blocks
(
8
);
assert
!
(
blocks
.is_some
());
assert_eq!
(
blocks
.unwrap
()
.len
(),
8
);
}
#[test]
fn
test_builder_with_custom_multi_lru_thresholds
()
{
let
registry
=
BlockRegistry
::
builder
()
.frequency_tracker
(
FrequencyTrackingCapacity
::
Medium
.create_tracker
())
.build
();
let
manager
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
100
)
.registry
(
registry
)
.with_multi_lru_backend_custom_thresholds
(
2
,
6
,
12
)
.build
()
.expect
(
"Should build with custom thresholds"
);
// Verify we can allocate blocks
let
blocks
=
manager
.allocate_blocks
(
4
);
assert
!
(
blocks
.is_some
());
assert_eq!
(
blocks
.unwrap
()
.len
(),
4
);
}
#[test]
fn
test_builder_with_duplication_policy
()
{
let
registry
=
BlockRegistry
::
new
();
let
manager
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
50
)
.registry
(
registry
)
.duplication_policy
(
BlockDuplicationPolicy
::
Reject
)
.with_lru_backend
()
.build
()
.expect
(
"Should build with duplication policy"
);
let
blocks
=
manager
.allocate_blocks
(
2
);
assert
!
(
blocks
.is_some
());
assert_eq!
(
blocks
.unwrap
()
.len
(),
2
);
}
#[test]
fn
test_builder_validation_zero_blocks
()
{
let
registry
=
BlockRegistry
::
new
();
let
result
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
0
)
.registry
(
registry
)
.build
();
assert
!
(
result
.is_err
());
if
let
Err
(
err
)
=
result
{
assert
!
(
err
.to_string
()
.contains
(
"block_count must be greater than 0"
)
);
}
}
#[test]
fn
test_builder_validation_missing_block_count
()
{
let
registry
=
BlockRegistry
::
new
();
let
result
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.registry
(
registry
)
.with_lru_backend
()
.build
();
assert
!
(
result
.is_err
());
if
let
Err
(
err
)
=
result
{
assert
!
(
err
.to_string
()
.contains
(
"block_count is required"
));
}
}
#[test]
fn
test_builder_validation_missing_registry
()
{
let
result
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
100
)
.with_lru_backend
()
.build
();
assert
!
(
result
.is_err
());
if
let
Err
(
err
)
=
result
{
assert
!
(
err
.to_string
()
.contains
(
"registry is required"
));
}
}
#[test]
#[should_panic(expected
=
"must be <= 15"
)]
fn
test_builder_invalid_threshold_too_high
()
{
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
100
)
.with_multi_lru_backend_custom_thresholds
(
2
,
6
,
20
);
// 20 > 15, should panic
}
#[test]
#[should_panic(expected
=
"must be in ascending order"
)]
fn
test_builder_invalid_threshold_order
()
{
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
100
)
.with_multi_lru_backend_custom_thresholds
(
6
,
2
,
10
);
// Not ascending, should panic
}
#[test]
fn
test_builder_multi_lru_requires_frequency_tracking
()
{
let
registry
=
BlockRegistry
::
new
();
// No frequency tracking
let
result
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
100
)
.registry
(
registry
)
.with_multi_lru_backend
()
.build
();
assert
!
(
result
.is_err
());
if
let
Err
(
err
)
=
result
{
assert
!
(
err
.to_string
()
.contains
(
"frequency tracking"
));
}
}
}
// ============================================================================
// BLOCK ALLOCATION TESTS
// ============================================================================
mod
allocation_tests
{
use
super
::
*
;
#[test]
fn
test_allocate_single_block
()
{
let
manager
=
create_test_manager
(
10
);
let
m
=
manager
.metrics
();
let
initial_available
=
manager
.available_blocks
();
let
initial_total
=
manager
.total_blocks
();
assert_eq!
(
initial_available
,
10
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.reset_pool_size
,
10
);
let
blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate 1 block"
);
assert_eq!
(
blocks
.len
(),
1
);
// Verify available blocks decreased
assert_eq!
(
manager
.available_blocks
(),
initial_available
-
1
);
assert_eq!
(
manager
.total_blocks
(),
initial_total
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.allocations
,
1
);
assert_eq!
(
snap
.inflight_mutable
,
1
);
assert_eq!
(
snap
.reset_pool_size
,
9
);
let
block
=
blocks
.into_iter
()
.next
()
.unwrap
();
// Verify block has a valid ID
let
_
block_id
=
block
.block_id
();
// Drop the block and verify it returns to pool
drop
(
block
);
assert_eq!
(
manager
.available_blocks
(),
initial_available
);
assert_eq!
(
manager
.total_blocks
(),
initial_total
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inflight_mutable
,
0
);
assert_eq!
(
snap
.reset_pool_size
,
10
);
}
#[test]
fn
test_allocate_multiple_blocks
()
{
let
manager
=
create_test_manager
(
20
);
let
m
=
manager
.metrics
();
let
initial_available
=
manager
.available_blocks
();
let
initial_total
=
manager
.total_blocks
();
assert_eq!
(
initial_available
,
20
);
let
blocks
=
manager
.allocate_blocks
(
5
)
.expect
(
"Should allocate 5 blocks"
);
assert_eq!
(
blocks
.len
(),
5
);
// Verify available blocks decreased correctly
assert_eq!
(
manager
.available_blocks
(),
initial_available
-
5
);
assert_eq!
(
manager
.total_blocks
(),
initial_total
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.allocations
,
5
);
assert_eq!
(
snap
.inflight_mutable
,
5
);
// Verify all blocks have unique IDs
let
mut
block_ids
=
Vec
::
new
();
for
block
in
blocks
{
let
id
=
block
.block_id
();
assert
!
(
!
block_ids
.contains
(
&
id
),
"Block IDs should be unique"
);
block_ids
.push
(
id
);
}
// All blocks should return to pool automatically on drop
assert_eq!
(
manager
.available_blocks
(),
initial_available
);
assert_eq!
(
manager
.total_blocks
(),
initial_total
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inflight_mutable
,
0
);
}
#[test]
fn
test_allocate_all_blocks
()
{
let
manager
=
create_test_manager
(
10
);
let
blocks
=
manager
.allocate_blocks
(
10
)
.expect
(
"Should allocate all blocks"
);
assert_eq!
(
blocks
.len
(),
10
);
}
#[test]
fn
test_allocate_more_than_available
()
{
let
manager
=
create_test_manager
(
5
);
let
result
=
manager
.allocate_blocks
(
10
);
assert
!
(
result
.is_none
(),
"Should not allocate more blocks than available"
);
}
#[test]
fn
test_allocate_zero_blocks
()
{
let
manager
=
create_test_manager
(
10
);
let
blocks
=
manager
.allocate_blocks
(
0
)
.expect
(
"Should allocate 0 blocks"
);
assert_eq!
(
blocks
.len
(),
0
);
}
#[test]
fn
test_sequential_allocations
()
{
let
manager
=
create_test_manager
(
10
);
let
m
=
manager
.metrics
();
let
total_blocks
=
manager
.total_blocks
();
assert_eq!
(
manager
.available_blocks
(),
total_blocks
);
assert_eq!
(
m
.snapshot
()
.reset_pool_size
,
10
);
let
blocks1
=
manager
.allocate_blocks
(
3
)
.expect
(
"First allocation"
);
assert_eq!
(
blocks1
.len
(),
3
);
assert_eq!
(
manager
.available_blocks
(),
total_blocks
-
3
);
assert_eq!
(
m
.snapshot
()
.reset_pool_size
,
7
);
let
blocks2
=
manager
.allocate_blocks
(
4
)
.expect
(
"Second allocation"
);
assert_eq!
(
blocks2
.len
(),
4
);
assert_eq!
(
manager
.available_blocks
(),
total_blocks
-
7
);
assert_eq!
(
m
.snapshot
()
.reset_pool_size
,
3
);
let
blocks3
=
manager
.allocate_blocks
(
3
)
.expect
(
"Third allocation"
);
assert_eq!
(
blocks3
.len
(),
3
);
assert_eq!
(
manager
.available_blocks
(),
0
);
assert_eq!
(
m
.snapshot
()
.reset_pool_size
,
0
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.allocations
,
10
);
assert_eq!
(
snap
.inflight_mutable
,
10
);
// Should have no blocks left
let
blocks4
=
manager
.allocate_blocks
(
1
);
assert
!
(
blocks4
.is_none
(),
"Should not have any blocks left"
);
// Drop blocks in reverse order and verify counts
drop
(
blocks3
);
assert_eq!
(
manager
.available_blocks
(),
3
);
assert_eq!
(
m
.snapshot
()
.reset_pool_size
,
3
);
drop
(
blocks2
);
assert_eq!
(
manager
.available_blocks
(),
7
);
assert_eq!
(
m
.snapshot
()
.reset_pool_size
,
7
);
drop
(
blocks1
);
assert_eq!
(
manager
.available_blocks
(),
total_blocks
);
assert_eq!
(
manager
.total_blocks
(),
total_blocks
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inflight_mutable
,
0
);
assert_eq!
(
snap
.reset_pool_size
,
10
);
}
}
// ============================================================================
// BLOCK LIFECYCLE AND POOL RETURN TESTS
// ============================================================================
mod
lifecycle_tests
{
use
super
::
*
;
#[test]
fn
test_mutable_block_returns_to_reset_pool
()
{
let
manager
=
create_test_manager
(
10
);
let
m
=
manager
.metrics
();
let
initial_available
=
manager
.available_blocks
();
let
initial_total
=
manager
.total_blocks
();
assert_eq!
(
initial_available
,
10
);
assert_eq!
(
initial_total
,
10
);
{
let
blocks
=
manager
.allocate_blocks
(
3
)
.expect
(
"Should allocate 3 blocks"
);
assert_eq!
(
blocks
.len
(),
3
);
// Available blocks should decrease
assert_eq!
(
manager
.available_blocks
(),
initial_available
-
3
);
assert_eq!
(
manager
.total_blocks
(),
initial_total
);
// Total never changes
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inflight_mutable
,
3
);
assert_eq!
(
snap
.reset_pool_size
,
7
);
}
// MutableBlocks dropped here - should return to reset pool
// Available blocks should return to original count
assert_eq!
(
manager
.available_blocks
(),
initial_available
);
assert_eq!
(
manager
.total_blocks
(),
initial_total
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inflight_mutable
,
0
);
assert_eq!
(
snap
.reset_pool_size
,
10
);
}
#[test]
fn
test_complete_block_returns_to_reset_pool
()
{
let
manager
=
create_test_manager
(
10
);
let
m
=
manager
.metrics
();
let
initial_available
=
manager
.available_blocks
();
let
initial_total
=
manager
.total_blocks
();
{
let
mutable_blocks
=
manager
.allocate_blocks
(
2
)
.expect
(
"Should allocate blocks"
);
assert_eq!
(
manager
.available_blocks
(),
initial_available
-
2
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.reset_pool_size
,
8
);
// Note: create_token_block uses 3 tokens but block_size is 4,
// so complete() returns Err(BlockSizeMismatch) for all blocks.
let
_
complete_blocks
:
Vec
<
_
>
=
mutable_blocks
.into_iter
()
.enumerate
()
.map
(|(
i
,
block
)|
{
let
tokens
=
vec!
[
400
+
i
as
u32
,
401
+
i
as
u32
,
402
+
i
as
u32
];
let
token_block
=
create_token_block
(
&
tokens
);
block
.complete
(
&
token_block
)
})
.collect
();
// Blocks are still unavailable while in Complete state
assert_eq!
(
manager
.available_blocks
(),
initial_available
-
2
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inflight_mutable
,
2
);
assert_eq!
(
snap
.stagings
,
0
);
assert_eq!
(
snap
.reset_pool_size
,
8
);
}
// CompleteBlocks dropped here - should return to reset pool
// Available blocks should return to original count since blocks weren't registered
assert_eq!
(
manager
.available_blocks
(),
initial_available
);
assert_eq!
(
manager
.total_blocks
(),
initial_total
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inflight_mutable
,
0
);
assert_eq!
(
snap
.reset_pool_size
,
10
);
}
#[test]
fn
test_registered_block_lifecycle
()
{
let
manager
=
create_test_manager
(
10
);
let
m
=
manager
.metrics
();
let
initial_available
=
manager
.available_blocks
();
let
initial_total
=
manager
.total_blocks
();
// Step 1: Allocate and complete blocks
let
token_block
=
create_test_token_block_from_iota
(
500
);
let
seq_hash
=
token_block
.kvbm_sequence_hash
();
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
assert_eq!
(
manager
.available_blocks
(),
initial_available
-
1
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.allocations
,
1
);
assert_eq!
(
snap
.inflight_mutable
,
1
);
assert_eq!
(
snap
.reset_pool_size
,
9
);
assert_eq!
(
snap
.inactive_pool_size
,
0
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
// Still unavailable while in Complete state
assert_eq!
(
manager
.available_blocks
(),
initial_available
-
1
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.stagings
,
1
);
assert_eq!
(
snap
.inflight_mutable
,
0
);
// Step 2: Register the block
let
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
assert_eq!
(
immutable_blocks
.len
(),
1
);
let
immutable_block
=
immutable_blocks
.into_iter
()
.next
()
.unwrap
();
// Block is still not available (it's now in active/inactive pools, not reset)
assert_eq!
(
manager
.available_blocks
(),
initial_available
-
1
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.registrations
,
1
);
assert_eq!
(
snap
.inflight_immutable
,
1
);
assert_eq!
(
snap
.reset_pool_size
,
9
);
assert_eq!
(
snap
.inactive_pool_size
,
0
);
{
// Step 3: Use the block and verify it can be matched
let
matched_blocks
=
manager
.match_blocks
(
&
[
seq_hash
]);
assert_eq!
(
matched_blocks
.len
(),
1
);
assert_eq!
(
matched_blocks
[
0
]
.sequence_hash
(),
seq_hash
);
// Still not available while being used
assert_eq!
(
manager
.available_blocks
(),
initial_available
-
1
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.match_hashes_requested
,
1
);
assert_eq!
(
snap
.match_blocks_returned
,
1
);
assert_eq!
(
snap
.inflight_immutable
,
2
);
}
// matched blocks dropped here
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inflight_immutable
,
1
);
// Step 4: Drop the original registered block → block moves to inactive
drop
(
immutable_block
);
// Block should now be available again (moved to inactive pool when ref count reached 0)
assert_eq!
(
manager
.available_blocks
(),
initial_available
);
assert_eq!
(
manager
.total_blocks
(),
initial_total
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inflight_immutable
,
0
);
assert_eq!
(
snap
.reset_pool_size
,
9
);
assert_eq!
(
snap
.inactive_pool_size
,
1
);
// Step 5: Re-match from inactive pool → pulls block out
{
let
re_matched
=
manager
.match_blocks
(
&
[
seq_hash
]);
assert_eq!
(
re_matched
.len
(),
1
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inactive_pool_size
,
0
);
}
// re_matched dropped → block returns to inactive
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inactive_pool_size
,
1
);
}
#[test]
fn
test_concurrent_allocation_and_return
()
{
use
std
::
sync
::
Arc
;
use
std
::
thread
;
let
manager
=
Arc
::
new
(
create_test_manager
(
20
));
let
initial_total
=
manager
.total_blocks
();
let
handles
:
Vec
<
_
>
=
(
0
..
5
)
.map
(|
i
|
{
let
manager_clone
=
Arc
::
clone
(
&
manager
);
thread
::
spawn
(
move
||
{
// Each thread allocates and drops some blocks
for
j
in
0
..
3
{
let
blocks
=
manager_clone
.allocate_blocks
(
2
);
if
let
Some
(
blocks
)
=
blocks
{
// Complete one block
let
token_block
=
create_test_token_block_from_iota
((
600
+
i
*
10
+
j
)
as
u32
);
let
complete_block
=
blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
// Register and drop
let
_
immutable_blocks
=
manager_clone
.register_blocks
(
vec!
[
complete_block
]);
// blocks automatically dropped at end of scope
}
}
})
})
.collect
();
// Wait for all threads to complete
for
handle
in
handles
{
handle
.join
()
.unwrap
();
}
// All blocks should eventually be available again
assert_eq!
(
manager
.total_blocks
(),
initial_total
);
// Available might be less than total if some blocks are in inactive pool,
// but total should be preserved
}
#[test]
fn
test_full_block_lifecycle
()
{
let
manager
=
create_test_manager
(
10
);
let
total_blocks
=
manager
.total_blocks
();
assert_eq!
(
manager
.available_blocks
(),
total_blocks
);
// Step 1: Allocate 5 blocks
let
mutable_blocks
=
manager
.allocate_blocks
(
5
)
.expect
(
"Should allocate 5 blocks"
);
assert_eq!
(
manager
.available_blocks
(),
total_blocks
-
5
);
assert_eq!
(
manager
.total_blocks
(),
total_blocks
);
// Step 2: Complete 3 blocks, drop 2 mutable blocks
let
mut
mutable_blocks_iter
=
mutable_blocks
.into_iter
();
let
complete_blocks
:
Vec
<
_
>
=
(
0
..
3
)
.map
(|
i
|
{
let
block
=
mutable_blocks_iter
.next
()
.unwrap
();
let
tokens
=
vec!
[
700
+
i
as
u32
,
701
+
i
as
u32
,
702
+
i
as
u32
,
703
+
i
as
u32
,
];
let
token_block
=
create_token_block
(
&
tokens
);
block
.complete
(
&
token_block
)
.expect
(
"Should complete block"
)
})
.collect
();
let
mutable_part
:
Vec
<
_
>
=
mutable_blocks_iter
.collect
();
drop
(
mutable_part
);
// Drop 2 mutable blocks
// Should have 2 blocks returned to reset pool
assert_eq!
(
manager
.available_blocks
(),
total_blocks
-
3
);
// Step 3: Register the 3 completed blocks
let
immutable_blocks
=
manager
.register_blocks
(
complete_blocks
);
assert_eq!
(
immutable_blocks
.len
(),
3
);
// Still 3 blocks unavailable (now in active pool)
assert_eq!
(
manager
.available_blocks
(),
total_blocks
-
3
);
// Step 4: Match and use one of the blocks
let
seq_hash
=
create_test_token_block_from_iota
(
700
)
.kvbm_sequence_hash
();
let
matched_blocks
=
manager
.match_blocks
(
&
[
seq_hash
]);
assert_eq!
(
matched_blocks
.len
(),
1
);
// Step 5: Drop one registered block, keep others
drop
(
immutable_blocks
.into_iter
()
.next
());
// Still have registered blocks in use, so available count depends on ref counting
let
available_after_drop
=
manager
.available_blocks
();
assert
!
(
available_after_drop
>=
total_blocks
-
3
);
assert
!
(
available_after_drop
<=
total_blocks
);
// Step 6: Drop everything
drop
(
matched_blocks
);
// Eventually all blocks should be available again
// (Some might be in inactive pool, but available_blocks counts both reset and inactive)
assert_eq!
(
manager
.total_blocks
(),
total_blocks
);
let
final_available
=
manager
.available_blocks
();
assert_eq!
(
final_available
,
total_blocks
);
// Allow for some blocks in inactive pool
}
}
// ============================================================================
// BLOCK SIZE VALIDATION TESTS
// ============================================================================
mod
block_size_tests
{
use
super
::
*
;
#[test]
fn
test_default_block_size
()
{
let
manager
=
create_test_manager
(
10
);
assert_eq!
(
manager
.block_size
(),
4
);
// create_test_manager uses block_size(4)
}
#[test]
fn
test_custom_block_size
()
{
let
registry
=
BlockRegistry
::
new
();
let
manager
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
10
)
.block_size
(
32
)
.registry
(
registry
)
.build
()
.expect
(
"Should build with custom block size"
);
assert_eq!
(
manager
.block_size
(),
32
);
}
#[test]
fn
test_block_size_validation_correct_size
()
{
let
manager
=
create_test_manager
(
10
);
let
token_block
=
create_test_token_block_from_iota
(
100
);
// 4 tokens
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
mutable_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
();
// Should succeed since token_block has exactly 4 tokens
let
result
=
mutable_block
.complete
(
&
token_block
);
assert
!
(
result
.is_ok
());
}
#[test]
fn
test_block_size_validation_wrong_size
()
{
// Create a manager expecting 8-token blocks
let
registry
=
BlockRegistry
::
new
();
let
manager
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
10
)
.block_size
(
8
)
.registry
(
registry
)
.with_lru_backend
()
.build
()
.expect
(
"Should build manager"
);
let
token_block
=
create_test_token_block_from_iota
(
1
);
// 4 tokens, expected 8
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
mutable_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
();
// Should fail since token_block has 4 tokens but manager expects 8
let
result
=
mutable_block
.complete
(
&
token_block
);
assert
!
(
result
.is_err
());
if
let
Err
(
BlockError
::
BlockSizeMismatch
{
expected
,
actual
,
block
:
_
,
})
=
result
{
assert_eq!
(
expected
,
8
);
assert_eq!
(
actual
,
4
);
}
else
{
panic!
(
"Expected BlockSizeMismatch error"
);
}
}
#[rstest]
#[case(
1
)]
#[case(
2
)]
#[case(
4
)]
#[case(
8
)]
#[case(
16
)]
#[case(
32
)]
#[case(
64
)]
#[case(
128
)]
#[case(
256
)]
#[case(
512
)]
#[case(
1024
)]
fn
test_builder_block_size_power_of_two
(
#[case]
size
:
usize
)
{
let
registry
=
BlockRegistry
::
new
();
let
result
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
10
)
.block_size
(
size
)
.registry
(
registry
)
.build
();
assert
!
(
result
.is_ok
(),
"Block size {} should be valid"
,
size
);
}
#[test]
#[should_panic(expected
=
"block_size must be a power of 2"
)]
fn
test_builder_block_size_not_power_of_two
()
{
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
10
)
.block_size
(
15
);
// Not a power of 2
}
#[test]
#[should_panic(expected
=
"block_size must be between 1 and 1024"
)]
fn
test_builder_block_size_too_large
()
{
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
10
)
.block_size
(
2048
);
// Too large
}
#[test]
#[should_panic(expected
=
"block_size must be between 1 and 1024"
)]
fn
test_builder_block_size_zero
()
{
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
10
)
.block_size
(
0
);
// Zero is invalid
}
#[test]
#[should_panic(expected
=
"block_size must be a power of 2"
)]
fn
test_builder_validation_invalid_block_size
()
{
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
10
)
.block_size
(
7
);
// Not a power of 2, panics immediately
}
#[test]
fn
test_different_block_sizes
()
{
// Test with block size 4
let
registry_4
=
BlockRegistry
::
new
();
let
manager_4
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
10
)
.block_size
(
4
)
.registry
(
registry_4
)
.build
()
.expect
(
"Should build with block size 4"
);
let
token_block_4
=
create_test_token_block_from_iota
(
10
);
// 4 tokens
let
mutable_blocks
=
manager_4
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
result
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block_4
);
assert
!
(
result
.is_ok
());
// Test with block size 8
let
registry_8
=
BlockRegistry
::
new
();
let
manager_8
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
10
)
.block_size
(
8
)
.registry
(
registry_8
)
.build
()
.expect
(
"Should build with block size 8"
);
let
token_block_8
=
create_test_token_block_8_from_iota
(
20
);
// 8 tokens
let
mutable_blocks
=
manager_8
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
result
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block_8
);
assert
!
(
result
.is_ok
());
}
}
// ============================================================================
// BLOCK REGISTRATION AND DEDUPLICATION TESTS
// ============================================================================
mod
registration_tests
{
use
super
::
*
;
#[test]
fn
test_register_single_block
()
{
let
manager
=
create_test_manager
(
10
);
let
m
=
manager
.metrics
();
let
token_block
=
create_test_token_block_from_iota
(
150
);
let
expected_hash
=
token_block
.kvbm_sequence_hash
();
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
let
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
assert_eq!
(
immutable_blocks
.len
(),
1
);
let
immutable_block
=
immutable_blocks
.into_iter
()
.next
()
.unwrap
();
assert_eq!
(
immutable_block
.sequence_hash
(),
expected_hash
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.registrations
,
1
);
assert_eq!
(
snap
.stagings
,
1
);
}
#[test]
fn
test_register_multiple_blocks
()
{
let
manager
=
create_test_manager
(
10
);
let
m
=
manager
.metrics
();
let
mut
complete_blocks
=
Vec
::
new
();
let
mut
expected_hashes
=
Vec
::
new
();
for
i
in
0
..
3
{
let
tokens
=
vec!
[
100
+
i
,
101
+
i
,
102
+
i
,
103
+
i
];
let
token_block
=
create_token_block
(
&
tokens
);
expected_hashes
.push
(
token_block
.kvbm_sequence_hash
());
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
complete_blocks
.push
(
complete_block
);
}
let
immutable_blocks
=
manager
.register_blocks
(
complete_blocks
);
assert_eq!
(
immutable_blocks
.len
(),
3
);
for
(
i
,
immutable_block
)
in
immutable_blocks
.iter
()
.enumerate
()
{
assert_eq!
(
immutable_block
.sequence_hash
(),
expected_hashes
[
i
]);
}
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.registrations
,
3
);
assert_eq!
(
snap
.stagings
,
3
);
}
#[rstest]
#[case(BlockDuplicationPolicy::Allow,
200
,
"allow"
,
false
)]
#[case(BlockDuplicationPolicy::Reject,
300
,
"reject"
,
true
)]
fn
test_deduplication_policy
(
#[case]
policy
:
BlockDuplicationPolicy
,
#[case]
iota_base
:
u32
,
#[case]
policy_name
:
&
str
,
#[case]
expect_same_block_id
:
bool
,
)
{
let
registry
=
BlockRegistry
::
new
();
let
manager
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
10
)
.block_size
(
4
)
.registry
(
registry
)
.duplication_policy
(
policy
)
.with_lru_backend
()
.build
()
.expect
(
"Should build manager"
);
let
token_block
=
create_test_token_block_from_iota
(
iota_base
);
let
seq_hash
=
token_block
.kvbm_sequence_hash
();
// Register the same sequence hash twice
let
complete_block1
=
{
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
)
};
let
complete_block2
=
{
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
)
};
let
immutable_blocks1
=
manager
.register_blocks
(
vec!
[
complete_block1
]);
let
immutable_blocks2
=
manager
.register_blocks
(
vec!
[
complete_block2
]);
assert_eq!
(
immutable_blocks1
.len
(),
1
);
assert_eq!
(
immutable_blocks2
.len
(),
1
);
// Both should have the same sequence hash
assert_eq!
(
immutable_blocks1
[
0
]
.sequence_hash
(),
seq_hash
);
assert_eq!
(
immutable_blocks2
[
0
]
.sequence_hash
(),
seq_hash
);
// Check block IDs based on policy
if
expect_same_block_id
{
// Duplicates are rejected - same block ID
assert_eq!
(
immutable_blocks1
[
0
]
.block_id
(),
immutable_blocks2
[
0
]
.block_id
(),
"With {} policy, duplicates should reuse the same block ID"
,
policy_name
);
let
snap
=
manager
.metrics
()
.snapshot
();
assert_eq!
(
snap
.registration_dedup
,
1
);
}
else
{
// Duplicates are allowed - different block IDs
assert_ne!
(
immutable_blocks1
[
0
]
.block_id
(),
immutable_blocks2
[
0
]
.block_id
(),
"With {} policy, duplicates should have different block IDs"
,
policy_name
);
let
snap
=
manager
.metrics
()
.snapshot
();
assert_eq!
(
snap
.duplicate_blocks
,
1
);
}
}
#[test]
fn
test_register_mutable_block_from_existing_reject_returns_block_to_reset_pool
()
{
let
registry
=
BlockRegistry
::
new
();
let
manager
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
2
)
.block_size
(
4
)
.registry
(
registry
)
.duplication_policy
(
BlockDuplicationPolicy
::
Reject
)
.build
()
.expect
(
"Should build manager"
);
let
blocks
=
manager
.allocate_blocks
(
2
)
.expect
(
"Should allocate two blocks"
);
let
mut
iter
=
blocks
.into_iter
();
let
primary_mutable
=
iter
.next
()
.expect
(
"Should have first block"
);
let
duplicate_mutable
=
iter
.next
()
.expect
(
"Should have second block"
);
let
primary_id
=
primary_mutable
.block_id
();
let
duplicate_id
=
duplicate_mutable
.block_id
();
let
token_block
=
create_test_token_block_from_iota
(
42
);
let
primary_complete
=
primary_mutable
.complete
(
&
token_block
)
.expect
(
"Should complete primary block"
);
let
mut
registered
=
manager
.register_blocks
(
vec!
[
primary_complete
]);
let
primary_immutable
=
registered
.pop
()
.expect
(
"Should register primary block"
);
let
duplicate_completed
=
duplicate_mutable
.stage
(
primary_immutable
.sequence_hash
(),
manager
.block_size
())
.expect
(
"block size should match"
);
let
result
=
manager
.register_block
(
duplicate_completed
);
assert_eq!
(
result
.block_id
(),
primary_id
,
"Should reuse existing primary when duplicates are rejected"
);
assert_eq!
(
manager
.available_blocks
(),
1
,
"Rejected duplicate should be returned to the reset pool"
);
let
mut
returned_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate returned reset block"
);
let
returned_block
=
returned_blocks
.pop
()
.expect
(
"Should contain one returned block"
);
assert_eq!
(
returned_block
.block_id
(),
duplicate_id
,
"Returned block should be the rejected duplicate"
);
let
snap
=
manager
.metrics
()
.snapshot
();
assert_eq!
(
snap
.registrations
,
2
);
assert_eq!
(
snap
.registration_dedup
,
1
);
// returned_block is still held, so reset pool is empty
assert_eq!
(
snap
.reset_pool_size
,
0
);
// Drop returned_block → back to reset pool
drop
(
returned_block
);
assert_eq!
(
manager
.metrics
()
.snapshot
()
.reset_pool_size
,
1
);
}
}
// ============================================================================
// BLOCK MATCHING TESTS
// ============================================================================
mod
matching_tests
{
use
super
::
*
;
#[test]
fn
test_match_no_blocks
()
{
let
manager
=
create_test_manager
(
10
);
let
seq_hashes
=
vec!
[
create_test_token_block_from_iota
(
400
)
.kvbm_sequence_hash
()];
let
matched_blocks
=
manager
.match_blocks
(
&
seq_hashes
);
assert_eq!
(
matched_blocks
.len
(),
0
);
}
#[test]
fn
test_match_single_block
()
{
let
manager
=
create_test_manager
(
10
);
let
m
=
manager
.metrics
();
let
token_block
=
create_test_token_block_from_iota
(
500
);
let
seq_hash
=
token_block
.kvbm_sequence_hash
();
// Register a block
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
let
_
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
// Try to match it
let
matched_blocks
=
manager
.match_blocks
(
&
[
seq_hash
]);
assert_eq!
(
matched_blocks
.len
(),
1
);
assert_eq!
(
matched_blocks
[
0
]
.sequence_hash
(),
seq_hash
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.match_hashes_requested
,
1
);
assert_eq!
(
snap
.match_blocks_returned
,
1
);
}
#[test]
fn
test_match_multiple_blocks
()
{
let
manager
=
create_test_manager
(
10
);
let
mut
seq_hashes
=
Vec
::
new
();
// Register multiple blocks
for
i
in
0
..
4
{
let
tokens
=
vec!
[
600
+
i
,
601
+
i
,
602
+
i
,
603
+
i
];
let
token_block
=
create_token_block
(
&
tokens
);
seq_hashes
.push
(
token_block
.kvbm_sequence_hash
());
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
let
_
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
}
// Match all blocks
let
matched_blocks
=
manager
.match_blocks
(
&
seq_hashes
);
assert_eq!
(
matched_blocks
.len
(),
4
);
for
(
i
,
matched_block
)
in
matched_blocks
.iter
()
.enumerate
()
{
assert_eq!
(
matched_block
.sequence_hash
(),
seq_hashes
[
i
]);
}
let
snap
=
manager
.metrics
()
.snapshot
();
assert_eq!
(
snap
.match_hashes_requested
,
4
);
assert_eq!
(
snap
.match_blocks_returned
,
4
);
}
#[test]
fn
test_match_partial_blocks
()
{
let
manager
=
create_test_manager
(
10
);
let
mut
seq_hashes
=
Vec
::
new
();
// Register only some blocks
for
i
in
0
..
3
{
let
tokens
=
vec!
[
700
+
i
,
701
+
i
,
702
+
i
,
703
+
i
];
let
token_block
=
create_token_block
(
&
tokens
);
seq_hashes
.push
(
token_block
.kvbm_sequence_hash
());
if
i
<
2
{
// Only register first 2 blocks
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
let
_
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
}
}
// Try to match all 3 - should only get 2
let
matched_blocks
=
manager
.match_blocks
(
&
seq_hashes
);
assert_eq!
(
matched_blocks
.len
(),
2
);
for
matched_block
in
matched_blocks
{
assert
!
(
seq_hashes
[
0
..
2
]
.contains
(
&
matched_block
.sequence_hash
()));
}
let
snap
=
manager
.metrics
()
.snapshot
();
assert_eq!
(
snap
.match_hashes_requested
,
3
);
assert_eq!
(
snap
.match_blocks_returned
,
2
);
}
#[test]
fn
test_match_blocks_returns_immutable_blocks
()
{
let
manager
=
create_test_manager
(
10
);
let
token_block
=
create_test_token_block_from_iota
(
800
);
let
seq_hash
=
token_block
.kvbm_sequence_hash
();
// Register a block
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
let
_
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
// Match and verify it's an ImmutableBlock
let
matched_blocks
=
manager
.match_blocks
(
&
[
seq_hash
]);
assert_eq!
(
matched_blocks
.len
(),
1
);
let
immutable_block
=
&
matched_blocks
[
0
];
assert_eq!
(
immutable_block
.sequence_hash
(),
seq_hash
);
// Test that we can downgrade it
let
weak_block
=
immutable_block
.downgrade
();
assert_eq!
(
weak_block
.sequence_hash
(),
seq_hash
);
}
}
// ============================================================================
// IMMUTABLE BLOCK AND WEAK BLOCK TESTS
// ============================================================================
mod
immutable_block_tests
{
use
super
::
*
;
#[test]
fn
test_immutable_block_downgrade_upgrade
()
{
let
manager
=
create_test_manager
(
10
);
let
token_block
=
create_test_token_block_from_iota
(
100
);
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
let
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
let
immutable_block
=
immutable_blocks
.into_iter
()
.next
()
.unwrap
();
// Test downgrade to WeakBlock
let
weak_block
=
immutable_block
.downgrade
();
assert_eq!
(
weak_block
.sequence_hash
(),
immutable_block
.sequence_hash
());
// Test upgrade from WeakBlock
let
upgraded_block
=
weak_block
.upgrade
()
.expect
(
"Should be able to upgrade"
);
assert_eq!
(
upgraded_block
.sequence_hash
(),
immutable_block
.sequence_hash
()
);
assert_eq!
(
upgraded_block
.block_id
(),
immutable_block
.block_id
());
}
#[test]
fn
test_weak_block_upgrade_after_drop
()
{
let
manager
=
create_test_manager
(
10
);
let
token_block
=
create_test_token_block_from_iota
(
200
);
let
seq_hash
=
token_block
.kvbm_sequence_hash
();
// Create a weak block
let
weak_block
=
{
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
let
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
let
immutable_block
=
immutable_blocks
.into_iter
()
.next
()
.unwrap
();
// Downgrade to weak
immutable_block
.downgrade
()
};
// immutable_block is dropped here
// The upgrade function should still find the block through the pools
let
upgraded_block
=
weak_block
.upgrade
();
// The result depends on whether the block is still in the pools
if
let
Some
(
block
)
=
upgraded_block
{
assert_eq!
(
block
.sequence_hash
(),
seq_hash
);
}
}
#[test]
fn
test_weak_block_upgrade_nonexistent
()
{
let
manager
=
create_test_manager
(
10
);
let
token_block
=
create_token_block
(
&
[
999
,
998
,
997
,
996
]);
// Keep non-sequential for this test
// Create an ImmutableBlock and immediately downgrade it
let
weak_block
=
{
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
let
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
let
immutable_block
=
immutable_blocks
.into_iter
()
.next
()
.unwrap
();
immutable_block
.downgrade
()
};
// Force eviction by filling up the pool with other blocks
for
i
in
0
..
10
{
let
tokens
=
vec!
[
1000
+
i
,
1001
+
i
,
1002
+
i
,
1003
+
i
];
let
token_block
=
create_token_block
(
&
tokens
);
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
let
_
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
}
// Try to upgrade - might fail if the original block was evicted
let
upgraded_block
=
weak_block
.upgrade
();
assert
!
(
upgraded_block
.is_none
());
// // This test just verifies that upgrade doesn't panic, result can be None
// if let Some(block) = upgraded_block {
// assert_eq!(
// block.sequence_hash(),
// create_token_block(&[999, 998, 997, 996]).sequence_hash()
// );
// }
}
#[test]
fn
test_multiple_weak_blocks_same_sequence
()
{
let
manager
=
create_test_manager
(
10
);
let
token_block
=
create_test_token_block_from_iota
(
150
);
let
seq_hash
=
token_block
.kvbm_sequence_hash
();
// Create multiple weak blocks from the same immutable block
let
(
weak1
,
weak2
,
weak3
)
=
{
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
let
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
let
immutable_block
=
immutable_blocks
.into_iter
()
.next
()
.unwrap
();
let
w1
=
immutable_block
.downgrade
();
let
w2
=
immutable_block
.downgrade
();
let
w3
=
immutable_block
.downgrade
();
(
w1
,
w2
,
w3
)
};
// All weak blocks should have the same sequence hash
assert_eq!
(
weak1
.sequence_hash
(),
seq_hash
);
assert_eq!
(
weak2
.sequence_hash
(),
seq_hash
);
assert_eq!
(
weak3
.sequence_hash
(),
seq_hash
);
// All should be able to upgrade
let
upgraded1
=
weak1
.upgrade
()
.expect
(
"Should upgrade"
);
let
upgraded2
=
weak2
.upgrade
()
.expect
(
"Should upgrade"
);
let
upgraded3
=
weak3
.upgrade
()
.expect
(
"Should upgrade"
);
assert_eq!
(
upgraded1
.sequence_hash
(),
seq_hash
);
assert_eq!
(
upgraded2
.sequence_hash
(),
seq_hash
);
assert_eq!
(
upgraded3
.sequence_hash
(),
seq_hash
);
}
}
// ============================================================================
// UPGRADE FUNCTION TESTS
// ============================================================================
mod
upgrade_function_tests
{
use
super
::
*
;
#[test]
fn
test_upgrade_function_finds_active_blocks
()
{
let
manager
=
create_test_manager
(
10
);
let
token_block
=
create_test_token_block_from_iota
(
250
);
let
seq_hash
=
token_block
.kvbm_sequence_hash
();
// Register a block (this puts it in active pool initially)
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
let
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
let
immutable_block
=
immutable_blocks
.into_iter
()
.next
()
.unwrap
();
// Create a weak block and test upgrade
let
weak_block
=
immutable_block
.downgrade
();
let
upgraded
=
weak_block
.upgrade
()
.expect
(
"Should find block in active pool"
);
assert_eq!
(
upgraded
.sequence_hash
(),
seq_hash
);
}
#[test]
fn
test_upgrade_function_finds_inactive_blocks
()
{
let
manager
=
create_test_manager
(
20
);
let
token_block
=
create_test_token_block_from_iota
(
350
);
let
seq_hash
=
token_block
.kvbm_sequence_hash
();
// Register a block
let
weak_block
=
{
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
let
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
let
immutable_block
=
immutable_blocks
.into_iter
()
.next
()
.unwrap
();
immutable_block
.downgrade
()
};
// Force the block to potentially move to inactive pool by creating many other blocks
for
i
in
0
..
10
{
let
tokens
=
vec!
[
400
+
i
,
401
+
i
,
402
+
i
,
403
+
i
];
let
token_block
=
create_token_block
(
&
tokens
);
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate blocks"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
let
_
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
}
// Try to upgrade - should still find the original block
let
upgraded
=
weak_block
.upgrade
();
if
let
Some
(
block
)
=
upgraded
{
assert_eq!
(
block
.sequence_hash
(),
seq_hash
);
}
}
}
// ============================================================================
// ERROR HANDLING AND EDGE CASE TESTS
// ============================================================================
mod
error_handling_tests
{
use
super
::
*
;
#[test]
fn
test_allocation_exhaustion
()
{
let
manager
=
create_test_manager
(
3
);
// Allocate all blocks
let
blocks1
=
manager
.allocate_blocks
(
2
)
.expect
(
"Should allocate 2 blocks"
);
let
blocks2
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate 1 block"
);
// Try to allocate more - should fail
let
blocks3
=
manager
.allocate_blocks
(
1
);
assert
!
(
blocks3
.is_none
(),
"Should not be able to allocate when pool is empty"
);
// Drop some blocks and try again
drop
(
blocks1
);
drop
(
blocks2
);
// Blocks should be returned to pool automatically
let
blocks4
=
manager
.allocate_blocks
(
1
);
assert
!
(
blocks4
.is_some
(),
"Should be able to allocate after blocks are returned"
);
}
#[test]
fn
test_empty_sequence_matching
()
{
let
manager
=
create_test_manager
(
10
);
let
matched_blocks
=
manager
.match_blocks
(
&
[]);
assert_eq!
(
matched_blocks
.len
(),
0
);
}
#[test]
fn
test_register_empty_block_list
()
{
let
manager
=
create_test_manager
(
10
);
let
immutable_blocks
=
manager
.register_blocks
(
vec!
[]);
assert_eq!
(
immutable_blocks
.len
(),
0
);
}
}
// ============================================================================
// INTEGRATION TESTS
// ============================================================================
mod
integration_tests
{
use
super
::
*
;
#[test]
fn
test_full_lifecycle_single_block
()
{
let
manager
=
create_test_manager
(
10
);
// 1. Allocate a mutable block
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate"
);
let
mutable_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
();
let
block_id
=
mutable_block
.block_id
();
// 2. Complete the block
let
token_block
=
create_test_token_block_from_iota
(
1
);
let
seq_hash
=
token_block
.kvbm_sequence_hash
();
let
complete_block
=
mutable_block
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
assert_eq!
(
complete_block
.block_id
(),
block_id
);
assert_eq!
(
complete_block
.sequence_hash
(),
seq_hash
);
// 3. Register the block
let
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
let
immutable_block
=
immutable_blocks
.into_iter
()
.next
()
.unwrap
();
assert_eq!
(
immutable_block
.block_id
(),
block_id
);
assert_eq!
(
immutable_block
.sequence_hash
(),
seq_hash
);
// 4. Match the block
let
matched_blocks
=
manager
.match_blocks
(
&
[
seq_hash
]);
assert_eq!
(
matched_blocks
.len
(),
1
);
assert_eq!
(
matched_blocks
[
0
]
.sequence_hash
(),
seq_hash
);
// 5. Create weak reference and upgrade
let
weak_block
=
immutable_block
.downgrade
();
let
upgraded_block
=
weak_block
.upgrade
()
.expect
(
"Should upgrade"
);
assert_eq!
(
upgraded_block
.sequence_hash
(),
seq_hash
);
}
#[rstest]
#[case(
"lru"
,
|
b:
BlockManagerConfigBuilder
<
TestBlockData
>|
b
.
with_lru_backend())]
#[case(
"multi_lru"
,
|
b:
BlockManagerConfigBuilder
<
TestBlockData
>|
b
.
with_multi_lru_backend())]
fn
test_multiple_blocks_different_backends
(
#[case]
backend_name
:
&
str
,
#[case]
backend_builder
:
fn
(
BlockManagerConfigBuilder
<
TestBlockData
>
,
)
->
BlockManagerConfigBuilder
<
TestBlockData
>
,
)
{
let
registry
=
BlockRegistry
::
builder
()
.frequency_tracker
(
FrequencyTrackingCapacity
::
default
()
.create_tracker
())
.build
();
let
manager
=
backend_builder
(
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
20
)
.block_size
(
4
)
.registry
(
registry
),
)
.build
()
.expect
(
"Should build"
);
// Allocate, complete, and register blocks using BlockSequenceBuilder
let
base
=
1000
;
// Use fixed base since we only test one backend per test now
let
tokens
:
Vec
<
u32
>
=
(
base
as
u32
..
base
as
u32
+
20
)
.collect
();
// 5 blocks * 4 tokens each = 20 tokens
let
mut
seq_hashes
=
Vec
::
new
();
let
mut
complete_blocks
=
Vec
::
new
();
// Create token blocks from sequence
let
token_blocks
=
{
let
token_seq
=
dynamo_tokens
::
TokenBlockSequence
::
from_slice
(
&
tokens
,
4
,
Some
(
42
));
token_seq
.blocks
()
.to_vec
()
};
for
token_block
in
token_blocks
.iter
()
{
let
seq_hash
=
token_block
.kvbm_sequence_hash
();
seq_hashes
.push
(
seq_hash
);
// Allocate mutable block and complete it
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
token_block
)
.expect
(
"Should complete block"
);
complete_blocks
.push
(
complete_block
);
}
// Register all blocks
let
_
immutable_blocks
=
manager
.register_blocks
(
complete_blocks
);
// Verify all blocks can be matched
let
matched_blocks
=
manager
.match_blocks
(
&
seq_hashes
);
assert_eq!
(
matched_blocks
.len
(),
5
,
"Manager with {} backend should match all blocks"
,
backend_name
);
}
#[test]
fn
test_concurrent_allocation_simulation
()
{
let
manager
=
create_test_manager
(
50
);
// Simulate concurrent allocations by interleaving operations
let
mut
all_blocks
=
Vec
::
new
();
let
mut
all_hashes
=
Vec
::
new
();
// Phase 1: Allocate and complete some blocks
for
i
in
0
..
10
{
let
tokens
=
vec!
[
2000
+
i
,
2001
+
i
,
2002
+
i
,
2003
+
i
];
let
token_block
=
create_token_block
(
&
tokens
);
all_hashes
.push
(
token_block
.kvbm_sequence_hash
());
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
all_blocks
.push
(
complete_block
);
}
// Phase 2: Register half the blocks
let
mut
remaining_blocks
=
all_blocks
.split_off
(
5
);
let
_
immutable_blocks1
=
manager
.register_blocks
(
all_blocks
);
// Phase 3: Allocate more blocks while some are registered
for
i
in
10
..
15
{
let
tokens
=
vec!
[
2000
+
i
,
2001
+
i
,
2002
+
i
,
2003
+
i
];
let
token_block
=
create_token_block
(
&
tokens
);
all_hashes
.push
(
token_block
.kvbm_sequence_hash
());
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete block"
);
remaining_blocks
.push
(
complete_block
);
}
// Phase 4: Register remaining blocks
let
_
immutable_blocks2
=
manager
.register_blocks
(
remaining_blocks
);
// Phase 5: Verify we can match all registered blocks
let
matched_blocks
=
manager
.match_blocks
(
&
all_hashes
);
assert_eq!
(
matched_blocks
.len
(),
15
,
"Should match all registered blocks"
);
}
#[test]
fn
test_shared_registry_across_managers
()
{
// Create shared registry with frequency tracking
let
tracker
=
FrequencyTrackingCapacity
::
Medium
.create_tracker
();
let
registry
=
BlockRegistry
::
builder
()
.frequency_tracker
(
tracker
)
.build
();
#[derive(Clone,
Debug)]
struct
G1
;
#[derive(Clone,
Debug)]
struct
G2
;
// Create two managers with different metadata types and policies
let
manager1
=
BlockManager
::
<
G1
>
::
builder
()
.block_count
(
100
)
.block_size
(
4
)
.registry
(
registry
.clone
())
.duplication_policy
(
BlockDuplicationPolicy
::
Allow
)
.with_multi_lru_backend
()
.build
()
.expect
(
"Should build manager1"
);
let
manager2
=
BlockManager
::
<
G2
>
::
builder
()
.block_count
(
100
)
.block_size
(
4
)
.registry
(
registry
.clone
())
.duplication_policy
(
BlockDuplicationPolicy
::
Reject
)
.with_multi_lru_backend
()
.build
()
.expect
(
"Should build manager2"
);
// Verify both managers work
assert_eq!
(
manager1
.total_blocks
(),
100
);
assert_eq!
(
manager2
.total_blocks
(),
100
);
// Verify they share the same registry (frequency tracking works across both)
let
token_block
=
create_test_token_block_from_iota
(
3000
);
let
seq_hash
=
token_block
.kvbm_sequence_hash
();
// Register in manager1
let
mutable_blocks1
=
manager1
.allocate_blocks
(
1
)
.expect
(
"Should allocate"
);
let
complete_block1
=
mutable_blocks1
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete"
);
let
_
immutable1
=
manager1
.register_blocks
(
vec!
[
complete_block1
]);
// Both managers should see the registered block count in shared registry
assert
!
(
registry
.is_registered
(
seq_hash
));
}
}
mod
capacity_lifecycle_tests
{
use
super
::
*
;
/// Build a BlockManager with any backend. Always includes frequency_tracker
/// so MultiLRU works; LRU/Lineage ignore it.
fn
create_backend_manager
(
block_count
:
usize
,
backend_builder
:
fn
(
BlockManagerConfigBuilder
<
TestBlockData
>
,
)
->
BlockManagerConfigBuilder
<
TestBlockData
>
,
)
->
BlockManager
<
TestBlockData
>
{
let
registry
=
BlockRegistry
::
builder
()
.frequency_tracker
(
FrequencyTrackingCapacity
::
default
()
.create_tracker
())
.build
();
backend_builder
(
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
block_count
)
.block_size
(
4
)
.registry
(
registry
),
)
.build
()
.expect
(
"Should build manager"
)
}
/// Allocate N, complete each with a unique token block, register all.
/// Returns the ImmutableBlocks.
fn
allocate_complete_register_all
(
manager
:
&
BlockManager
<
TestBlockData
>
,
block_count
:
usize
,
iota_base
:
u32
,
)
->
Vec
<
ImmutableBlock
<
TestBlockData
>>
{
let
mutable
=
manager
.allocate_blocks
(
block_count
)
.expect
(
"allocate failed"
);
let
complete
:
Vec
<
_
>
=
mutable
.into_iter
()
.enumerate
()
.map
(|(
i
,
mb
)|
{
let
tb
=
create_iota_token_block
(
iota_base
+
(
i
as
u32
*
4
),
4
);
mb
.complete
(
&
tb
)
.expect
(
"complete failed"
)
})
.collect
();
manager
.register_blocks
(
complete
)
}
// ====================================================================
// 1. Full capacity register and return to inactive
// ====================================================================
#[rstest]
#[case(
"lru"
,
|
b:
BlockManagerConfigBuilder
<
TestBlockData
>|
b
.
with_lru_backend())]
#[case(
"multi_lru"
,
|
b:
BlockManagerConfigBuilder
<
TestBlockData
>|
b
.
with_multi_lru_backend())]
#[case(
"lineage"
,
|
b:
BlockManagerConfigBuilder
<
TestBlockData
>|
b
.
with_lineage_backend())]
fn
test_full_capacity_register_and_return_to_inactive
(
#[case]
_
backend_name
:
&
str
,
#[case]
backend_builder
:
fn
(
BlockManagerConfigBuilder
<
TestBlockData
>
,
)
->
BlockManagerConfigBuilder
<
TestBlockData
>
,
)
{
let
manager
=
create_backend_manager
(
32
,
backend_builder
);
// Allocate, complete, register all 32
let
immutable
=
allocate_complete_register_all
(
&
manager
,
32
,
5000
);
assert_eq!
(
manager
.inactive_pool
.len
(),
0
);
assert_eq!
(
manager
.reset_pool
.len
(),
0
);
let
snap
=
manager
.metrics
.snapshot
();
assert_eq!
(
snap
.reset_pool_size
,
0
);
assert_eq!
(
snap
.inactive_pool_size
,
0
);
// Drop all ImmutableBlocks → should all land in inactive pool
drop
(
immutable
);
assert_eq!
(
manager
.inactive_pool
.len
(),
32
);
assert_eq!
(
manager
.reset_pool
.len
(),
0
);
// Check metrics
let
snap
=
manager
.metrics
.snapshot
();
assert_eq!
(
snap
.allocations
,
32
);
assert_eq!
(
snap
.registrations
,
32
);
assert_eq!
(
snap
.inflight_immutable
,
0
);
assert_eq!
(
snap
.inflight_mutable
,
0
);
assert_eq!
(
snap
.inactive_pool_size
,
32
);
assert_eq!
(
snap
.reset_pool_size
,
0
);
// Check totals
assert_eq!
(
manager
.available_blocks
(),
32
);
assert_eq!
(
manager
.total_blocks
(),
32
);
}
// ====================================================================
// 2. Full capacity eviction cycle
// ====================================================================
#[rstest]
#[case(
"lru"
,
|
b:
BlockManagerConfigBuilder
<
TestBlockData
>|
b
.
with_lru_backend())]
#[case(
"multi_lru"
,
|
b:
BlockManagerConfigBuilder
<
TestBlockData
>|
b
.
with_multi_lru_backend())]
#[case(
"lineage"
,
|
b:
BlockManagerConfigBuilder
<
TestBlockData
>|
b
.
with_lineage_backend())]
fn
test_full_capacity_eviction_cycle
(
#[case]
_
backend_name
:
&
str
,
#[case]
backend_builder
:
fn
(
BlockManagerConfigBuilder
<
TestBlockData
>
,
)
->
BlockManagerConfigBuilder
<
TestBlockData
>
,
)
{
let
manager
=
create_backend_manager
(
16
,
backend_builder
);
// Allocate, register all 16
let
immutable
=
allocate_complete_register_all
(
&
manager
,
16
,
6000
);
assert_eq!
(
manager
.reset_pool
.len
(),
0
);
assert_eq!
(
manager
.inactive_pool
.len
(),
0
);
// Drop all → inactive pool
drop
(
immutable
);
assert_eq!
(
manager
.inactive_pool
.len
(),
16
);
assert_eq!
(
manager
.reset_pool
.len
(),
0
);
let
snap
=
manager
.metrics
.snapshot
();
assert_eq!
(
snap
.inactive_pool_size
,
16
);
assert_eq!
(
snap
.reset_pool_size
,
0
);
// Allocate 16 again (evicts from inactive)
let
mutable
=
manager
.allocate_blocks
(
16
)
.expect
(
"second allocate failed"
);
assert_eq!
(
manager
.inactive_pool
.len
(),
0
);
assert_eq!
(
manager
.reset_pool
.len
(),
0
);
let
snap
=
manager
.metrics
.snapshot
();
assert_eq!
(
snap
.inactive_pool_size
,
0
);
assert_eq!
(
snap
.reset_pool_size
,
0
);
// Drop mutable blocks → reset pool
drop
(
mutable
);
assert_eq!
(
manager
.reset_pool
.len
(),
16
);
assert_eq!
(
manager
.inactive_pool
.len
(),
0
);
// Check metrics
let
snap
=
manager
.metrics
.snapshot
();
assert_eq!
(
snap
.evictions
,
16
);
assert_eq!
(
snap
.allocations
,
32
);
assert_eq!
(
snap
.reset_pool_size
,
16
);
assert_eq!
(
snap
.inactive_pool_size
,
0
);
}
// ====================================================================
// 3. Mutable drops go to reset, not inactive
// ====================================================================
#[test]
fn
test_mutable_drops_go_to_reset_not_inactive
()
{
let
manager
=
create_backend_manager
(
16
,
|
b
|
b
.with_lru_backend
());
let
mutable
=
manager
.allocate_blocks
(
16
)
.expect
(
"allocate failed"
);
assert_eq!
(
manager
.reset_pool
.len
(),
0
);
assert_eq!
(
manager
.inactive_pool
.len
(),
0
);
// Drop all mutable blocks → reset pool
drop
(
mutable
);
assert_eq!
(
manager
.reset_pool
.len
(),
16
);
assert_eq!
(
manager
.inactive_pool
.len
(),
0
);
let
snap
=
manager
.metrics
.snapshot
();
assert_eq!
(
snap
.inflight_mutable
,
0
);
assert_eq!
(
snap
.registrations
,
0
);
}
// ====================================================================
// 4. Complete drops go to reset, not inactive
// ====================================================================
#[test]
fn
test_complete_drops_go_to_reset_not_inactive
()
{
let
manager
=
create_backend_manager
(
16
,
|
b
|
b
.with_lru_backend
());
let
mutable
=
manager
.allocate_blocks
(
16
)
.expect
(
"allocate failed"
);
let
complete
:
Vec
<
_
>
=
mutable
.into_iter
()
.enumerate
()
.map
(|(
i
,
mb
)|
{
let
tb
=
create_iota_token_block
(
7000
+
(
i
as
u32
*
4
),
4
);
mb
.complete
(
&
tb
)
.expect
(
"complete failed"
)
})
.collect
();
assert_eq!
(
manager
.reset_pool
.len
(),
0
);
// Drop all CompleteBlocks (not registered) → reset pool
drop
(
complete
);
assert_eq!
(
manager
.reset_pool
.len
(),
16
);
assert_eq!
(
manager
.inactive_pool
.len
(),
0
);
let
snap
=
manager
.metrics
.snapshot
();
assert_eq!
(
snap
.stagings
,
16
);
assert_eq!
(
snap
.registrations
,
0
);
}
// ====================================================================
// 5. Mixed return paths
// ====================================================================
#[rstest]
#[case(
"lru"
,
|
b:
BlockManagerConfigBuilder
<
TestBlockData
>|
b
.
with_lru_backend())]
#[case(
"multi_lru"
,
|
b:
BlockManagerConfigBuilder
<
TestBlockData
>|
b
.
with_multi_lru_backend())]
#[case(
"lineage"
,
|
b:
BlockManagerConfigBuilder
<
TestBlockData
>|
b
.
with_lineage_backend())]
fn
test_mixed_return_paths
(
#[case]
_
backend_name
:
&
str
,
#[case]
backend_builder
:
fn
(
BlockManagerConfigBuilder
<
TestBlockData
>
,
)
->
BlockManagerConfigBuilder
<
TestBlockData
>
,
)
{
let
manager
=
create_backend_manager
(
24
,
backend_builder
);
let
mutable
=
manager
.allocate_blocks
(
24
)
.expect
(
"allocate failed"
);
let
mut
mutable_iter
=
mutable
.into_iter
();
// Group A (8): drop as MutableBlocks
{
let
group_a
:
Vec
<
_
>
=
mutable_iter
.by_ref
()
.take
(
8
)
.collect
();
drop
(
group_a
);
}
assert_eq!
(
manager
.reset_pool
.len
(),
8
);
assert_eq!
(
manager
.metrics
.snapshot
()
.reset_pool_size
,
8
);
// Group B (8): complete, drop as CompleteBlocks
{
let
group_b
:
Vec
<
_
>
=
mutable_iter
.by_ref
()
.take
(
8
)
.enumerate
()
.map
(|(
i
,
mb
)|
{
let
tb
=
create_iota_token_block
(
8000
+
(
i
as
u32
*
4
),
4
);
mb
.complete
(
&
tb
)
.expect
(
"complete failed"
)
})
.collect
();
drop
(
group_b
);
}
assert_eq!
(
manager
.reset_pool
.len
(),
16
);
assert_eq!
(
manager
.metrics
.snapshot
()
.reset_pool_size
,
16
);
// Group C (8): complete, register, hold ImmutableBlocks
let
group_c_complete
:
Vec
<
_
>
=
mutable_iter
.enumerate
()
.map
(|(
i
,
mb
)|
{
let
tb
=
create_iota_token_block
(
8100
+
(
i
as
u32
*
4
),
4
);
mb
.complete
(
&
tb
)
.expect
(
"complete failed"
)
})
.collect
();
let
group_c_immutable
=
manager
.register_blocks
(
group_c_complete
);
assert_eq!
(
manager
.inactive_pool
.len
(),
0
);
// Drop Group C → inactive pool
drop
(
group_c_immutable
);
assert_eq!
(
manager
.inactive_pool
.len
(),
8
);
assert_eq!
(
manager
.reset_pool
.len
(),
16
);
// Check totals
assert_eq!
(
manager
.available_blocks
(),
24
);
// Check metrics
let
snap
=
manager
.metrics
.snapshot
();
assert_eq!
(
snap
.allocations
,
24
);
assert_eq!
(
snap
.stagings
,
16
);
// Group B (8) + Group C (8)
assert_eq!
(
snap
.registrations
,
8
);
assert_eq!
(
snap
.inflight_mutable
,
0
);
assert_eq!
(
snap
.inflight_immutable
,
0
);
assert_eq!
(
snap
.inactive_pool_size
,
8
);
assert_eq!
(
snap
.reset_pool_size
,
16
);
}
// ====================================================================
// 6. MultiLRU all cold blocks at capacity (regression)
// ====================================================================
#[test]
fn
test_multi_lru_all_cold_blocks_at_capacity
()
{
let
manager
=
create_backend_manager
(
64
,
|
b
|
b
.with_multi_lru_backend
());
// Allocate, register all 64 (no frequency touches → all cold)
let
immutable
=
allocate_complete_register_all
(
&
manager
,
64
,
9000
);
// Drop all → all go to level 0 (cold). With old div_ceil(4)=16
// per-level capacity this would panic at block 17.
drop
(
immutable
);
assert_eq!
(
manager
.inactive_pool
.len
(),
64
);
let
snap
=
manager
.metrics
.snapshot
();
assert_eq!
(
snap
.evictions
,
0
);
assert_eq!
(
snap
.allocations
,
64
);
}
// ====================================================================
// 7. MultiLRU mixed frequency levels
// ====================================================================
#[test]
fn
test_multi_lru_mixed_frequency_levels
()
{
// thresholds [3, 8, 15]: cold=0-2, warm=3-7, hot=8-14, very_hot=15
let
registry
=
BlockRegistry
::
builder
()
.frequency_tracker
(
FrequencyTrackingCapacity
::
default
()
.create_tracker
())
.build
();
let
manager
=
BlockManager
::
<
TestBlockData
>
::
builder
()
.block_count
(
32
)
.block_size
(
4
)
.registry
(
registry
)
.with_multi_lru_backend
()
.build
()
.expect
(
"Should build manager"
);
// Allocate, register all 32
let
immutable
=
allocate_complete_register_all
(
&
manager
,
32
,
10000
);
// Touch frequency tracker for different blocks to spread across levels
let
tracker
=
manager
.block_registry
()
.frequency_tracker
()
.unwrap
();
for
block
in
&
immutable
{
let
hash
=
block
.sequence_hash
();
let
idx
=
block
.block_id
();
let
touches
=
if
idx
<
8
{
0
// cold: 0-7 untouched
}
else
if
idx
<
16
{
3
// warm: 8-15
}
else
if
idx
<
24
{
8
// hot: 16-23
}
else
{
15
// very hot: 24-31
};
for
_
in
0
..
touches
{
tracker
.touch
(
hash
.as_u128
());
}
}
// Drop all → distributed across 4 levels
drop
(
immutable
);
assert_eq!
(
manager
.inactive_pool
.len
(),
32
);
// Allocate 32 again → evicts from all levels
let
mutable
=
manager
.allocate_blocks
(
32
)
.expect
(
"eviction allocate"
);
assert_eq!
(
manager
.inactive_pool
.len
(),
0
);
drop
(
mutable
);
let
snap
=
manager
.metrics
.snapshot
();
assert_eq!
(
snap
.evictions
,
32
);
assert_eq!
(
snap
.allocations
,
64
);
}
// ====================================================================
// 8. Double lifecycle cycle
// ====================================================================
#[rstest]
#[case(
"lru"
,
|
b:
BlockManagerConfigBuilder
<
TestBlockData
>|
b
.
with_lru_backend())]
#[case(
"multi_lru"
,
|
b:
BlockManagerConfigBuilder
<
TestBlockData
>|
b
.
with_multi_lru_backend())]
#[case(
"lineage"
,
|
b:
BlockManagerConfigBuilder
<
TestBlockData
>|
b
.
with_lineage_backend())]
fn
test_double_lifecycle_cycle
(
#[case]
_
backend_name
:
&
str
,
#[case]
backend_builder
:
fn
(
BlockManagerConfigBuilder
<
TestBlockData
>
,
)
->
BlockManagerConfigBuilder
<
TestBlockData
>
,
)
{
let
manager
=
create_backend_manager
(
16
,
backend_builder
);
let
m
=
&
manager
.metrics
;
// Cycle 1: allocate, register, drop → inactive
{
let
immutable
=
allocate_complete_register_all
(
&
manager
,
16
,
11000
);
drop
(
immutable
);
}
assert_eq!
(
manager
.inactive_pool
.len
(),
16
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inactive_pool_size
,
16
);
assert_eq!
(
snap
.reset_pool_size
,
0
);
// Evict all: allocate from inactive, drop mutable → reset
{
let
mutable
=
manager
.allocate_blocks
(
16
)
.expect
(
"eviction allocate"
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inactive_pool_size
,
0
);
assert_eq!
(
snap
.reset_pool_size
,
0
);
drop
(
mutable
);
}
assert_eq!
(
manager
.reset_pool
.len
(),
16
);
assert_eq!
(
manager
.inactive_pool
.len
(),
0
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.reset_pool_size
,
16
);
assert_eq!
(
snap
.inactive_pool_size
,
0
);
// Cycle 2: allocate, register (different tokens), drop → inactive
{
let
immutable
=
allocate_complete_register_all
(
&
manager
,
16
,
12000
);
drop
(
immutable
);
}
assert_eq!
(
manager
.inactive_pool
.len
(),
16
);
// Check metrics
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.allocations
,
48
);
assert_eq!
(
snap
.registrations
,
32
);
assert_eq!
(
snap
.evictions
,
16
);
assert_eq!
(
snap
.inactive_pool_size
,
16
);
assert_eq!
(
snap
.reset_pool_size
,
0
);
// Check totals
assert_eq!
(
manager
.available_blocks
(),
16
);
assert_eq!
(
manager
.total_blocks
(),
16
);
}
}
// ============================================================================
// SCAN MATCHES POOL SIZE GAUGE TESTS
// ============================================================================
mod
scan_matches_tests
{
use
super
::
*
;
#[test]
fn
test_scan_matches_with_pool_size_gauges
()
{
let
manager
=
create_test_manager
(
10
);
let
m
=
manager
.metrics
();
// Register 3 blocks with distinct hashes
let
mut
seq_hashes
=
Vec
::
new
();
for
i
in
0
..
3
{
let
tb
=
create_iota_token_block
(
13000
+
(
i
as
u32
*
4
),
4
);
seq_hashes
.push
(
tb
.kvbm_sequence_hash
());
let
mutable
=
manager
.allocate_blocks
(
1
)
.expect
(
"allocate"
);
let
complete
=
mutable
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
tb
)
.expect
(
"complete"
);
let
immutable
=
manager
.register_blocks
(
vec!
[
complete
]);
drop
(
immutable
);
}
// All 3 should be in inactive pool
assert_eq!
(
manager
.inactive_pool
.len
(),
3
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inactive_pool_size
,
3
);
assert_eq!
(
snap
.reset_pool_size
,
7
);
// scan_matches with 2 matching + 1 missing hash
let
missing_hash
=
create_iota_token_block
(
99000
,
4
)
.kvbm_sequence_hash
();
let
scan_hashes
=
vec!
[
seq_hashes
[
0
],
missing_hash
,
seq_hashes
[
2
]];
let
found
=
manager
.scan_matches
(
&
scan_hashes
,
true
);
assert_eq!
(
found
.len
(),
2
);
// inactive_pool_size decreased by 2
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inactive_pool_size
,
1
);
// Drop scanned blocks → they return to inactive
drop
(
found
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inactive_pool_size
,
3
);
}
}
lib/kvbm-logical/src/metrics/collector.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Custom `prometheus::core::Collector` that reads raw atomics at scrape time.
//!
//! External labels (e.g. `instance_id`, `worker_id`) are appended at collection time,
//! not baked in at metric creation time.
use
std
::
sync
::{
Arc
,
RwLock
};
use
prometheus
::
core
::{
Collector
,
Desc
};
use
prometheus
::
proto
::{
Gauge
,
LabelPair
,
Metric
,
MetricFamily
,
MetricType
};
use
super
::
pool_metrics
::
BlockPoolMetrics
;
/// Metric definitions: (name, help, type).
const
COUNTER_DEFS
:
&
[(
&
str
,
&
str
)]
=
&
[
(
"kvbm_allocations_total"
,
"Total blocks allocated from pools"
,
),
(
"kvbm_allocations_from_reset_total"
,
"Total blocks allocated from the reset pool"
,
),
(
"kvbm_evictions_total"
,
"Total blocks evicted from inactive pool"
,
),
(
"kvbm_registrations_total"
,
"Total blocks registered (CompleteBlock -> ImmutableBlock)"
,
),
(
"kvbm_duplicate_blocks_total"
,
"Total duplicate blocks created (Allow policy)"
,
),
(
"kvbm_registration_dedup_total"
,
"Total block registrations deduplicated (Reject policy)"
,
),
(
"kvbm_stagings_total"
,
"Total MutableBlock -> CompleteBlock transitions"
,
),
(
"kvbm_match_hashes_requested_total"
,
"Total hashes requested in match_blocks calls"
,
),
(
"kvbm_match_blocks_returned_total"
,
"Total blocks returned from match_blocks calls"
,
),
(
"kvbm_scan_hashes_requested_total"
,
"Total hashes requested in scan_matches calls"
,
),
(
"kvbm_scan_blocks_returned_total"
,
"Total blocks returned from scan_matches calls"
,
),
];
const
GAUGE_DEFS
:
&
[(
&
str
,
&
str
)]
=
&
[
(
"kvbm_inflight_mutable"
,
"Current MutableBlocks held outside pool"
,
),
(
"kvbm_inflight_immutable"
,
"Current ImmutableBlocks held outside pool"
,
),
(
"kvbm_reset_pool_size"
,
"Current reset pool size"
),
(
"kvbm_inactive_pool_size"
,
"Current inactive pool size"
),
];
/// Aggregates metrics from multiple `BlockPoolMetrics` sources and exports
/// them as Prometheus `MetricFamily` protos with per-pool-type labels.
#[derive(Clone)]
pub
struct
MetricsAggregator
{
inner
:
Arc
<
Inner
>
,
}
struct
Inner
{
sources
:
RwLock
<
Vec
<
Arc
<
BlockPoolMetrics
>>>
,
external_labels
:
RwLock
<
Vec
<
(
String
,
String
)
>>
,
descs
:
Vec
<
Desc
>
,
}
impl
MetricsAggregator
{
/// Create a new `MetricsAggregator`.
pub
fn
new
()
->
Self
{
let
mut
descs
=
Vec
::
with_capacity
(
COUNTER_DEFS
.len
()
+
GAUGE_DEFS
.len
());
for
(
name
,
help
)
in
COUNTER_DEFS
{
descs
.push
(
Desc
::
new
(
name
.to_string
(),
help
.to_string
(),
vec!
[
"pool"
.to_string
()],
Default
::
default
(),
)
.expect
(
"valid desc"
),
);
}
for
(
name
,
help
)
in
GAUGE_DEFS
{
descs
.push
(
Desc
::
new
(
name
.to_string
(),
help
.to_string
(),
vec!
[
"pool"
.to_string
()],
Default
::
default
(),
)
.expect
(
"valid desc"
),
);
}
Self
{
inner
:
Arc
::
new
(
Inner
{
sources
:
RwLock
::
new
(
Vec
::
new
()),
external_labels
:
RwLock
::
new
(
Vec
::
new
()),
descs
,
}),
}
}
/// Register a `BlockPoolMetrics` source (called by `BlockManager::build()`).
pub
fn
register_source
(
&
self
,
source
:
Arc
<
BlockPoolMetrics
>
)
{
self
.inner
.sources
.write
()
.expect
(
"sources lock poisoned"
)
.push
(
source
);
}
/// Set external labels appended at scrape time (e.g. `instance_id`, `worker_id`).
pub
fn
set_external_labels
(
&
self
,
labels
:
Vec
<
(
String
,
String
)
>
)
{
*
self
.inner
.external_labels
.write
()
.expect
(
"external_labels lock poisoned"
)
=
labels
;
}
/// Register this collector with a `prometheus::Registry`.
pub
fn
register_with
(
&
self
,
registry
:
&
prometheus
::
Registry
)
->
Result
<
(),
prometheus
::
Error
>
{
registry
.register
(
Box
::
new
(
self
.clone
()))
}
}
impl
Default
for
MetricsAggregator
{
fn
default
()
->
Self
{
Self
::
new
()
}
}
impl
Collector
for
MetricsAggregator
{
fn
desc
(
&
self
)
->
Vec
<&
Desc
>
{
self
.inner.descs
.iter
()
.collect
()
}
fn
collect
(
&
self
)
->
Vec
<
MetricFamily
>
{
let
sources
=
self
.inner.sources
.read
()
.expect
(
"sources lock poisoned"
);
let
ext_labels
=
self
.inner
.external_labels
.read
()
.expect
(
"external_labels lock poisoned"
);
let
mut
families
:
Vec
<
MetricFamily
>
=
Vec
::
new
();
for
source
in
sources
.iter
()
{
let
snap
=
source
.snapshot
();
let
pool_label
=
source
.type_label
();
let
mut
base_labels
:
Vec
<
LabelPair
>
=
Vec
::
with_capacity
(
1
+
ext_labels
.len
());
let
mut
pool_lp
=
LabelPair
::
default
();
pool_lp
.set_name
(
"pool"
.to_string
());
pool_lp
.set_value
(
pool_label
.to_string
());
base_labels
.push
(
pool_lp
);
for
(
k
,
v
)
in
ext_labels
.iter
()
{
let
mut
lp
=
LabelPair
::
default
();
lp
.set_name
(
k
.clone
());
lp
.set_value
(
v
.clone
());
base_labels
.push
(
lp
);
}
// Counter values in order matching COUNTER_DEFS
let
counter_values
:
[
u64
;
11
]
=
[
snap
.allocations
,
snap
.allocations_from_reset
,
snap
.evictions
,
snap
.registrations
,
snap
.duplicate_blocks
,
snap
.registration_dedup
,
snap
.stagings
,
snap
.match_hashes_requested
,
snap
.match_blocks_returned
,
snap
.scan_hashes_requested
,
snap
.scan_blocks_returned
,
];
for
(
i
,
(
name
,
help
))
in
COUNTER_DEFS
.iter
()
.enumerate
()
{
let
mut
m
=
Metric
::
default
();
m
.set_label
(
base_labels
.clone
());
let
mut
c
=
prometheus
::
proto
::
Counter
::
default
();
c
.set_value
(
counter_values
[
i
]
as
f64
);
m
.set_counter
(
c
);
let
mut
mf
=
MetricFamily
::
default
();
mf
.set_name
(
name
.to_string
());
mf
.set_help
(
help
.to_string
());
mf
.set_field_type
(
MetricType
::
COUNTER
);
mf
.set_metric
(
vec!
[
m
]);
families
.push
(
mf
);
}
// Gauge values in order matching GAUGE_DEFS
let
gauge_values
:
[
i64
;
4
]
=
[
snap
.inflight_mutable
,
snap
.inflight_immutable
,
snap
.reset_pool_size
,
snap
.inactive_pool_size
,
];
for
(
i
,
(
name
,
help
))
in
GAUGE_DEFS
.iter
()
.enumerate
()
{
let
mut
m
=
Metric
::
default
();
m
.set_label
(
base_labels
.clone
());
let
mut
g
=
Gauge
::
default
();
g
.set_value
(
gauge_values
[
i
]
as
f64
);
m
.set_gauge
(
g
);
let
mut
mf
=
MetricFamily
::
default
();
mf
.set_name
(
name
.to_string
());
mf
.set_help
(
help
.to_string
());
mf
.set_field_type
(
MetricType
::
GAUGE
);
mf
.set_metric
(
vec!
[
m
]);
families
.push
(
mf
);
}
}
// Merge families with the same name (when multiple sources)
if
sources
.len
()
>
1
{
let
mut
merged
:
Vec
<
MetricFamily
>
=
Vec
::
new
();
for
mut
family
in
families
{
if
let
Some
(
existing
)
=
merged
.iter_mut
()
.find
(|
f
|
f
.name
()
==
family
.name
())
{
existing
.mut_metric
()
.extend
(
family
.take_metric
());
}
else
{
merged
.push
(
family
);
}
}
merged
}
else
{
families
}
}
}
#[cfg(test)]
#[allow(deprecated)]
mod
tests
{
use
super
::
*
;
use
prometheus
::
core
::
Collector
;
#[test]
fn
test_empty_aggregator_collects_nothing
()
{
let
agg
=
MetricsAggregator
::
new
();
let
families
=
agg
.collect
();
assert
!
(
families
.is_empty
());
}
#[test]
fn
test_single_source_collect
()
{
let
agg
=
MetricsAggregator
::
new
();
let
metrics
=
Arc
::
new
(
BlockPoolMetrics
::
new
(
"G1"
.to_string
()));
metrics
.inc_allocations
(
10
);
metrics
.inc_evictions
(
3
);
metrics
.set_reset_pool_size
(
42
);
agg
.register_source
(
metrics
);
let
families
=
agg
.collect
();
assert_eq!
(
families
.len
(),
COUNTER_DEFS
.len
()
+
GAUGE_DEFS
.len
());
// Find allocations counter
let
alloc_family
=
families
.iter
()
.find
(|
f
|
f
.get_name
()
==
"kvbm_allocations_total"
)
.expect
(
"should have allocations family"
);
assert_eq!
(
alloc_family
.get_field_type
(),
MetricType
::
COUNTER
);
let
m
=
&
alloc_family
.get_metric
()[
0
];
assert_eq!
(
m
.get_counter
()
.value
(),
10.0
);
assert_eq!
(
m
.get_label
()[
0
]
.get_name
(),
"pool"
);
assert_eq!
(
m
.get_label
()[
0
]
.get_value
(),
"G1"
);
// Find reset_pool_size gauge
let
reset_family
=
families
.iter
()
.find
(|
f
|
f
.get_name
()
==
"kvbm_reset_pool_size"
)
.expect
(
"should have reset pool size family"
);
assert_eq!
(
reset_family
.get_field_type
(),
MetricType
::
GAUGE
);
assert_eq!
(
reset_family
.get_metric
()[
0
]
.get_gauge
()
.value
(),
42.0
);
}
#[test]
fn
test_external_labels
()
{
let
agg
=
MetricsAggregator
::
new
();
let
metrics
=
Arc
::
new
(
BlockPoolMetrics
::
new
(
"G1"
.to_string
()));
agg
.register_source
(
metrics
);
agg
.set_external_labels
(
vec!
[
(
"instance_id"
.to_string
(),
"node-1"
.to_string
()),
(
"worker_id"
.to_string
(),
"w0"
.to_string
()),
]);
let
families
=
agg
.collect
();
let
alloc_family
=
families
.iter
()
.find
(|
f
|
f
.get_name
()
==
"kvbm_allocations_total"
)
.unwrap
();
let
labels
=
alloc_family
.get_metric
()[
0
]
.get_label
();
assert_eq!
(
labels
.len
(),
3
);
// pool + 2 external
assert_eq!
(
labels
[
1
]
.get_name
(),
"instance_id"
);
assert_eq!
(
labels
[
1
]
.get_value
(),
"node-1"
);
assert_eq!
(
labels
[
2
]
.get_name
(),
"worker_id"
);
assert_eq!
(
labels
[
2
]
.get_value
(),
"w0"
);
}
#[test]
fn
test_multiple_sources_merged
()
{
let
agg
=
MetricsAggregator
::
new
();
let
m1
=
Arc
::
new
(
BlockPoolMetrics
::
new
(
"G1"
.to_string
()));
let
m2
=
Arc
::
new
(
BlockPoolMetrics
::
new
(
"G2"
.to_string
()));
m1
.inc_allocations
(
5
);
m2
.inc_allocations
(
10
);
agg
.register_source
(
m1
);
agg
.register_source
(
m2
);
let
families
=
agg
.collect
();
// Families should be merged by name
let
alloc_family
=
families
.iter
()
.find
(|
f
|
f
.get_name
()
==
"kvbm_allocations_total"
)
.expect
(
"should have allocations family"
);
assert_eq!
(
alloc_family
.get_metric
()
.len
(),
2
);
let
values
:
Vec
<
f64
>
=
alloc_family
.get_metric
()
.iter
()
.map
(|
m
|
m
.get_counter
()
.value
())
.collect
();
assert
!
(
values
.contains
(
&
5.0
));
assert
!
(
values
.contains
(
&
10.0
));
}
#[test]
fn
test_register_with_prometheus_registry
()
{
let
agg
=
MetricsAggregator
::
new
();
let
metrics
=
Arc
::
new
(
BlockPoolMetrics
::
new
(
"G1"
.to_string
()));
metrics
.inc_allocations
(
42
);
agg
.register_source
(
metrics
);
let
registry
=
prometheus
::
Registry
::
new
();
agg
.register_with
(
&
registry
)
.expect
(
"should register successfully"
);
let
gathered
=
registry
.gather
();
assert
!
(
!
gathered
.is_empty
());
let
alloc_family
=
gathered
.iter
()
.find
(|
f
|
f
.get_name
()
==
"kvbm_allocations_total"
)
.expect
(
"should find allocations in gathered metrics"
);
assert_eq!
(
alloc_family
.get_metric
()[
0
]
.get_counter
()
.value
(),
42.0
);
}
#[test]
fn
test_descs_match_definitions
()
{
let
agg
=
MetricsAggregator
::
new
();
let
descs
=
agg
.desc
();
assert_eq!
(
descs
.len
(),
COUNTER_DEFS
.len
()
+
GAUGE_DEFS
.len
());
}
}
lib/kvbm-logical/src/metrics/mod.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Prometheus metrics for block pool observability.
//!
//! Architecture:
//! - [`BlockPoolMetrics`]: Raw `AtomicU64`/`AtomicI64` counters and gauges. Zero overhead on hot path.
//! - [`MetricsAggregator`]: Custom `prometheus::core::Collector` that reads atomics at scrape time.
//! - [`StatsCollector`]: Optional periodic sampler computing rates, gradients, and hit ratios.
mod
collector
;
mod
pool_metrics
;
mod
stats
;
pub
use
collector
::
MetricsAggregator
;
pub
use
pool_metrics
::{
BlockPoolMetrics
,
MetricsSnapshot
};
pub
use
stats
::{
StatsCollector
,
StatsConfig
,
StatsSnapshot
};
/// Returns the short (unqualified) type name for `T`.
///
/// Strips generic parameters and the module path, returning only the base
/// type name. May still be imperfect for deeply nested or anonymous types.
pub
fn
short_type_name
<
T
:
'static
>
()
->
String
{
let
full
=
std
::
any
::
type_name
::
<
T
>
();
let
base
=
full
.split_once
(
'<'
)
.map
(|(
b
,
_
)|
b
)
.unwrap_or
(
full
);
base
.rsplit
(
"::"
)
.next
()
.unwrap_or
(
base
)
.to_string
()
}
#[cfg(test)]
mod
tests
{
use
super
::
*
;
struct
MyType
;
#[test]
fn
test_short_type_name
()
{
let
name
=
short_type_name
::
<
MyType
>
();
assert_eq!
(
name
,
"MyType"
);
}
#[test]
fn
test_short_type_name_primitive
()
{
let
name
=
short_type_name
::
<
u32
>
();
assert_eq!
(
name
,
"u32"
);
}
#[test]
fn
test_short_type_name_generic
()
{
let
name
=
short_type_name
::
<
Vec
<
String
>>
();
assert_eq!
(
name
,
"Vec"
);
}
}
lib/kvbm-logical/src/metrics/pool_metrics.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Raw atomic counters and gauges for a single block pool type.
//!
//! All increment/decrement methods use `Ordering::Relaxed` for zero overhead on the hot path.
//! The [`MetricsAggregator`] reads these atomics at scrape time and builds Prometheus protos.
use
std
::
sync
::
atomic
::{
AtomicI64
,
AtomicU64
,
Ordering
};
/// Raw atomic metrics for a single block pool (one per `BlockManager<T>`).
///
/// Counters are monotonically increasing `AtomicU64`.
/// Gauges are bidirectional `AtomicI64`.
pub
struct
BlockPoolMetrics
{
type_label
:
String
,
// Counters (monotonic)
allocations
:
AtomicU64
,
allocations_from_reset
:
AtomicU64
,
evictions
:
AtomicU64
,
registrations
:
AtomicU64
,
duplicate_blocks
:
AtomicU64
,
registration_dedup
:
AtomicU64
,
stagings
:
AtomicU64
,
match_hashes_requested
:
AtomicU64
,
match_blocks_returned
:
AtomicU64
,
scan_hashes_requested
:
AtomicU64
,
scan_blocks_returned
:
AtomicU64
,
// Gauges (bidirectional)
inflight_mutable
:
AtomicI64
,
inflight_immutable
:
AtomicI64
,
reset_pool_size
:
AtomicI64
,
inactive_pool_size
:
AtomicI64
,
}
impl
BlockPoolMetrics
{
/// Create a new `BlockPoolMetrics` with the given type label (e.g. `"G1"`).
pub
fn
new
(
type_label
:
String
)
->
Self
{
Self
{
type_label
,
allocations
:
AtomicU64
::
new
(
0
),
allocations_from_reset
:
AtomicU64
::
new
(
0
),
evictions
:
AtomicU64
::
new
(
0
),
registrations
:
AtomicU64
::
new
(
0
),
duplicate_blocks
:
AtomicU64
::
new
(
0
),
registration_dedup
:
AtomicU64
::
new
(
0
),
stagings
:
AtomicU64
::
new
(
0
),
match_hashes_requested
:
AtomicU64
::
new
(
0
),
match_blocks_returned
:
AtomicU64
::
new
(
0
),
scan_hashes_requested
:
AtomicU64
::
new
(
0
),
scan_blocks_returned
:
AtomicU64
::
new
(
0
),
inflight_mutable
:
AtomicI64
::
new
(
0
),
inflight_immutable
:
AtomicI64
::
new
(
0
),
reset_pool_size
:
AtomicI64
::
new
(
0
),
inactive_pool_size
:
AtomicI64
::
new
(
0
),
}
}
/// The pool type label (e.g. `"G1"`, `"G2"`).
#[inline(always)]
pub
fn
type_label
(
&
self
)
->
&
str
{
&
self
.type_label
}
// ---- Counter increments ----
#[inline(always)]
pub
fn
inc_allocations
(
&
self
,
n
:
u64
)
{
self
.allocations
.fetch_add
(
n
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
inc_allocations_from_reset
(
&
self
,
n
:
u64
)
{
self
.allocations_from_reset
.fetch_add
(
n
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
inc_evictions
(
&
self
,
n
:
u64
)
{
self
.evictions
.fetch_add
(
n
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
inc_registrations
(
&
self
)
{
self
.registrations
.fetch_add
(
1
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
inc_duplicate_blocks
(
&
self
)
{
self
.duplicate_blocks
.fetch_add
(
1
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
inc_registration_dedup
(
&
self
)
{
self
.registration_dedup
.fetch_add
(
1
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
inc_stagings
(
&
self
)
{
self
.stagings
.fetch_add
(
1
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
inc_match_hashes_requested
(
&
self
,
n
:
u64
)
{
self
.match_hashes_requested
.fetch_add
(
n
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
inc_match_blocks_returned
(
&
self
,
n
:
u64
)
{
self
.match_blocks_returned
.fetch_add
(
n
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
inc_scan_hashes_requested
(
&
self
,
n
:
u64
)
{
self
.scan_hashes_requested
.fetch_add
(
n
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
inc_scan_blocks_returned
(
&
self
,
n
:
u64
)
{
self
.scan_blocks_returned
.fetch_add
(
n
,
Ordering
::
Relaxed
);
}
// ---- Gauge operations ----
#[inline(always)]
pub
fn
inc_inflight_mutable
(
&
self
)
{
self
.inflight_mutable
.fetch_add
(
1
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
dec_inflight_mutable
(
&
self
)
{
self
.inflight_mutable
.fetch_sub
(
1
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
inc_inflight_immutable
(
&
self
)
{
self
.inflight_immutable
.fetch_add
(
1
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
dec_inflight_immutable
(
&
self
)
{
self
.inflight_immutable
.fetch_sub
(
1
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
set_reset_pool_size
(
&
self
,
size
:
i64
)
{
self
.reset_pool_size
.store
(
size
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
inc_reset_pool_size
(
&
self
)
{
self
.reset_pool_size
.fetch_add
(
1
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
dec_reset_pool_size
(
&
self
)
{
self
.reset_pool_size
.fetch_sub
(
1
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
set_inactive_pool_size
(
&
self
,
size
:
i64
)
{
self
.inactive_pool_size
.store
(
size
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
inc_inactive_pool_size
(
&
self
)
{
self
.inactive_pool_size
.fetch_add
(
1
,
Ordering
::
Relaxed
);
}
#[inline(always)]
pub
fn
dec_inactive_pool_size
(
&
self
)
{
self
.inactive_pool_size
.fetch_sub
(
1
,
Ordering
::
Relaxed
);
}
// ---- Snapshot for stats collector ----
/// Take a point-in-time snapshot of all metrics.
pub
fn
snapshot
(
&
self
)
->
MetricsSnapshot
{
MetricsSnapshot
{
allocations
:
self
.allocations
.load
(
Ordering
::
Relaxed
),
allocations_from_reset
:
self
.allocations_from_reset
.load
(
Ordering
::
Relaxed
),
evictions
:
self
.evictions
.load
(
Ordering
::
Relaxed
),
registrations
:
self
.registrations
.load
(
Ordering
::
Relaxed
),
duplicate_blocks
:
self
.duplicate_blocks
.load
(
Ordering
::
Relaxed
),
registration_dedup
:
self
.registration_dedup
.load
(
Ordering
::
Relaxed
),
stagings
:
self
.stagings
.load
(
Ordering
::
Relaxed
),
match_hashes_requested
:
self
.match_hashes_requested
.load
(
Ordering
::
Relaxed
),
match_blocks_returned
:
self
.match_blocks_returned
.load
(
Ordering
::
Relaxed
),
scan_hashes_requested
:
self
.scan_hashes_requested
.load
(
Ordering
::
Relaxed
),
scan_blocks_returned
:
self
.scan_blocks_returned
.load
(
Ordering
::
Relaxed
),
inflight_mutable
:
self
.inflight_mutable
.load
(
Ordering
::
Relaxed
),
inflight_immutable
:
self
.inflight_immutable
.load
(
Ordering
::
Relaxed
),
reset_pool_size
:
self
.reset_pool_size
.load
(
Ordering
::
Relaxed
),
inactive_pool_size
:
self
.inactive_pool_size
.load
(
Ordering
::
Relaxed
),
}
}
}
/// Point-in-time snapshot of all atomic metrics, used by the stats collector and prometheus collector.
#[derive(Debug,
Clone,
Copy)]
pub
struct
MetricsSnapshot
{
pub
allocations
:
u64
,
pub
allocations_from_reset
:
u64
,
pub
evictions
:
u64
,
pub
registrations
:
u64
,
pub
duplicate_blocks
:
u64
,
pub
registration_dedup
:
u64
,
pub
stagings
:
u64
,
pub
match_hashes_requested
:
u64
,
pub
match_blocks_returned
:
u64
,
pub
scan_hashes_requested
:
u64
,
pub
scan_blocks_returned
:
u64
,
pub
inflight_mutable
:
i64
,
pub
inflight_immutable
:
i64
,
pub
reset_pool_size
:
i64
,
pub
inactive_pool_size
:
i64
,
}
#[cfg(test)]
mod
tests
{
use
super
::
*
;
#[test]
fn
test_counter_increments
()
{
let
m
=
BlockPoolMetrics
::
new
(
"G1"
.to_string
());
m
.inc_allocations
(
5
);
m
.inc_allocations
(
3
);
m
.inc_evictions
(
2
);
m
.inc_registrations
();
m
.inc_duplicate_blocks
();
m
.inc_registration_dedup
();
m
.inc_stagings
();
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.allocations
,
8
);
assert_eq!
(
snap
.evictions
,
2
);
assert_eq!
(
snap
.registrations
,
1
);
assert_eq!
(
snap
.duplicate_blocks
,
1
);
assert_eq!
(
snap
.registration_dedup
,
1
);
assert_eq!
(
snap
.stagings
,
1
);
}
#[test]
fn
test_gauge_bidirectional
()
{
let
m
=
BlockPoolMetrics
::
new
(
"G2"
.to_string
());
m
.inc_inflight_mutable
();
m
.inc_inflight_mutable
();
m
.dec_inflight_mutable
();
m
.inc_inflight_immutable
();
m
.inc_inflight_immutable
();
m
.inc_inflight_immutable
();
m
.dec_inflight_immutable
();
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inflight_mutable
,
1
);
assert_eq!
(
snap
.inflight_immutable
,
2
);
}
#[test]
fn
test_pool_size_gauges
()
{
let
m
=
BlockPoolMetrics
::
new
(
"G1"
.to_string
());
m
.set_reset_pool_size
(
100
);
m
.set_inactive_pool_size
(
50
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.reset_pool_size
,
100
);
assert_eq!
(
snap
.inactive_pool_size
,
50
);
m
.set_reset_pool_size
(
80
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.reset_pool_size
,
80
);
// Test inc/dec for reset pool size
m
.inc_reset_pool_size
();
m
.inc_reset_pool_size
();
m
.dec_reset_pool_size
();
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.reset_pool_size
,
81
);
// Test inc/dec for inactive pool size
m
.inc_inactive_pool_size
();
m
.inc_inactive_pool_size
();
m
.inc_inactive_pool_size
();
m
.dec_inactive_pool_size
();
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.inactive_pool_size
,
52
);
}
#[test]
fn
test_type_label
()
{
let
m
=
BlockPoolMetrics
::
new
(
"MyPool"
.to_string
());
assert_eq!
(
m
.type_label
(),
"MyPool"
);
}
#[test]
fn
test_match_scan_counters
()
{
let
m
=
BlockPoolMetrics
::
new
(
"G1"
.to_string
());
m
.inc_match_hashes_requested
(
10
);
m
.inc_match_blocks_returned
(
7
);
m
.inc_scan_hashes_requested
(
20
);
m
.inc_scan_blocks_returned
(
15
);
let
snap
=
m
.snapshot
();
assert_eq!
(
snap
.match_hashes_requested
,
10
);
assert_eq!
(
snap
.match_blocks_returned
,
7
);
assert_eq!
(
snap
.scan_hashes_requested
,
20
);
assert_eq!
(
snap
.scan_blocks_returned
,
15
);
}
}
lib/kvbm-logical/src/metrics/stats.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Toggleable periodic sampler for computing rates, gradients, and hit ratios.
//!
//! The raw counters/gauges in [`BlockPoolMetrics`] are always active (zero-cost atomics).
//! This `StatsCollector` is the optional layer that computes derived statistics.
use
std
::
collections
::
VecDeque
;
use
std
::
sync
::
atomic
::{
AtomicBool
,
Ordering
};
use
std
::
sync
::{
Arc
,
RwLock
};
use
std
::
time
::{
Duration
,
Instant
};
use
super
::
pool_metrics
::{
BlockPoolMetrics
,
MetricsSnapshot
};
/// Configuration for the stats collector.
#[derive(Debug,
Clone)]
pub
struct
StatsConfig
{
/// Maximum number of snapshots to retain in the ring buffer.
pub
window_size
:
usize
,
/// Interval between periodic samples when using `spawn()`.
pub
sample_interval
:
Duration
,
}
impl
Default
for
StatsConfig
{
fn
default
()
->
Self
{
Self
{
window_size
:
60
,
sample_interval
:
Duration
::
from_secs
(
1
),
}
}
}
/// Derived statistics computed from counter deltas over time.
#[derive(Debug,
Clone,
Copy)]
pub
struct
StatsSnapshot
{
/// Allocations per second.
pub
allocation_rate
:
f64
,
/// Evictions per second.
pub
eviction_rate
:
f64
,
/// Ratio of blocks returned to hashes requested in match_blocks.
pub
match_hit_rate
:
f64
,
/// Ratio of blocks returned to hashes requested in scan_matches.
pub
scan_hit_rate
:
f64
,
/// Rate of change of allocation_rate (d(alloc_rate)/dt).
pub
allocation_gradient
:
f64
,
/// Rate of change of eviction_rate (d(eviction_rate)/dt).
pub
eviction_gradient
:
f64
,
/// Current inflight mutable blocks.
pub
inflight_mutable
:
i64
,
/// Current inflight immutable blocks.
pub
inflight_immutable
:
i64
,
/// Current reset pool size.
pub
reset_pool_size
:
i64
,
/// Current inactive pool size.
pub
inactive_pool_size
:
i64
,
}
struct
TimedEntry
{
timestamp
:
Instant
,
raw
:
MetricsSnapshot
,
stats
:
StatsSnapshot
,
}
/// Toggleable periodic sampler that computes rates and gradients from raw metrics.
pub
struct
StatsCollector
{
enabled
:
AtomicBool
,
metrics
:
Arc
<
BlockPoolMetrics
>
,
entries
:
RwLock
<
VecDeque
<
TimedEntry
>>
,
config
:
StatsConfig
,
}
impl
StatsCollector
{
/// Create a new `StatsCollector` for the given metrics source.
/// Disabled by default — call `set_enabled(true)` to activate.
pub
fn
new
(
metrics
:
Arc
<
BlockPoolMetrics
>
,
config
:
StatsConfig
)
->
Self
{
Self
{
enabled
:
AtomicBool
::
new
(
false
),
metrics
,
entries
:
RwLock
::
new
(
VecDeque
::
with_capacity
(
config
.window_size
)),
config
,
}
}
/// Enable or disable the stats collector.
pub
fn
set_enabled
(
&
self
,
enabled
:
bool
)
{
self
.enabled
.store
(
enabled
,
Ordering
::
Relaxed
);
}
/// Whether the stats collector is currently enabled.
pub
fn
is_enabled
(
&
self
)
->
bool
{
self
.enabled
.load
(
Ordering
::
Relaxed
)
}
/// Manually take a sample, computing rates/gradients from the delta since the last snapshot.
/// No-op if disabled.
pub
fn
sample
(
&
self
)
{
if
!
self
.is_enabled
()
{
return
;
}
let
now
=
Instant
::
now
();
let
raw
=
self
.metrics
.snapshot
();
let
mut
entries
=
self
.entries
.write
()
.expect
(
"entries lock poisoned"
);
let
stats
=
if
let
Some
(
prev
)
=
entries
.back
()
{
let
dt
=
now
.duration_since
(
prev
.timestamp
)
.as_secs_f64
();
if
dt
>
0.0
{
let
alloc_rate
=
(
raw
.allocations
-
prev
.raw.allocations
)
as
f64
/
dt
;
let
eviction_rate
=
(
raw
.evictions
-
prev
.raw.evictions
)
as
f64
/
dt
;
let
match_hit_rate
=
{
let
delta_req
=
raw
.match_hashes_requested
-
prev
.raw.match_hashes_requested
;
let
delta_ret
=
raw
.match_blocks_returned
-
prev
.raw.match_blocks_returned
;
if
delta_req
>
0
{
delta_ret
as
f64
/
delta_req
as
f64
}
else
{
0.0
}
};
let
scan_hit_rate
=
{
let
delta_req
=
raw
.scan_hashes_requested
-
prev
.raw.scan_hashes_requested
;
let
delta_ret
=
raw
.scan_blocks_returned
-
prev
.raw.scan_blocks_returned
;
if
delta_req
>
0
{
delta_ret
as
f64
/
delta_req
as
f64
}
else
{
0.0
}
};
let
allocation_gradient
=
(
alloc_rate
-
prev
.stats.allocation_rate
)
/
dt
;
let
eviction_gradient
=
(
eviction_rate
-
prev
.stats.eviction_rate
)
/
dt
;
StatsSnapshot
{
allocation_rate
:
alloc_rate
,
eviction_rate
,
match_hit_rate
,
scan_hit_rate
,
allocation_gradient
,
eviction_gradient
,
inflight_mutable
:
raw
.inflight_mutable
,
inflight_immutable
:
raw
.inflight_immutable
,
reset_pool_size
:
raw
.reset_pool_size
,
inactive_pool_size
:
raw
.inactive_pool_size
,
}
}
else
{
zero_stats
(
&
raw
)
}
}
else
{
zero_stats
(
&
raw
)
};
entries
.push_back
(
TimedEntry
{
timestamp
:
now
,
raw
,
stats
,
});
while
entries
.len
()
>
self
.config.window_size
{
entries
.pop_front
();
}
}
/// Get the latest computed stats snapshot.
pub
fn
latest
(
&
self
)
->
Option
<
StatsSnapshot
>
{
let
entries
=
self
.entries
.read
()
.expect
(
"entries lock poisoned"
);
entries
.back
()
.map
(|
e
|
e
.stats
)
}
/// Get the last `n` stats snapshots (most recent last).
pub
fn
window
(
&
self
,
n
:
usize
)
->
Vec
<
StatsSnapshot
>
{
let
entries
=
self
.entries
.read
()
.expect
(
"entries lock poisoned"
);
entries
.iter
()
.rev
()
.take
(
n
)
.map
(|
e
|
e
.stats
)
.collect
::
<
Vec
<
_
>>
()
.into_iter
()
.rev
()
.collect
()
}
/// Spawn a tokio task that periodically calls `sample()`.
/// Returns a `JoinHandle` that can be used to abort the task.
pub
fn
spawn
(
self
:
&
Arc
<
Self
>
)
->
tokio
::
task
::
JoinHandle
<
()
>
{
let
this
=
Arc
::
clone
(
self
);
let
interval
=
this
.config.sample_interval
;
tokio
::
spawn
(
async
move
{
let
mut
ticker
=
tokio
::
time
::
interval
(
interval
);
loop
{
ticker
.tick
()
.await
;
this
.sample
();
}
})
}
}
fn
zero_stats
(
raw
:
&
MetricsSnapshot
)
->
StatsSnapshot
{
StatsSnapshot
{
allocation_rate
:
0.0
,
eviction_rate
:
0.0
,
match_hit_rate
:
0.0
,
scan_hit_rate
:
0.0
,
allocation_gradient
:
0.0
,
eviction_gradient
:
0.0
,
inflight_mutable
:
raw
.inflight_mutable
,
inflight_immutable
:
raw
.inflight_immutable
,
reset_pool_size
:
raw
.reset_pool_size
,
inactive_pool_size
:
raw
.inactive_pool_size
,
}
}
#[cfg(test)]
mod
tests
{
use
super
::
*
;
#[test]
fn
test_disabled_by_default
()
{
let
metrics
=
Arc
::
new
(
BlockPoolMetrics
::
new
(
"G1"
.to_string
()));
let
stats
=
StatsCollector
::
new
(
metrics
,
StatsConfig
::
default
());
assert
!
(
!
stats
.is_enabled
());
stats
.sample
();
// should be a no-op
assert
!
(
stats
.latest
()
.is_none
());
}
#[test]
fn
test_enable_and_sample
()
{
let
metrics
=
Arc
::
new
(
BlockPoolMetrics
::
new
(
"G1"
.to_string
()));
let
stats
=
StatsCollector
::
new
(
metrics
.clone
(),
StatsConfig
::
default
());
stats
.set_enabled
(
true
);
// First sample — baseline
stats
.sample
();
let
snap
=
stats
.latest
()
.expect
(
"should have a snapshot"
);
assert_eq!
(
snap
.allocation_rate
,
0.0
);
// Simulate some allocations
metrics
.inc_allocations
(
100
);
std
::
thread
::
sleep
(
Duration
::
from_millis
(
50
));
// Second sample — should compute rate
stats
.sample
();
let
snap
=
stats
.latest
()
.expect
(
"should have a snapshot"
);
assert
!
(
snap
.allocation_rate
>
0.0
,
"allocation rate should be > 0"
);
}
#[test]
fn
test_window
()
{
let
metrics
=
Arc
::
new
(
BlockPoolMetrics
::
new
(
"G1"
.to_string
()));
let
config
=
StatsConfig
{
window_size
:
5
,
..
Default
::
default
()
};
let
stats
=
StatsCollector
::
new
(
metrics
.clone
(),
config
);
stats
.set_enabled
(
true
);
for
i
in
0
..
10
{
metrics
.inc_allocations
(
i
+
1
);
stats
.sample
();
std
::
thread
::
sleep
(
Duration
::
from_millis
(
10
));
}
let
window
=
stats
.window
(
3
);
assert_eq!
(
window
.len
(),
3
);
// Ring buffer should have at most 5 entries
let
all
=
stats
.window
(
100
);
assert_eq!
(
all
.len
(),
5
);
}
#[test]
fn
test_hit_rate_computation
()
{
let
metrics
=
Arc
::
new
(
BlockPoolMetrics
::
new
(
"G1"
.to_string
()));
let
stats
=
StatsCollector
::
new
(
metrics
.clone
(),
StatsConfig
::
default
());
stats
.set_enabled
(
true
);
stats
.sample
();
// baseline
std
::
thread
::
sleep
(
Duration
::
from_millis
(
10
));
metrics
.inc_match_hashes_requested
(
10
);
metrics
.inc_match_blocks_returned
(
7
);
stats
.sample
();
let
snap
=
stats
.latest
()
.unwrap
();
assert
!
((
snap
.match_hit_rate
-
0.7
)
.abs
()
<
0.01
);
}
#[test]
fn
test_gradient_computation
()
{
let
metrics
=
Arc
::
new
(
BlockPoolMetrics
::
new
(
"G1"
.to_string
()));
let
stats
=
StatsCollector
::
new
(
metrics
.clone
(),
StatsConfig
::
default
());
stats
.set_enabled
(
true
);
// Sample 1: baseline
stats
.sample
();
std
::
thread
::
sleep
(
Duration
::
from_millis
(
50
));
// Sample 2: some allocations → positive rate
metrics
.inc_allocations
(
100
);
stats
.sample
();
let
snap2
=
stats
.latest
()
.unwrap
();
assert
!
(
snap2
.allocation_rate
>
0.0
);
std
::
thread
::
sleep
(
Duration
::
from_millis
(
50
));
// Sample 3: more allocations → rate changes → gradient
metrics
.inc_allocations
(
500
);
stats
.sample
();
let
snap3
=
stats
.latest
()
.unwrap
();
// Gradient should be non-zero since rate changed
// (exact value depends on timing, but should be positive since rate increased)
assert
!
(
snap3
.allocation_gradient
!=
0.0
||
snap3
.allocation_rate
>
snap2
.allocation_rate
,
"allocation gradient or rate should reflect increasing allocations"
);
}
}
lib/kvbm-logical/src/pools/active.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Active pool for managing blocks that are currently in use (have strong references).
//!
//! This pool provides a layer of abstraction over the BlockRegistry for finding
//! active blocks. Active blocks are those that have been registered and are
//! currently being used, as opposed to inactive blocks which are available
//! for reuse.
use
std
::
sync
::
Arc
;
use
super
::{
BlockMetadata
,
RegisteredBlock
,
SequenceHash
};
use
crate
::
blocks
::
RegisteredReturnFn
;
use
crate
::
registry
::
BlockRegistry
;
/// Pool for managing active (in-use) blocks.
///
/// This is a simple wrapper around BlockRegistry that encapsulates the logic
/// for finding blocks that are currently active (have strong references).
pub
(
crate
)
struct
ActivePool
<
T
:
BlockMetadata
>
{
block_registry
:
BlockRegistry
,
return_fn
:
RegisteredReturnFn
<
T
>
,
}
impl
<
T
:
BlockMetadata
>
ActivePool
<
T
>
{
/// Create a new ActivePool with the given registry and return function.
pub
(
crate
)
fn
new
(
block_registry
:
BlockRegistry
,
return_fn
:
RegisteredReturnFn
<
T
>
)
->
Self
{
Self
{
block_registry
,
return_fn
,
}
}
/// Find multiple blocks by sequence hashes, stopping on first miss.
///
/// This searches for active blocks in the registry and returns them as
/// RegisteredBlock guards. If any hash is not found or the block cannot
/// be retrieved, the search stops and returns only the blocks found so far.
#[inline]
pub
(
crate
)
fn
find_matches
(
&
self
,
hashes
:
&
[
SequenceHash
],
touch
:
bool
,
)
->
Vec
<
Arc
<
dyn
RegisteredBlock
<
T
>>>
{
let
mut
matches
=
Vec
::
with_capacity
(
hashes
.len
());
for
hash
in
hashes
{
if
let
Some
(
handle
)
=
self
.block_registry
.match_sequence_hash
(
*
hash
,
touch
)
{
if
let
Some
(
block
)
=
handle
.try_get_block
::
<
T
>
(
self
.return_fn
.clone
())
{
matches
.push
(
block
);
}
else
{
break
;
// Stop on first miss
}
}
else
{
break
;
// Stop on first miss
}
}
matches
}
/// Scan for blocks in the active pool (doesn't stop on miss).
///
/// Unlike `find_matches`, this continues scanning even when a hash is not found.
/// Returns all found blocks with their corresponding sequence hashes.
#[inline]
pub
(
crate
)
fn
scan_matches
(
&
self
,
hashes
:
&
[
SequenceHash
],
)
->
Vec
<
(
SequenceHash
,
Arc
<
dyn
RegisteredBlock
<
T
>>
)
>
{
hashes
.iter
()
.filter_map
(|
hash
|
{
self
.block_registry
.match_sequence_hash
(
*
hash
,
false
)
.and_then
(|
handle
|
{
handle
.try_get_block
::
<
T
>
(
self
.return_fn
.clone
())
.map
(|
block
|
(
*
hash
,
block
))
})
})
.collect
()
}
// /// Find a single block by sequence hash.
// ///
// /// Returns the block if found and active, None otherwise.
// #[inline]
// pub(crate) fn find_match(&self, seq_hash: SequenceHash) -> Option<Arc<dyn RegisteredBlock<T>>> {
// self.block_registry
// .match_sequence_hash(seq_hash, true)
// .and_then(|handle| handle.try_get_block::<T>(self.return_fn.clone()))
// }
// /// Check if a block with the given sequence hash is currently active.
// pub(crate) fn has_block(&self, seq_hash: SequenceHash) -> bool {
// self.find_match(seq_hash).is_some()
// }
}
#[cfg(test)]
mod
tests
{
use
super
::
*
;
use
crate
::
blocks
::{
Block
,
PrimaryBlock
,
state
::
Reset
};
use
crate
::
pools
::
backends
::{
FifoReusePolicy
,
HashMapBackend
};
use
crate
::
pools
::
inactive
::
InactivePool
;
use
crate
::
pools
::
reset
::
ResetPool
;
use
crate
::
testing
::{
TestMeta
,
create_staged_block
,
tokens_for_id
};
fn
create_test_setup
()
->
(
ActivePool
<
TestMeta
>
,
InactivePool
<
TestMeta
>
,
BlockRegistry
,
ResetPool
<
TestMeta
>
,
)
{
let
registry
=
BlockRegistry
::
new
();
let
reset_blocks
:
Vec
<
Block
<
TestMeta
,
Reset
>>
=
(
0
..
10_u
size
)
.map
(|
i
|
Block
::
new
(
i
,
4
))
.collect
();
let
reset_pool
=
ResetPool
::
new
(
reset_blocks
,
4
,
None
);
let
reuse_policy
=
Box
::
new
(
FifoReusePolicy
::
new
());
let
backend
=
Box
::
new
(
HashMapBackend
::
new
(
reuse_policy
));
let
inactive_pool
=
InactivePool
::
new
(
backend
,
&
reset_pool
,
None
);
let
active_pool
=
ActivePool
::
new
(
registry
.clone
(),
inactive_pool
.return_fn
());
(
active_pool
,
inactive_pool
,
registry
,
reset_pool
)
}
/// Register a staged block and hold a strong reference to make it "active".
fn
make_active_block
(
registry
:
&
BlockRegistry
,
return_fn
:
&
RegisteredReturnFn
<
TestMeta
>
,
id
:
usize
,
tokens
:
&
[
u32
],
)
->
(
Arc
<
PrimaryBlock
<
TestMeta
>>
,
SequenceHash
)
{
let
staged
=
create_staged_block
::
<
TestMeta
>
(
id
,
tokens
);
let
seq_hash
=
staged
.sequence_hash
();
let
handle
=
registry
.register_sequence_hash
(
seq_hash
);
let
registered
=
staged
.register_with_handle
(
handle
);
let
primary
=
PrimaryBlock
::
new_attached
(
Arc
::
new
(
registered
),
return_fn
.clone
());
(
primary
,
seq_hash
)
}
#[test]
fn
test_find_matches
()
{
let
(
active_pool
,
inactive_pool
,
registry
,
_
reset_pool
)
=
create_test_setup
();
let
return_fn
=
inactive_pool
.return_fn
();
let
(
_
hold1
,
hash1
)
=
make_active_block
(
&
registry
,
&
return_fn
,
1
,
&
tokens_for_id
(
1
));
let
(
_
hold2
,
hash2
)
=
make_active_block
(
&
registry
,
&
return_fn
,
2
,
&
tokens_for_id
(
2
));
let
(
_
hold3
,
hash3
)
=
make_active_block
(
&
registry
,
&
return_fn
,
3
,
&
tokens_for_id
(
3
));
let
found
=
active_pool
.find_matches
(
&
[
hash1
,
hash2
,
hash3
],
true
);
assert_eq!
(
found
.len
(),
3
);
assert_eq!
(
found
[
0
]
.block_id
(),
1
);
assert_eq!
(
found
[
1
]
.block_id
(),
2
);
assert_eq!
(
found
[
2
]
.block_id
(),
3
);
}
#[test]
fn
test_find_matches_stops_on_miss
()
{
let
(
active_pool
,
inactive_pool
,
registry
,
_
reset_pool
)
=
create_test_setup
();
let
return_fn
=
inactive_pool
.return_fn
();
let
(
_
hold1
,
hash1
)
=
make_active_block
(
&
registry
,
&
return_fn
,
1
,
&
tokens_for_id
(
1
));
let
(
_
hold3
,
hash3
)
=
make_active_block
(
&
registry
,
&
return_fn
,
3
,
&
tokens_for_id
(
3
));
// Create a hash that's not in the registry
let
missing_hash
=
{
let
staged
=
create_staged_block
::
<
TestMeta
>
(
999
,
&
[
9999
,
9998
,
9997
,
9996
]);
staged
.sequence_hash
()
};
let
found
=
active_pool
.find_matches
(
&
[
hash1
,
missing_hash
,
hash3
],
true
);
assert_eq!
(
found
.len
(),
1
);
assert_eq!
(
found
[
0
]
.block_id
(),
1
);
}
#[test]
fn
test_scan_matches
()
{
let
(
active_pool
,
inactive_pool
,
registry
,
_
reset_pool
)
=
create_test_setup
();
let
return_fn
=
inactive_pool
.return_fn
();
let
(
_
hold1
,
hash1
)
=
make_active_block
(
&
registry
,
&
return_fn
,
1
,
&
tokens_for_id
(
1
));
let
(
_
hold3
,
hash3
)
=
make_active_block
(
&
registry
,
&
return_fn
,
3
,
&
tokens_for_id
(
3
));
let
missing_hash
=
{
let
staged
=
create_staged_block
::
<
TestMeta
>
(
999
,
&
[
9999
,
9998
,
9997
,
9996
]);
staged
.sequence_hash
()
};
// scan_matches doesn't stop on miss — should find both 1 and 3
let
found
=
active_pool
.scan_matches
(
&
[
hash1
,
missing_hash
,
hash3
]);
assert_eq!
(
found
.len
(),
2
);
assert_eq!
(
found
[
0
]
.0
,
hash1
);
assert_eq!
(
found
[
0
]
.1
.block_id
(),
1
);
assert_eq!
(
found
[
1
]
.0
,
hash3
);
assert_eq!
(
found
[
1
]
.1
.block_id
(),
3
);
}
#[test]
fn
test_find_matches_empty
()
{
let
(
active_pool
,
_
inactive_pool
,
_
registry
,
_
reset_pool
)
=
create_test_setup
();
let
found
=
active_pool
.find_matches
(
&
[],
true
);
assert
!
(
found
.is_empty
());
}
#[test]
fn
test_find_matches_no_active_blocks
()
{
let
(
active_pool
,
_
inactive_pool
,
_
registry
,
_
reset_pool
)
=
create_test_setup
();
let
missing_hash
=
{
let
staged
=
create_staged_block
::
<
TestMeta
>
(
999
,
&
[
9999
,
9998
,
9997
,
9996
]);
staged
.sequence_hash
()
};
let
found
=
active_pool
.find_matches
(
&
[
missing_hash
],
true
);
assert
!
(
found
.is_empty
());
}
}
lib/kvbm-logical/src/pools/block_proptest.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Property-based tests for block state machine.
//!
//! These tests verify critical invariants of the Block<T, State> type-state pattern
//! using property-based testing with proptest. The focus is on the state transitions
//! and block size validation logic.
use
super
::{
super
::
blocks
::
*
,
tests
::
*
,
*
};
#[cfg(test)]
pub
(
crate
)
mod
tests
{
use
super
::
*
;
use
crate
::
testing
::
config
::{
COMMON_TEST_BLOCK_SIZES
,
constants
,
generate_test_tokens
,
validate_test_block_size
,
};
use
dynamo_tokens
::{
TokenBlock
,
TokenBlockSequence
};
use
proptest
::
prelude
::
*
;
/// Helper function to create a TokenBlock from a token sequence
fn
create_token_block_from_sequence
(
tokens
:
&
[
u32
])
->
Option
<
TokenBlock
>
{
if
tokens
.is_empty
()
{
return
None
;
}
let
sequence
=
TokenBlockSequence
::
from_slice
(
tokens
,
tokens
.len
()
as
u32
,
Some
(
42
));
// If we have a complete block, return it
if
let
Some
(
block
)
=
sequence
.blocks
()
.first
()
{
return
Some
(
block
.clone
());
}
// Otherwise try to commit the partial block
let
(
_
,
mut
partial
)
=
sequence
.into_parts
();
partial
.commit
()
.ok
()
}
proptest!
{
/// Property: Block size validation is consistent across all operations
///
/// This test verifies that:
/// 1. Blocks can only be completed when token count matches block size
/// 2. Block size is preserved across all state transitions
/// 3. Error handling returns the original block to prevent leaks
#[test]
fn
prop_block_size_validation_consistency
(
block_size
in
prop
::
sample
::
select
(
COMMON_TEST_BLOCK_SIZES
),
token_count
in
1u
size
..=
128u
size
,
block_id
in
any
::
<
BlockId
>
(),
)
{
prop_assume!
(
validate_test_block_size
(
block_size
));
prop_assume!
(
token_count
<=
128
);
// Keep token sequences reasonable
let
block
=
Block
::
<
TestData
,
Reset
>
::
new
(
block_id
,
block_size
);
prop_assert_eq!
(
block
.block_size
(),
block_size
);
prop_assert_eq!
(
block
.block_id
(),
block_id
);
let
tokens
=
generate_test_tokens
(
100
,
token_count
);
// Only test if we can create a valid token block
if
let
Some
(
token_block
)
=
create_token_block_from_sequence
(
&
tokens
)
{
let
actual_token_size
=
token_block
.block_size
();
let
result
=
block
.complete
(
&
token_block
);
if
actual_token_size
==
block_size
{
// Should succeed when sizes match
let
complete_block
=
result
.expect
(
"Should complete when sizes match"
);
prop_assert_eq!
(
complete_block
.block_size
(),
block_size
);
prop_assert_eq!
(
complete_block
.block_id
(),
block_id
);
}
else
{
// Should fail when sizes don't match
prop_assert!
(
matches!
(
result
,
Err
(
BlockError
::
BlockSizeMismatch
{
..
})),
"Expected BlockSizeMismatch error when sizes don't match"
);
// Verify error contains the original block
if
let
Err
(
BlockError
::
BlockSizeMismatch
{
expected
,
actual
,
block
:
returned_block
})
=
result
{
prop_assert_eq!
(
expected
,
block_size
);
prop_assert_eq!
(
actual
,
actual_token_size
);
prop_assert_eq!
(
returned_block
.block_size
(),
block_size
);
prop_assert_eq!
(
returned_block
.block_id
(),
block_id
);
}
}
}
}
/// Property: Complete state transitions preserve all block properties
///
/// This test verifies that going from Reset → Complete → Registered → Reset
/// preserves the block ID and block size throughout all transitions.
#[test]
fn
prop_state_transitions_preserve_properties
(
block_size
in
prop
::
sample
::
select
(
&
[
1u
size
,
4
,
16
,
64
]),
block_id
in
any
::
<
BlockId
>
(),
base_token
in
0u32
..
1000u32
,
)
{
prop_assume!
(
validate_test_block_size
(
block_size
));
// Start with Reset block
let
reset_block
=
Block
::
<
TestData
,
Reset
>
::
new
(
block_id
,
block_size
);
prop_assert_eq!
(
reset_block
.block_size
(),
block_size
);
prop_assert_eq!
(
reset_block
.block_id
(),
block_id
);
// Generate matching tokens and create token block
let
tokens
=
generate_test_tokens
(
base_token
,
block_size
);
let
token_block
=
create_token_block_from_sequence
(
&
tokens
)
.expect
(
"Should be able to create token block"
);
// Transition to Complete
let
complete_block
=
reset_block
.complete
(
&
token_block
)
.expect
(
"Should complete with matching size"
);
prop_assert_eq!
(
complete_block
.block_size
(),
block_size
);
prop_assert_eq!
(
complete_block
.block_id
(),
block_id
);
// Transition to Registered
let
registry
=
BlockRegistry
::
new
();
let
seq_hash
=
complete_block
.sequence_hash
();
let
handle
=
registry
.register_sequence_hash
(
seq_hash
);
let
registered_block
=
complete_block
.register_with_handle
(
handle
);
prop_assert_eq!
(
registered_block
.block_size
(),
block_size
);
prop_assert_eq!
(
registered_block
.block_id
(),
block_id
);
prop_assert_eq!
(
registered_block
.sequence_hash
(),
seq_hash
);
// Transition back to Reset
let
reset_again
=
registered_block
.reset
();
prop_assert_eq!
(
reset_again
.block_size
(),
block_size
);
prop_assert_eq!
(
reset_again
.block_id
(),
block_id
);
}
/// Property: Block IDs are preserved but can be arbitrary values
///
/// This test verifies that block IDs are not constrained by the block manager
/// and are preserved through all operations.
#[test]
fn
prop_block_id_preservation
(
block_id
in
any
::
<
BlockId
>
(),
block_size
in
prop
::
sample
::
select
(
&
[
constants
::
SMALL
,
constants
::
MEDIUM
]),
)
{
let
block
=
Block
::
<
TestData
,
Reset
>
::
new
(
block_id
,
block_size
);
prop_assert_eq!
(
block
.block_id
(),
block_id
);
// Test with edge case block IDs
let
edge_cases
=
[
0
,
1
,
BlockId
::
MAX
/
2
,
BlockId
::
MAX
-
1
,
BlockId
::
MAX
];
for
&
test_id
in
&
edge_cases
{
let
edge_block
=
Block
::
<
TestData
,
Reset
>
::
new
(
test_id
,
block_size
);
prop_assert_eq!
(
edge_block
.block_id
(),
test_id
);
}
}
// /// Property: Error handling prevents resource leaks
// ///
// /// This test verifies that when block completion fails, the original block
// /// is returned in the error, preventing resource leaks.
// #[test]
// fn prop_error_handling_prevents_leaks(
// block_size in prop::sample::select(&[4usize, 8, 16]),
// wrong_token_count in prop::sample::select(&[1usize, 2, 3, 7, 9, 15, 17, 32]),
// block_id in 0u64..1000u64,
// ) {
// prop_assume!(validate_test_block_size(block_size));
// prop_assume!(wrong_token_count != block_size); // Ensure mismatch
// prop_assume!(wrong_token_count <= 32); // Keep reasonable
// let original_block = Block::<TestData, Reset>::new(block_id, block_size);
// let tokens = generate_test_tokens(500, wrong_token_count);
// if let Some(token_block) = create_token_block_from_sequence(&tokens) {
// let result = original_block.complete(token_block);
// // Should always fail due to size mismatch
// prop_assert!(result.is_err());
// if let Err(BlockError::BlockSizeMismatch { expected, actual, block: returned_block }) = result {
// // Verify error details
// prop_assert_eq!(expected, block_size);
// prop_assert_ne!(actual, block_size); // Should be different
// // Verify returned block is equivalent to original
// prop_assert_eq!(returned_block.block_id(), block_id);
// prop_assert_eq!(returned_block.block_size(), block_size);
// // We can still use the returned block
// let correct_tokens = generate_test_tokens(600, block_size);
// if let Some(correct_token_block) = create_token_block_from_sequence(&correct_tokens) {
// let success_result = returned_block.complete(correct_token_block);
// prop_assert!(success_result.is_ok());
// }
// }
// }
// }
/// Property: Block size constraints are enforced at construction
///
/// This test verifies that all valid block sizes work correctly,
/// providing coverage across the full range of allowed values.
#[test]
fn
prop_valid_block_sizes_work
(
block_size
in
prop
::
sample
::
select
(
&
[
1u
size
,
2
,
4
,
8
,
16
,
32
,
64
,
128
,
256
,
512
,
1024
]),
block_id
in
any
::
<
BlockId
>
(),
)
{
// All these sizes should be valid
prop_assert!
(
validate_test_block_size
(
block_size
));
// Should be able to create blocks with any valid size
let
block
=
Block
::
<
TestData
,
Reset
>
::
new
(
block_id
,
block_size
);
prop_assert_eq!
(
block
.block_size
(),
block_size
);
// Should be able to complete with matching token count
let
tokens
=
generate_test_tokens
(
0
,
block_size
);
if
let
Some
(
token_block
)
=
create_token_block_from_sequence
(
&
tokens
)
{
let
result
=
block
.complete
(
&
token_block
);
prop_assert!
(
result
.is_ok
(),
"Failed to complete block with size {}"
,
block_size
);
}
}
}
/// Additional focused property tests for specific scenarios
mod
focused_properties
{
use
super
::
*
;
proptest!
{
/// Property: Sequence hash is deterministic for identical token sequences
#[test]
fn
prop_sequence_hash_deterministic
(
tokens
in
prop
::
collection
::
vec
(
any
::
<
u32
>
(),
4
..=
4
),
// Always 4 tokens
block_id1
in
any
::
<
BlockId
>
(),
block_id2
in
any
::
<
BlockId
>
(),
)
{
prop_assume!
(
block_id1
!=
block_id2
);
// Different block IDs
if
let
Some
(
token_block1
)
=
create_token_block_from_sequence
(
&
tokens
)
&&
let
Some
(
token_block2
)
=
create_token_block_from_sequence
(
&
tokens
)
{
// Same tokens should produce same sequence hash regardless of block ID
let
block1
=
Block
::
<
TestData
,
Reset
>
::
new
(
block_id1
,
4
);
let
block2
=
Block
::
<
TestData
,
Reset
>
::
new
(
block_id2
,
4
);
let
complete1
=
block1
.complete
(
&
token_block1
)
.expect
(
"Should complete"
);
let
complete2
=
block2
.complete
(
&
token_block2
)
.expect
(
"Should complete"
);
prop_assert_eq!
(
complete1
.sequence_hash
(),
complete2
.sequence_hash
());
}
}
/// Property: Different token sequences produce different hashes (with high probability)
#[test]
fn
prop_different_tokens_different_hashes
(
tokens1
in
prop
::
collection
::
vec
(
0u32
..
100u32
,
4
..=
4
),
tokens2
in
prop
::
collection
::
vec
(
100u32
..
200u32
,
4
..=
4
),
// Different range
)
{
if
let
(
Some
(
token_block1
),
Some
(
token_block2
))
=
(
create_token_block_from_sequence
(
&
tokens1
),
create_token_block_from_sequence
(
&
tokens2
)
)
{
let
block1
=
Block
::
<
TestData
,
Reset
>
::
new
(
1
,
4
);
let
block2
=
Block
::
<
TestData
,
Reset
>
::
new
(
2
,
4
);
let
complete1
=
block1
.complete
(
&
token_block1
)
.expect
(
"Should complete"
);
let
complete2
=
block2
.complete
(
&
token_block2
)
.expect
(
"Should complete"
);
// Different token sequences should produce different hashes
// (This is probabilistic but very likely with our token ranges)
if
tokens1
!=
tokens2
{
prop_assert_ne!
(
complete1
.sequence_hash
(),
complete2
.sequence_hash
());
}
}
}
}
}
}
lib/kvbm-logical/src/pools/inactive/backends/fifo.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! FIFO reuse policy for inactive registered blocks.
//!
//! Allocates blocks in first-in-first-out order using a monotonic counter.
//! Uses BTreeMap for O(log n) insertion/removal with priority key ordering.
#![allow(dead_code)]
use
super
::
reuse_policy
::
*
;
use
std
::
collections
::{
BTreeMap
,
HashMap
};
use
super
::{
BlockId
,
InactiveBlock
};
/// Monotonic sequence number used as priority key for FIFO ordering.
pub
type
PriorityKey
=
u64
;
/// FIFO reuse policy
#[derive(Debug)]
pub
struct
FifoReusePolicy
{
keys
:
HashMap
<
BlockId
,
PriorityKey
>
,
blocks
:
BTreeMap
<
PriorityKey
,
InactiveBlock
>
,
next_seq
:
u64
,
}
impl
Default
for
FifoReusePolicy
{
fn
default
()
->
Self
{
Self
::
new
()
}
}
impl
FifoReusePolicy
{
pub
fn
new
()
->
Self
{
Self
{
keys
:
HashMap
::
new
(),
blocks
:
BTreeMap
::
new
(),
next_seq
:
0
,
}
}
}
impl
ReusePolicy
for
FifoReusePolicy
{
fn
insert
(
&
mut
self
,
inactive_block
:
InactiveBlock
)
->
Result
<
(),
ReusePolicyError
>
{
assert
!
(
!
self
.keys
.contains_key
(
&
inactive_block
.block_id
),
"block already exists"
);
let
priority_key
=
self
.next_seq
;
self
.next_seq
+=
1
;
self
.keys
.insert
(
inactive_block
.block_id
,
priority_key
);
self
.blocks
.insert
(
priority_key
,
inactive_block
);
Ok
(())
}
fn
remove
(
&
mut
self
,
block_id
:
BlockId
)
->
Result
<
(),
ReusePolicyError
>
{
let
priority_key
=
self
.keys
.remove
(
&
block_id
)
.ok_or
(
ReusePolicyError
::
BlockNotFound
(
block_id
))
?
;
assert
!
(
self
.blocks
.remove
(
&
priority_key
)
.is_some
(),
"block not found"
);
Ok
(())
}
fn
next_free
(
&
mut
self
)
->
Option
<
InactiveBlock
>
{
let
next_block
=
self
.blocks
.pop_first
();
if
let
Some
((
_
,
block
))
=
next_block
{
assert
!
(
self
.keys
.remove
(
&
block
.block_id
)
.is_some
(),
"block not found"
);
Some
(
block
)
}
else
{
None
}
}
fn
is_empty
(
&
self
)
->
bool
{
self
.blocks
.is_empty
()
}
fn
len
(
&
self
)
->
usize
{
self
.blocks
.len
()
}
}
#[cfg(test)]
mod
tests
{
use
super
::
*
;
use
crate
::
testing
::{
TestMeta
,
create_staged_block
};
// Use TestMeta instead of local TestData
type
TestData
=
TestMeta
;
// Wrapper for backward compatibility
fn
create_completed_block
<
T
:
crate
::
blocks
::
BlockMetadata
+
std
::
fmt
::
Debug
>
(
tokens
:
&
[
u32
],
block_id
:
BlockId
,
)
->
crate
::
blocks
::
Block
<
T
,
crate
::
blocks
::
state
::
Staged
>
{
create_staged_block
::
<
T
>
(
block_id
,
tokens
)
}
/// Helper function to create InactiveBlock instances for testing
fn
create_inactive_block
(
block_id
:
BlockId
,
seq_hash
:
u64
)
->
InactiveBlock
{
let
complete_block
=
create_completed_block
::
<
TestData
>
(
&
[
seq_hash
as
u32
],
block_id
);
InactiveBlock
{
block_id
,
seq_hash
:
complete_block
.sequence_hash
(),
}
}
#[test]
fn
test_fifo_ordering_basic
()
{
let
mut
policy
=
FifoReusePolicy
::
new
();
// Insert blocks
let
block1
=
create_inactive_block
(
1
,
100
);
let
block2
=
create_inactive_block
(
2
,
200
);
let
block3
=
create_inactive_block
(
3
,
300
);
let
seq_hash1
=
block1
.seq_hash
;
let
seq_hash2
=
block2
.seq_hash
;
let
seq_hash3
=
block3
.seq_hash
;
policy
.insert
(
block1
)
.unwrap
();
policy
.insert
(
block2
)
.unwrap
();
policy
.insert
(
block3
)
.unwrap
();
// Verify FIFO order - first inserted should come out first
assert_eq!
(
policy
.len
(),
3
);
assert
!
(
!
policy
.is_empty
());
let
retrieved1
=
policy
.next_free
()
.unwrap
();
assert_eq!
(
retrieved1
.block_id
,
1
);
assert_eq!
(
retrieved1
.seq_hash
,
seq_hash1
);
let
retrieved2
=
policy
.next_free
()
.unwrap
();
assert_eq!
(
retrieved2
.block_id
,
2
);
assert_eq!
(
retrieved2
.seq_hash
,
seq_hash2
);
let
retrieved3
=
policy
.next_free
()
.unwrap
();
assert_eq!
(
retrieved3
.block_id
,
3
);
assert_eq!
(
retrieved3
.seq_hash
,
seq_hash3
);
assert
!
(
policy
.is_empty
());
assert_eq!
(
policy
.len
(),
0
);
}
#[test]
fn
test_fifo_ordering_with_delays
()
{
let
mut
policy
=
FifoReusePolicy
::
new
();
// Insert blocks
let
blocks
=
vec!
[
create_inactive_block
(
10
,
1000
),
create_inactive_block
(
20
,
2000
),
create_inactive_block
(
30
,
3000
),
create_inactive_block
(
40
,
4000
),
];
for
block
in
blocks
{
policy
.insert
(
block
)
.unwrap
();
}
// Retrieve all blocks and verify FIFO order
let
expected_order
=
vec!
[
10
,
20
,
30
,
40
];
let
mut
retrieved_order
=
Vec
::
new
();
while
let
Some
(
block
)
=
policy
.next_free
()
{
retrieved_order
.push
(
block
.block_id
);
}
assert_eq!
(
retrieved_order
,
expected_order
);
}
#[test]
fn
test_insert_and_remove
()
{
let
mut
policy
=
FifoReusePolicy
::
new
();
// Insert several blocks
let
blocks
=
vec!
[
create_inactive_block
(
1
,
100
),
create_inactive_block
(
2
,
200
),
create_inactive_block
(
3
,
300
),
create_inactive_block
(
4
,
400
),
];
for
block
in
blocks
{
policy
.insert
(
block
)
.unwrap
();
}
assert_eq!
(
policy
.len
(),
4
);
// Remove block 2 (second inserted)
policy
.remove
(
2
)
.unwrap
();
assert_eq!
(
policy
.len
(),
3
);
// Retrieve remaining blocks - should be 1, 3, 4 in that order
let
retrieved1
=
policy
.next_free
()
.unwrap
();
assert_eq!
(
retrieved1
.block_id
,
1
);
let
retrieved2
=
policy
.next_free
()
.unwrap
();
assert_eq!
(
retrieved2
.block_id
,
3
);
let
retrieved3
=
policy
.next_free
()
.unwrap
();
assert_eq!
(
retrieved3
.block_id
,
4
);
assert
!
(
policy
.is_empty
());
}
#[test]
fn
test_empty_operations
()
{
let
mut
policy
=
FifoReusePolicy
::
new
();
// Test empty state
assert
!
(
policy
.is_empty
());
assert_eq!
(
policy
.len
(),
0
);
assert
!
(
policy
.next_free
()
.is_none
());
// Insert and remove a block
let
block
=
create_inactive_block
(
1
,
100
);
policy
.insert
(
block
)
.unwrap
();
assert
!
(
!
policy
.is_empty
());
assert_eq!
(
policy
.len
(),
1
);
let
retrieved
=
policy
.next_free
()
.unwrap
();
assert_eq!
(
retrieved
.block_id
,
1
);
// Should be empty again
assert
!
(
policy
.is_empty
());
assert_eq!
(
policy
.len
(),
0
);
assert
!
(
policy
.next_free
()
.is_none
());
}
#[test]
#[should_panic(expected
=
"block already exists"
)]
fn
test_duplicate_block_panic
()
{
let
mut
policy
=
FifoReusePolicy
::
new
();
let
block
=
create_inactive_block
(
1
,
100
);
policy
.insert
(
block
)
.unwrap
();
// Inserting the same block ID again should panic
let
duplicate_block
=
create_inactive_block
(
1
,
200
);
// Same ID, different hash
policy
.insert
(
duplicate_block
)
.unwrap
();
}
#[test]
fn
test_remove_nonexistent_block
()
{
let
mut
policy
=
FifoReusePolicy
::
new
();
// Try to remove from empty policy
let
result
=
policy
.remove
(
999
);
assert
!
(
matches!
(
result
,
Err
(
ReusePolicyError
::
BlockNotFound
(
_
))));
// Insert a block and try to remove a different one
let
block
=
create_inactive_block
(
1
,
100
);
policy
.insert
(
block
)
.unwrap
();
let
result
=
policy
.remove
(
999
);
assert
!
(
matches!
(
result
,
Err
(
ReusePolicyError
::
BlockNotFound
(
_
))));
// Verify the original block is still there
assert_eq!
(
policy
.len
(),
1
);
let
retrieved
=
policy
.next_free
()
.unwrap
();
assert_eq!
(
retrieved
.block_id
,
1
);
}
#[test]
fn
test_interleaved_operations
()
{
let
mut
policy
=
FifoReusePolicy
::
new
();
// Insert some blocks
policy
.insert
(
create_inactive_block
(
1
,
100
))
.unwrap
();
policy
.insert
(
create_inactive_block
(
2
,
200
))
.unwrap
();
policy
.insert
(
create_inactive_block
(
3
,
300
))
.unwrap
();
// Remove the first one
let
first
=
policy
.next_free
()
.unwrap
();
assert_eq!
(
first
.block_id
,
1
);
// Insert another block
policy
.insert
(
create_inactive_block
(
4
,
400
))
.unwrap
();
// Remove a specific block by ID
policy
.remove
(
3
)
.unwrap
();
// Insert another block
policy
.insert
(
create_inactive_block
(
5
,
500
))
.unwrap
();
// The remaining blocks should come out in order: 2, 4, 5
let
second
=
policy
.next_free
()
.unwrap
();
assert_eq!
(
second
.block_id
,
2
);
let
third
=
policy
.next_free
()
.unwrap
();
assert_eq!
(
third
.block_id
,
4
);
let
fourth
=
policy
.next_free
()
.unwrap
();
assert_eq!
(
fourth
.block_id
,
5
);
assert
!
(
policy
.is_empty
());
}
#[test]
fn
test_priority_key_ordering
()
{
let
mut
policy
=
FifoReusePolicy
::
new
();
// Insert blocks — monotonic counter guarantees unique, ordered keys
for
i
in
1
..=
5
{
policy
.insert
(
create_inactive_block
(
i
,
i
as
u64
*
100
))
.unwrap
();
}
// Verify each key is unique and strictly increasing
let
keys
:
Vec
<
PriorityKey
>
=
(
1
..=
5
)
.map
(|
id
|
*
policy
.keys
.get
(
&
id
)
.unwrap
())
.collect
();
for
window
in
keys
.windows
(
2
)
{
assert
!
(
window
[
0
]
<
window
[
1
],
"keys must be strictly increasing"
);
}
// Retrieve all blocks and verify they come out in insertion order
let
mut
retrieval_order
:
Vec
<
BlockId
>
=
Vec
::
new
();
while
let
Some
(
block
)
=
policy
.next_free
()
{
retrieval_order
.push
(
block
.block_id
);
}
let
expected_order
:
Vec
<
BlockId
>
=
(
1
..=
5
)
.collect
();
assert_eq!
(
retrieval_order
,
expected_order
);
}
#[test]
fn
test_btreemap_ordering_assumption
()
{
use
std
::
collections
::
BTreeMap
;
// Verify our assumption about BTreeMap ordering with u64 keys
let
mut
map
=
BTreeMap
::
new
();
// Insert keys in non-sorted order
map
.insert
(
100u64
,
"hundred"
);
map
.insert
(
10u64
,
"ten"
);
map
.insert
(
50u64
,
"fifty"
);
map
.insert
(
1u64
,
"one"
);
map
.insert
(
200u64
,
"two_hundred"
);
// pop_first should return the smallest key first
assert_eq!
(
map
.pop_first
(),
Some
((
1
,
"one"
)));
assert_eq!
(
map
.pop_first
(),
Some
((
10
,
"ten"
)));
assert_eq!
(
map
.pop_first
(),
Some
((
50
,
"fifty"
)));
assert_eq!
(
map
.pop_first
(),
Some
((
100
,
"hundred"
)));
assert_eq!
(
map
.pop_first
(),
Some
((
200
,
"two_hundred"
)));
assert_eq!
(
map
.pop_first
(),
None
);
}
}
lib/kvbm-logical/src/pools/inactive/backends/hashmap_backend.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Backend trait for InactivePool storage strategies.
use
std
::
collections
::
HashMap
;
use
super
::{
Block
,
BlockMetadata
,
InactiveBlock
,
Registered
,
SequenceHash
};
use
super
::
super
::
InactivePoolBackend
;
use
super
::
ReusePolicy
;
pub
struct
HashMapBackend
<
T
:
BlockMetadata
>
{
blocks
:
HashMap
<
SequenceHash
,
Block
<
T
,
Registered
>>
,
reuse_policy
:
Box
<
dyn
ReusePolicy
>
,
}
impl
<
T
:
BlockMetadata
>
HashMapBackend
<
T
>
{
pub
fn
new
(
reuse_policy
:
Box
<
dyn
ReusePolicy
>
)
->
Self
{
Self
{
blocks
:
HashMap
::
new
(),
reuse_policy
,
}
}
}
impl
<
T
:
BlockMetadata
>
InactivePoolBackend
<
T
>
for
HashMapBackend
<
T
>
{
fn
find_matches
(
&
mut
self
,
hashes
:
&
[
SequenceHash
],
_
touch
:
bool
)
->
Vec
<
Block
<
T
,
Registered
>>
{
let
mut
matches
=
Vec
::
with_capacity
(
hashes
.len
());
for
hash
in
hashes
{
if
let
Some
(
block
)
=
self
.blocks
.remove
(
hash
)
{
let
_
=
self
.reuse_policy
.remove
(
block
.block_id
());
matches
.push
(
block
);
}
else
{
break
;
}
}
matches
}
fn
scan_matches
(
&
mut
self
,
hashes
:
&
[
SequenceHash
],
_
touch
:
bool
,
)
->
Vec
<
(
SequenceHash
,
Block
<
T
,
Registered
>
)
>
{
let
mut
matches
=
Vec
::
new
();
for
hash
in
hashes
{
if
let
Some
(
block
)
=
self
.blocks
.remove
(
hash
)
{
let
_
=
self
.reuse_policy
.remove
(
block
.block_id
());
matches
.push
((
*
hash
,
block
));
}
// Unlike find_matches: NO break on miss - continue scanning
}
matches
}
fn
allocate
(
&
mut
self
,
count
:
usize
)
->
Vec
<
Block
<
T
,
Registered
>>
{
let
mut
allocated
=
Vec
::
with_capacity
(
count
);
for
_
in
0
..
count
{
if
let
Some
(
InactiveBlock
{
seq_hash
,
..
})
=
self
.reuse_policy
.next_free
()
{
if
let
Some
(
block
)
=
self
.blocks
.remove
(
&
seq_hash
)
{
allocated
.push
(
block
);
}
else
{
debug_assert!
(
false
,
"reuse_policy yielded seq_hash {:?} not found in blocks"
,
seq_hash
);
}
}
else
{
break
;
}
}
allocated
}
fn
insert
(
&
mut
self
,
block
:
Block
<
T
,
Registered
>
)
{
let
seq_hash
=
block
.sequence_hash
();
let
_
=
self
.reuse_policy
.insert
(
InactiveBlock
{
block_id
:
block
.block_id
(),
seq_hash
,
});
self
.blocks
.insert
(
seq_hash
,
block
);
}
fn
len
(
&
self
)
->
usize
{
self
.blocks
.len
()
}
fn
has_block
(
&
self
,
seq_hash
:
SequenceHash
)
->
bool
{
self
.blocks
.contains_key
(
&
seq_hash
)
}
fn
allocate_all
(
&
mut
self
)
->
Vec
<
Block
<
T
,
Registered
>>
{
// Drain reuse policy by consuming all entries
while
self
.reuse_policy
.next_free
()
.is_some
()
{}
// Drain and return all blocks
self
.blocks
.drain
()
.map
(|(
_
,
block
)|
block
)
.collect
()
}
}
lib/kvbm-logical/src/pools/inactive/backends/lineage.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
use
std
::
collections
::{
BTreeMap
,
HashMap
,
HashSet
};
use
dynamo_tokens
::
PositionalLineageHash
;
use
super
::
super
::{
Block
,
BlockMetadata
,
InactivePoolBackend
,
Registered
};
/// The data stored in a lineage node - either a real block or a ghost placeholder.
enum
LineageNodeData
<
T
:
BlockMetadata
>
{
/// A real block with timestamp.
Real
{
block
:
Block
<
T
,
Registered
>
,
last_used
:
u64
,
},
/// A ghost node created for out-of-order insertions.
Ghost
,
}
/// A node in the lineage graph.
struct
LineageNode
<
T
:
BlockMetadata
>
{
/// The data stored in this node (real block or ghost).
data
:
LineageNodeData
<
T
>
,
/// The parent fragment (at position - 1), if any.
parent_fragment
:
Option
<
u64
>
,
/// Children fragments (at position + 1).
children
:
HashSet
<
u64
>
,
}
impl
<
T
:
BlockMetadata
>
LineageNode
<
T
>
{
fn
new
(
block
:
Block
<
T
,
Registered
>
,
lineage_hash
:
PositionalLineageHash
,
tick
:
u64
)
->
Self
{
let
parent_fragment
=
if
lineage_hash
.position
()
>
0
{
Some
(
lineage_hash
.parent_hash_fragment
())
}
else
{
None
};
Self
{
data
:
LineageNodeData
::
Real
{
block
,
last_used
:
tick
,
},
parent_fragment
,
children
:
HashSet
::
new
(),
}
}
fn
is_leaf
(
&
self
)
->
bool
{
self
.children
.is_empty
()
}
}
/// A backend that manages blocks using a lineage graph and evicts from the leaves.
pub
struct
LineageBackend
<
T
:
BlockMetadata
>
{
/// Map from (position, fragment) to Node.
nodes
:
HashMap
<
u64
,
HashMap
<
u64
,
LineageNode
<
T
>>>
,
/// Sorted queue of leaf nodes, keyed by (last_used, position, fragment).
/// Smallest key (oldest tick) is popped first.
leaf_queue
:
BTreeMap
<
(
u64
,
u64
,
u64
),
()
>
,
/// Total number of blocks currently stored (excluding ghost nodes).
count
:
usize
,
/// Monotonic counter for insertion ordering.
current_tick
:
u64
,
}
impl
<
T
:
BlockMetadata
>
Default
for
LineageBackend
<
T
>
{
fn
default
()
->
Self
{
Self
::
new
()
}
}
impl
<
T
:
BlockMetadata
>
LineageBackend
<
T
>
{
/// Creates a new LineageBackend.
pub
fn
new
()
->
Self
{
Self
{
nodes
:
HashMap
::
new
(),
leaf_queue
:
BTreeMap
::
new
(),
count
:
0
,
current_tick
:
0
,
}
}
/// Inserts a block into the lineage graph.
/// Panics on collision or duplicate insertion.
pub
fn
insert
(
&
mut
self
,
block
:
Block
<
T
,
Registered
>
)
{
let
lineage_hash
=
block
.sequence_hash
();
let
position
=
lineage_hash
.position
();
let
fragment
=
lineage_hash
.current_hash_fragment
();
let
full_hash
=
lineage_hash
.as_u128
();
let
parent_fragment
=
if
position
>
0
{
Some
(
lineage_hash
.parent_hash_fragment
())
}
else
{
None
};
let
increment_count
:
bool
;
let
tick
=
self
.current_tick
;
self
.current_tick
+=
1
;
// 1. Create or update the node
let
level
=
self
.nodes
.entry
(
position
)
.or_default
();
match
level
.entry
(
fragment
)
{
std
::
collections
::
hash_map
::
Entry
::
Vacant
(
e
)
=>
{
increment_count
=
true
;
let
node
=
LineageNode
::
new
(
block
,
lineage_hash
,
tick
);
e
.insert
(
node
);
}
std
::
collections
::
hash_map
::
Entry
::
Occupied
(
mut
e
)
=>
{
let
node
=
e
.get_mut
();
match
&
node
.data
{
LineageNodeData
::
Ghost
=>
{
// Fill ghost with real block data
increment_count
=
true
;
node
.data
=
LineageNodeData
::
Real
{
block
,
last_used
:
tick
,
};
node
.parent_fragment
=
parent_fragment
;
}
LineageNodeData
::
Real
{
block
:
existing_block
,
..
}
=>
{
let
existing_hash
=
existing_block
.sequence_hash
()
.as_u128
();
if
existing_hash
==
full_hash
{
panic!
(
"Duplicate insertion detected! position={}, fragment={:#x}, hash={:#032x}.
\
The same block was inserted twice."
,
position
,
fragment
,
full_hash
);
}
else
{
panic!
(
"Hash collision detected! position={}, fragment={:#x},
\
existing_hash={:#032x}, new_hash={:#032x}.
\
Different blocks mapped to same position+fragment."
,
position
,
fragment
,
existing_hash
,
full_hash
);
}
}
}
}
}
if
increment_count
{
self
.count
+=
1
;
}
// 2. Link to parent
if
let
Some
(
p_frag
)
=
parent_fragment
{
let
p_pos
=
position
-
1
;
let
parent_level
=
self
.nodes
.entry
(
p_pos
)
.or_default
();
let
parent_node
=
parent_level
.entry
(
p_frag
)
.or_insert_with
(||
{
LineageNode
{
data
:
LineageNodeData
::
Ghost
,
parent_fragment
:
None
,
// We don't know the parent's parent yet
children
:
HashSet
::
new
(),
}
});
let
was_parent_leaf
=
parent_node
.is_leaf
();
parent_node
.children
.insert
(
fragment
);
if
was_parent_leaf
{
// Parent was a leaf, now has a child. Remove from queue.
// Note: Ghost nodes are never in queue.
if
let
LineageNodeData
::
Real
{
last_used
,
..
}
=
parent_node
.data
{
self
.leaf_queue
.remove
(
&
(
last_used
,
p_pos
,
p_frag
));
}
}
}
// 3. Update LRU status for this node
let
node
=
self
.nodes
.get
(
&
position
)
.unwrap
()
.get
(
&
fragment
)
.unwrap
();
if
node
.is_leaf
()
&&
let
LineageNodeData
::
Real
{
last_used
,
..
}
=
node
.data
{
self
.leaf_queue
.insert
((
last_used
,
position
,
fragment
),
());
}
}
/// Allocates (removes) a block from the pool, preferring leaves in LRU order.
pub
fn
allocate
(
&
mut
self
,
count
:
usize
)
->
Vec
<
Block
<
T
,
Registered
>>
{
let
mut
allocated
=
Vec
::
with_capacity
(
count
);
while
allocated
.len
()
<
count
{
if
let
Some
((
&
(
_
tick
,
pos
,
frag
),
_
))
=
self
.leaf_queue
.iter
()
.next
()
{
// Need to remove from map using the key we just found
let
key
=
(
_
tick
,
pos
,
frag
);
self
.leaf_queue
.remove
(
&
key
);
if
let
Some
(
b
)
=
self
.remove_block
(
pos
,
frag
)
{
allocated
.push
(
b
);
}
}
else
{
break
;
// No more leaves
}
}
allocated
}
/// Removes a specific block by its lineage hash (for cache hits).
pub
fn
remove
(
&
mut
self
,
lineage_hash
:
&
PositionalLineageHash
)
->
Option
<
Block
<
T
,
Registered
>>
{
let
position
=
lineage_hash
.position
();
let
fragment
=
lineage_hash
.current_hash_fragment
();
let
node_data
=
self
.nodes
.get
(
&
position
)
.and_then
(|
level
|
level
.get
(
&
fragment
))
.and_then
(|
node
|
match
&
node
.data
{
LineageNodeData
::
Real
{
last_used
,
..
}
=>
Some
(
*
last_used
),
LineageNodeData
::
Ghost
=>
None
,
});
if
let
Some
(
tick
)
=
node_data
{
// Remove from queue if present (might be present if it's a leaf)
self
.leaf_queue
.remove
(
&
(
tick
,
position
,
fragment
));
self
.remove_block
(
position
,
fragment
)
}
else
{
None
}
}
/// Internal method to remove a block from the graph.
/// Returns the block if one existed at that node.
/// Handles ghost cleanup iteratively.
fn
remove_block
(
&
mut
self
,
position
:
u64
,
fragment
:
u64
)
->
Option
<
Block
<
T
,
Registered
>>
{
let
node_block
=
{
let
level
=
self
.nodes
.get_mut
(
&
position
)
?
;
let
node
=
level
.get_mut
(
&
fragment
)
?
;
match
&
mut
node
.data
{
LineageNodeData
::
Real
{
..
}
=>
{
// Replace Real with Ghost, taking ownership of the block
let
block_val
=
std
::
mem
::
replace
(
&
mut
node
.data
,
LineageNodeData
::
Ghost
);
if
let
LineageNodeData
::
Real
{
block
,
..
}
=
block_val
{
Some
(
block
)
}
else
{
unreachable!
()
}
}
LineageNodeData
::
Ghost
=>
None
,
}
};
if
node_block
.is_some
()
{
self
.count
-=
1
;
}
let
mut
current_pos
=
position
;
let
mut
current_frag
=
fragment
;
// Loop for iterative cleanup upwards
loop
{
let
mut
should_remove_node
=
false
;
let
mut
parent_info
=
None
;
if
let
Some
(
level
)
=
self
.nodes
.get
(
&
current_pos
)
&&
let
Some
(
node
)
=
level
.get
(
&
current_frag
)
{
let
is_ghost
=
matches!
(
node
.data
,
LineageNodeData
::
Ghost
);
if
node
.children
.is_empty
()
&&
is_ghost
{
// It's a ghost leaf (no block, no children). Prune it.
should_remove_node
=
true
;
parent_info
=
node
.parent_fragment
.map
(|
pf
|
(
current_pos
.saturating_sub
(
1
),
pf
));
}
}
if
should_remove_node
{
if
let
Some
(
level
)
=
self
.nodes
.get_mut
(
&
current_pos
)
{
level
.remove
(
&
current_frag
);
if
level
.is_empty
()
{
self
.nodes
.remove
(
&
current_pos
);
}
}
if
let
Some
((
p_pos
,
p_frag
))
=
parent_info
{
let
mut
parent_became_leaf
=
false
;
let
mut
parent_has_block
=
false
;
let
mut
parent_tick
=
0
;
if
let
Some
(
level
)
=
self
.nodes
.get_mut
(
&
p_pos
)
&&
let
Some
(
parent
)
=
level
.get_mut
(
&
p_frag
)
{
parent
.children
.remove
(
&
current_frag
);
if
parent
.children
.is_empty
()
{
parent_became_leaf
=
true
;
match
&
parent
.data
{
LineageNodeData
::
Real
{
last_used
,
..
}
=>
{
parent_has_block
=
true
;
parent_tick
=
*
last_used
;
}
LineageNodeData
::
Ghost
=>
{
parent_has_block
=
false
;
}
}
}
}
if
parent_became_leaf
{
if
parent_has_block
{
// Parent is a real block leaf -> add to queue using its OLD tick
self
.leaf_queue
.insert
((
parent_tick
,
p_pos
,
p_frag
),
());
break
;
}
else
{
current_pos
=
p_pos
;
current_frag
=
p_frag
;
continue
;
}
}
else
{
break
;
}
}
else
{
break
;
}
}
else
{
break
;
}
}
node_block
}
}
impl
<
T
:
BlockMetadata
>
InactivePoolBackend
<
T
>
for
LineageBackend
<
T
>
{
fn
find_matches
(
&
mut
self
,
hashes
:
&
[
PositionalLineageHash
],
_
touch
:
bool
,
)
->
Vec
<
Block
<
T
,
Registered
>>
{
let
mut
matches
=
Vec
::
with_capacity
(
hashes
.len
());
for
hash
in
hashes
{
if
let
Some
(
block
)
=
self
.remove
(
hash
)
{
matches
.push
(
block
);
}
else
{
break
;
// Stop on first miss
}
}
matches
}
fn
scan_matches
(
&
mut
self
,
hashes
:
&
[
PositionalLineageHash
],
_
touch
:
bool
,
)
->
Vec
<
(
PositionalLineageHash
,
Block
<
T
,
Registered
>
)
>
{
let
mut
matches
=
Vec
::
new
();
for
hash
in
hashes
{
if
let
Some
(
block
)
=
self
.remove
(
hash
)
{
matches
.push
((
*
hash
,
block
));
}
// Unlike find_matches: NO break on miss - continue scanning
}
matches
}
fn
allocate
(
&
mut
self
,
count
:
usize
)
->
Vec
<
Block
<
T
,
Registered
>>
{
// Delegate to the inherent method
LineageBackend
::
allocate
(
self
,
count
)
}
fn
insert
(
&
mut
self
,
block
:
Block
<
T
,
Registered
>
)
{
// Delegate to the inherent method
LineageBackend
::
insert
(
self
,
block
)
}
fn
len
(
&
self
)
->
usize
{
self
.count
}
fn
has_block
(
&
self
,
seq_hash
:
PositionalLineageHash
)
->
bool
{
let
position
=
seq_hash
.position
();
let
fragment
=
seq_hash
.current_hash_fragment
();
self
.nodes
.get
(
&
position
)
.and_then
(|
level
|
level
.get
(
&
fragment
))
.is_some_and
(|
node
|
matches!
(
node
.data
,
LineageNodeData
::
Real
{
..
}))
}
}
#[cfg(test)]
mod
tests
{
use
super
::
*
;
use
crate
::
SequenceHash
;
use
crate
::
blocks
::
Block
;
use
crate
::
pools
::
tests
::
TestData
;
use
crate
::
pools
::
tests
::
fixtures
::
BlockSequenceBuilder
;
impl
<
T
:
BlockMetadata
>
LineageBackend
<
T
>
{
/// Test helper: get the number of entries in the leaf queue.
pub
fn
get_queue_len
(
&
self
)
->
usize
{
self
.leaf_queue
.len
()
}
}
// Helper to create test blocks with proper lineage using BlockSequenceBuilder
// Returns a vector of (Block, SequenceHash) tuples
// offset: starting token value (use different offsets for independent chains)
fn
create_blocks_with_offset
(
count
:
usize
,
offset
:
u32
,
)
->
Vec
<
(
Block
<
TestData
,
Registered
>
,
SequenceHash
)
>
{
let
tokens
:
Vec
<
u32
>
=
(
offset
..
offset
+
count
as
u32
)
.collect
();
BlockSequenceBuilder
::
from_tokens
(
tokens
)
.with_block_size
(
1
)
.build
()
}
// Helper to create test blocks starting from token 0
fn
create_blocks
(
count
:
usize
)
->
Vec
<
(
Block
<
TestData
,
Registered
>
,
SequenceHash
)
>
{
create_blocks_with_offset
(
count
,
0
)
}
// Helper for single block creation (root block at position 0)
fn
create_block
(
id
:
u32
)
->
(
Block
<
TestData
,
Registered
>
,
SequenceHash
)
{
let
tokens
=
vec!
[
id
];
let
blocks
=
BlockSequenceBuilder
::
from_tokens
(
tokens
)
.with_block_size
(
1
)
.build
();
blocks
.into_iter
()
.next
()
.unwrap
()
}
#[test]
fn
test_leaf_insertion
()
{
let
mut
backend
=
LineageBackend
::
<
TestData
>
::
new
();
let
(
b1
,
_
)
=
create_block
(
1
);
backend
.insert
(
b1
);
assert_eq!
(
backend
.len
(),
1
);
assert_eq!
(
backend
.get_queue_len
(),
1
);
// It is a leaf (no children)
let
allocated
=
backend
.allocate
(
1
);
assert_eq!
(
allocated
.len
(),
1
);
assert_eq!
(
allocated
[
0
]
.block_id
(),
0
);
// Block ID is 0 (first block in sequence)
assert_eq!
(
backend
.len
(),
0
);
}
#[test]
fn
test_parent_child_insertion
()
{
let
mut
backend
=
LineageBackend
::
<
TestData
>
::
new
();
// Create a sequence of 2 blocks with proper parent-child relationship
let
mut
blocks
=
create_blocks
(
2
);
let
(
b1
,
_
)
=
blocks
.remove
(
0
);
// Parent at position 0
let
(
b2
,
_
)
=
blocks
.remove
(
0
);
// Child at position 1
// Insert parent first
backend
.insert
(
b1
);
assert_eq!
(
backend
.get_queue_len
(),
1
);
// b1 is leaf
// Insert child
backend
.insert
(
b2
);
assert_eq!
(
backend
.len
(),
2
);
// b1 is no longer leaf (has child b2). b2 is leaf.
// LRU should contain only b2.
assert_eq!
(
backend
.get_queue_len
(),
1
);
let
allocated
=
backend
.allocate
(
1
);
assert_eq!
(
allocated
.len
(),
1
);
assert_eq!
(
allocated
[
0
]
.block_id
(),
1
);
// Should allocate b2 (leaf, block_id=1)
// Now b1 should be a leaf again and added to LRU
assert_eq!
(
backend
.get_queue_len
(),
1
);
let
allocated2
=
backend
.allocate
(
1
);
assert_eq!
(
allocated2
.len
(),
1
);
assert_eq!
(
allocated2
[
0
]
.block_id
(),
0
);
// b1 has block_id=0
}
#[test]
fn
test_out_of_order_insertion
()
{
let
mut
backend
=
LineageBackend
::
<
TestData
>
::
new
();
// Insert child first (from blocks2)
let
mut
blocks2_mut
=
create_blocks
(
2
);
backend
.insert
(
blocks2_mut
.remove
(
1
)
.0
);
// Created ghost node for parent b1.
// b2 is leaf.
assert_eq!
(
backend
.len
(),
1
);
// Only 1 actual block
assert_eq!
(
backend
.get_queue_len
(),
1
);
// Insert parent (from blocks1)
let
mut
blocks1_mut
=
create_blocks
(
2
);
backend
.insert
(
blocks1_mut
.remove
(
0
)
.0
);
// Parent b1 fills ghost. It has child b2, so it's NOT a leaf.
// b2 is still leaf.
assert_eq!
(
backend
.len
(),
2
);
assert_eq!
(
backend
.get_queue_len
(),
1
);
// Only b2
let
allocated
=
backend
.allocate
(
1
);
assert_eq!
(
allocated
[
0
]
.block_id
(),
1
);
// b2
// Now b1 becomes leaf
assert_eq!
(
backend
.get_queue_len
(),
1
);
let
allocated2
=
backend
.allocate
(
1
);
assert_eq!
(
allocated2
[
0
]
.block_id
(),
0
);
// b1
}
#[test]
fn
test_branching
()
{
let
mut
backend
=
LineageBackend
::
<
TestData
>
::
new
();
// Test that multiple independent chains can coexist and be allocated independently
let
seq1
=
create_blocks_with_offset
(
3
,
0
);
// chain1: 0 -> 1 -> 2
let
seq2
=
create_blocks_with_offset
(
3
,
5000
);
// chain2: 0 -> 1 -> 2
// Insert all blocks from both chains
for
(
block
,
_
)
in
seq1
{
backend
.insert
(
block
);
}
for
(
block
,
_
)
in
seq2
{
backend
.insert
(
block
);
}
// Should have 6 blocks total
assert_eq!
(
backend
.len
(),
6
);
// Leaves are position 2 from each chain (2 leaves)
assert_eq!
(
backend
.get_queue_len
(),
2
);
// Allocate one leaf
let
alloc1
=
backend
.allocate
(
1
);
assert_eq!
(
alloc1
.len
(),
1
);
assert_eq!
(
backend
.len
(),
5
);
// Now position 1 from one chain should be a leaf, plus position 2 from the other chain
assert_eq!
(
backend
.get_queue_len
(),
2
);
}
#[test]
fn
test_interleaved_chains
()
{
// Chain 1: A(0) -> B(1)
// Chain 2: X(0) -> Y(1)
// We want strict consumption based on insertion order (ticks).
let
mut
backend
=
LineageBackend
::
<
TestData
>
::
new
();
let
mut
chain1
=
create_blocks_with_offset
(
2
,
0
);
let
(
a
,
_
)
=
chain1
.remove
(
0
);
let
(
b
,
_
)
=
chain1
.remove
(
0
);
let
mut
chain2
=
create_blocks_with_offset
(
2
,
1000
);
let
(
x
,
_
)
=
chain2
.remove
(
0
);
let
(
y
,
_
)
=
chain2
.remove
(
0
);
// Insert in order: A, B, X, Y
// insert(A) tick 0
// insert(B) tick 1
// insert(X) tick 2
// insert(Y) tick 3
// So Chain 1 is older.
backend
.insert
(
a
);
backend
.insert
(
b
);
backend
.insert
(
x
);
backend
.insert
(
y
);
assert_eq!
(
backend
.len
(),
4
);
assert_eq!
(
backend
.get_queue_len
(),
2
);
// Leaves: B, Y
// B (tick 1) is older than Y (tick 3). Expect B.
let
alloc1
=
backend
.allocate
(
1
);
assert_eq!
(
alloc1
[
0
]
.block_id
(),
1
);
// B (block_id 1 from chain1)
// Now A becomes leaf. A has tick 0.
// Queue: A(0), Y(3).
// Expect A.
let
alloc2
=
backend
.allocate
(
1
);
assert_eq!
(
alloc2
[
0
]
.block_id
(),
0
);
// A (block_id 0 from chain1)
// Now Y(3).
let
alloc3
=
backend
.allocate
(
1
);
assert_eq!
(
alloc3
[
0
]
.block_id
(),
1
);
// Y (block_id 1 from chain2)
// Now X becomes leaf. X has tick 2.
let
alloc4
=
backend
.allocate
(
1
);
assert_eq!
(
alloc4
[
0
]
.block_id
(),
0
);
// X (block_id 0 from chain2)
}
#[test]
fn
test_remove_by_hash
()
{
let
mut
backend
=
LineageBackend
::
<
TestData
>
::
new
();
let
(
b1
,
seq_hash
)
=
create_block
(
1
);
backend
.insert
(
b1
);
assert_eq!
(
backend
.len
(),
1
);
let
removed
=
backend
.remove
(
&
seq_hash
);
assert
!
(
removed
.is_some
());
assert_eq!
(
removed
.unwrap
()
.block_id
(),
0
);
assert_eq!
(
backend
.len
(),
0
);
}
#[test]
fn
test_deep_chain_cleanup_iterative
()
{
// Create deep chain: 0 -> 1 -> 2 ... -> 999
let
depth
=
1000
;
let
mut
backend
=
LineageBackend
::
<
TestData
>
::
new
();
// Create a deep chain of blocks
let
blocks
=
create_blocks
(
depth
);
let
last_hash
=
blocks
[
depth
-
1
]
.1
;
for
(
block
,
_
)
in
blocks
{
backend
.insert
(
block
);
}
assert_eq!
(
backend
.len
(),
depth
);
// Only last one is leaf
assert_eq!
(
backend
.get_queue_len
(),
1
);
backend
.remove
(
&
last_hash
);
assert_eq!
(
backend
.len
(),
depth
-
1
);
// Now depth-2 is leaf
assert_eq!
(
backend
.get_queue_len
(),
1
);
// Test out-of-order insertion to create ghosts
backend
=
LineageBackend
::
<
TestData
>
::
new
();
// Create a chain and insert only the leaf at position 100
let
mut
chain
=
create_blocks
(
101
);
// 0..100
let
(
b_leaf
,
h_leaf
)
=
chain
.remove
(
100
);
// Insert leaf at depth 100. This creates a ghost parent at position 99.
backend
.insert
(
b_leaf
);
assert_eq!
(
backend
.len
(),
1
);
// Only 1 real block
// Ghost nodes exist but are not counted in len
// Remove leaf. This should clean up the ghost at position 99.
backend
.remove
(
&
h_leaf
);
assert_eq!
(
backend
.len
(),
0
);
assert
!
(
backend
.nodes
.is_empty
());
}
#[test]
fn
test_split_sequence_eviction
()
{
// Test eviction ordering with two independent chains
// Branch 1: A(0)->B(1)->C(2)->D(3)->E(4)
// Branch 2: X(0)->Y(1)->Z(2)->W(3)->V(4)
let
mut
backend
=
LineageBackend
::
<
TestData
>
::
new
();
// Create two separate 5-block chains with different tokens
let
mut
branch1
=
create_blocks_with_offset
(
5
,
0
);
let
mut
branch2
=
create_blocks_with_offset
(
5
,
3000
);
// Insert all 5 blocks from branch1
for
_
i
in
0
..
5
{
backend
.insert
(
branch1
.remove
(
0
)
.0
);
}
// Insert all 5 blocks from branch2
for
_
i
in
0
..
5
{
backend
.insert
(
branch2
.remove
(
0
)
.0
);
}
assert_eq!
(
backend
.len
(),
10
);
// Leaves are E(4) from branch1 and V(4) from branch2
assert_eq!
(
backend
.get_queue_len
(),
2
);
// Allocate first leaf (oldest)
let
alloc1
=
backend
.allocate
(
1
);
assert_eq!
(
alloc1
.len
(),
1
);
assert_eq!
(
backend
.len
(),
9
);
// Allocate second leaf
let
alloc2
=
backend
.allocate
(
1
);
assert_eq!
(
alloc2
.len
(),
1
);
assert_eq!
(
backend
.len
(),
8
);
// Now D(3) from branch1 and W(3) from branch2 should be leaves
assert_eq!
(
backend
.get_queue_len
(),
2
);
// Allocate both
backend
.allocate
(
2
);
assert_eq!
(
backend
.len
(),
6
);
// C(2) from both chains should now be leaves
assert_eq!
(
backend
.get_queue_len
(),
2
);
}
#[test]
fn
test_duplicate_insertion_would_panic
()
{
// Note: This test documents that duplicate insertions would be detected.
// We cannot easily test this because Block<T, Registered> does not implement Clone,
// and the first insert() consumes the block, making it impossible to insert twice.
//
// The duplicate detection logic in insert() checks if a node already exists at
// position+fragment with the same full_hash:
// - If the node exists and is Real with matching full_hash, it panics with
// "Duplicate insertion detected!"
//
// This is the expected behavior: attempting to insert the same block twice would
// panic if we could somehow obtain a second copy of the block with identical hash.
let
mut
backend
=
LineageBackend
::
<
TestData
>
::
new
();
let
(
b1
,
_
)
=
create_block
(
1
);
backend
.insert
(
b1
);
assert_eq!
(
backend
.len
(),
1
);
// Any future insert of a block with matching position+fragment+full_hash would
// trigger the duplicate panic. Since Block doesn't implement Clone and is consumed
// on insert, this test serves as documentation of the expected behavior.
}
#[test]
fn
test_collision_would_be_detected
()
{
// Note: This test documents that hash collisions would be detected.
// We cannot easily create a real collision (two different u128 values with
// the same position+fragment) without constructing invalid PositionalLineageHash
// values directly, which would bypass the normal construction logic.
//
// The collision detection logic in insert() compares full_hash values:
// - If position+fragment match but full_hash differs, it panics with
// "Hash collision detected!"
//
// This is tested implicitly by ensuring that all insertions with the same
// position+fragment must have identical full hashes, otherwise they panic.
let
mut
backend
=
LineageBackend
::
<
TestData
>
::
new
();
let
(
b1
,
_
)
=
create_block
(
1
);
backend
.insert
(
b1
);
assert_eq!
(
backend
.len
(),
1
);
// Any future insert with matching position+fragment but different full_hash
// would trigger the collision panic. Since we can't construct such a case
// without bypassing PositionalLineageHash invariants, this test serves as
// documentation.
}
}
lib/kvbm-logical/src/pools/inactive/backends/lru_backend.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
use
std
::
num
::
NonZeroUsize
;
use
lru
::
LruCache
;
use
super
::{
Block
,
BlockMetadata
,
Registered
,
SequenceHash
};
use
super
::
super
::
InactivePoolBackend
;
pub
struct
LruBackend
<
T
:
BlockMetadata
>
{
cache
:
LruCache
<
SequenceHash
,
Block
<
T
,
Registered
>>
,
}
impl
<
T
:
BlockMetadata
>
LruBackend
<
T
>
{
pub
fn
new
(
capacity
:
NonZeroUsize
)
->
Self
{
Self
{
cache
:
LruCache
::
new
(
capacity
),
}
}
}
impl
<
T
:
BlockMetadata
>
InactivePoolBackend
<
T
>
for
LruBackend
<
T
>
{
fn
find_matches
(
&
mut
self
,
hashes
:
&
[
SequenceHash
],
_
touch
:
bool
)
->
Vec
<
Block
<
T
,
Registered
>>
{
let
mut
matches
=
Vec
::
with_capacity
(
hashes
.len
());
for
hash
in
hashes
{
if
let
Some
(
block
)
=
self
.cache
.pop
(
hash
)
{
matches
.push
(
block
);
}
else
{
break
;
}
}
matches
}
fn
scan_matches
(
&
mut
self
,
hashes
:
&
[
SequenceHash
],
_
touch
:
bool
,
)
->
Vec
<
(
SequenceHash
,
Block
<
T
,
Registered
>
)
>
{
let
mut
matches
=
Vec
::
new
();
for
hash
in
hashes
{
if
let
Some
(
block
)
=
self
.cache
.pop
(
hash
)
{
matches
.push
((
*
hash
,
block
));
}
// Unlike find_matches: NO break on miss - continue scanning
}
matches
}
fn
allocate
(
&
mut
self
,
count
:
usize
)
->
Vec
<
Block
<
T
,
Registered
>>
{
let
mut
allocated
=
Vec
::
with_capacity
(
count
);
for
_
in
0
..
count
{
if
let
Some
((
_
seq_hash
,
block
))
=
self
.cache
.pop_lru
()
{
allocated
.push
(
block
);
}
else
{
break
;
}
}
allocated
}
fn
insert
(
&
mut
self
,
block
:
Block
<
T
,
Registered
>
)
{
let
seq_hash
=
block
.sequence_hash
();
// Assert we're not causing an eviction
assert
!
(
self
.cache
.len
()
<
self
.cache
.cap
()
.get
(),
"LRU backend insert would cause eviction! len={}, cap={}.
\
This indicates insufficient capacity for all blocks."
,
self
.cache
.len
(),
self
.cache
.cap
()
.get
()
);
self
.cache
.put
(
seq_hash
,
block
);
}
fn
len
(
&
self
)
->
usize
{
self
.cache
.len
()
}
fn
has_block
(
&
self
,
seq_hash
:
SequenceHash
)
->
bool
{
self
.cache
.peek
(
&
seq_hash
)
.is_some
()
}
fn
allocate_all
(
&
mut
self
)
->
Vec
<
Block
<
T
,
Registered
>>
{
let
mut
allocated
=
Vec
::
with_capacity
(
self
.cache
.len
());
while
let
Some
((
_
seq_hash
,
block
))
=
self
.cache
.pop_lru
()
{
allocated
.push
(
block
);
}
allocated
}
}
#[cfg(test)]
mod
tests
{
use
super
::
*
;
use
crate
::
pools
::
tests
::
fixtures
::
*
;
#[test]
fn
test_lru_eviction_order
()
{
let
mut
backend
=
LruBackend
::
new
(
NonZeroUsize
::
new
(
3
)
.unwrap
());
let
(
block1
,
hash1
)
=
create_registered_block
(
1
,
&
tokens_for_id
(
1
));
let
(
block2
,
hash2
)
=
create_registered_block
(
2
,
&
tokens_for_id
(
2
));
let
(
block3
,
hash3
)
=
create_registered_block
(
3
,
&
tokens_for_id
(
3
));
backend
.insert
(
block1
);
backend
.insert
(
block2
);
backend
.insert
(
block3
);
assert_eq!
(
backend
.len
(),
3
);
let
allocated
=
backend
.allocate
(
1
);
assert_eq!
(
allocated
.len
(),
1
);
assert_eq!
(
allocated
[
0
]
.block_id
(),
1
);
assert
!
(
!
backend
.has_block
(
hash1
));
assert
!
(
backend
.has_block
(
hash2
));
assert
!
(
backend
.has_block
(
hash3
));
}
#[test]
fn
test_lru_peek_doesnt_affect_order
()
{
let
mut
backend
=
LruBackend
::
new
(
NonZeroUsize
::
new
(
3
)
.unwrap
());
let
(
block1
,
hash1
)
=
create_registered_block
(
1
,
&
tokens_for_id
(
1
));
let
(
block2
,
hash2
)
=
create_registered_block
(
2
,
&
tokens_for_id
(
2
));
backend
.insert
(
block1
);
backend
.insert
(
block2
);
assert_eq!
(
backend
.len
(),
2
);
// Peek at block1 (should not affect LRU order)
assert
!
(
backend
.has_block
(
hash1
));
assert
!
(
backend
.has_block
(
hash2
));
// Allocate blocks - should still follow insertion order (block1 first, then block2)
// despite the peek at block1
let
allocated
=
backend
.allocate
(
2
);
assert_eq!
(
allocated
.len
(),
2
);
assert_eq!
(
allocated
[
0
]
.block_id
(),
1
);
// block1 allocated first (oldest)
assert_eq!
(
allocated
[
1
]
.block_id
(),
2
);
// block2 allocated second
assert_eq!
(
backend
.len
(),
0
);
}
#[test]
fn
test_lru_allocate_more_than_available
()
{
let
mut
backend
=
LruBackend
::
new
(
NonZeroUsize
::
new
(
10
)
.unwrap
());
let
(
block1
,
_
)
=
create_registered_block
(
1
,
&
tokens_for_id
(
1
));
let
(
block2
,
_
)
=
create_registered_block
(
2
,
&
tokens_for_id
(
2
));
backend
.insert
(
block1
);
backend
.insert
(
block2
);
let
allocated
=
backend
.allocate
(
5
);
assert_eq!
(
allocated
.len
(),
2
);
assert_eq!
(
backend
.len
(),
0
);
}
}
lib/kvbm-logical/src/pools/inactive/backends/mod.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Backend storage strategies for InactivePool.
use
super
::
*
;
mod
fifo
;
mod
hashmap_backend
;
mod
lineage
;
mod
lru_backend
;
mod
multi_lru_backend
;
mod
reuse_policy
;
#[cfg(test)]
mod
tests
;
#[allow(unused_imports)]
pub
use
fifo
::
FifoReusePolicy
;
pub
(
crate
)
use
hashmap_backend
::
HashMapBackend
;
pub
(
crate
)
use
lineage
::
LineageBackend
;
pub
(
crate
)
use
lru_backend
::
LruBackend
;
pub
(
crate
)
use
multi_lru_backend
::
MultiLruBackend
;
// Not used widely yet
pub
use
reuse_policy
::{
ReusePolicy
,
ReusePolicyError
};
use
super
::
SequenceHash
;
lib/kvbm-logical/src/pools/inactive/backends/multi_lru_backend.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
use
std
::
num
::
NonZeroUsize
;
use
std
::
sync
::
Arc
;
use
lru
::
LruCache
;
use
super
::
super
::
InactivePoolBackend
;
use
super
::{
Block
,
BlockMetadata
,
Registered
,
SequenceHash
};
use
crate
::
tinylfu
::
FrequencyTracker
;
use
anyhow
::{
Result
,
bail
};
pub
struct
MultiLruBackend
<
T
:
BlockMetadata
>
{
priority_pools
:
[
LruCache
<
SequenceHash
,
Block
<
T
,
Registered
>>
;
4
],
frequency_tracker
:
Arc
<
dyn
FrequencyTracker
<
u128
>>
,
frequency_thresholds
:
[
u8
;
3
],
}
impl
<
T
:
BlockMetadata
>
MultiLruBackend
<
T
>
{
/// Create with custom frequency thresholds
/// The 4 levels are fixed, but thresholds can be customized
///
/// # Arguments
/// * `block_count` - Number of blocks in the pool
/// * `thresholds` - Array of 3 thresholds: [cold->warm, warm->hot, hot->very_hot]
/// * `frequency_tracker` - Shared frequency tracker
pub
fn
new_with_thresholds
(
block_count
:
NonZeroUsize
,
thresholds
:
&
[
u8
;
3
],
frequency_tracker
:
Arc
<
dyn
FrequencyTracker
<
u128
>>
,
)
->
Result
<
Self
>
{
// Validate thresholds
if
!
(
thresholds
[
0
]
<
thresholds
[
1
]
&&
thresholds
[
1
]
<
thresholds
[
2
])
{
bail!
(
"Thresholds must be in ascending order: {:?}"
,
thresholds
);
}
if
thresholds
[
2
]
>
15
{
bail!
(
"Maximum threshold cannot exceed 15 (4-bit counter limit), got: {:?}"
,
thresholds
);
}
if
thresholds
[
0
]
<
1
{
bail!
(
"Cold threshold must be >= 1 to distinguish from never-accessed blocks, got: {:?}"
,
thresholds
);
}
Ok
(
Self
{
priority_pools
:
[
LruCache
::
new
(
block_count
),
LruCache
::
new
(
block_count
),
LruCache
::
new
(
block_count
),
LruCache
::
new
(
block_count
),
],
frequency_tracker
,
frequency_thresholds
:
*
thresholds
,
})
}
fn
calculate_priority_level
(
&
self
,
seq_hash
:
SequenceHash
)
->
usize
{
let
frequency
=
self
.frequency_tracker
.count
(
seq_hash
.as_u128
());
let
[
t1
,
t2
,
t3
]
=
self
.frequency_thresholds
;
if
frequency
<
t1
as
u32
{
0
// Cold: 0 to (t1 - 1)
}
else
if
frequency
<
t2
as
u32
{
1
// Warm: t1 to (t2 - 1)
}
else
if
frequency
<
t3
as
u32
{
2
// Hot: t2 to (t3 - 1)
}
else
{
3
// Very Hot: t3 to 15
}
}
}
impl
<
T
:
BlockMetadata
>
InactivePoolBackend
<
T
>
for
MultiLruBackend
<
T
>
{
fn
find_matches
(
&
mut
self
,
hashes
:
&
[
SequenceHash
],
touch
:
bool
)
->
Vec
<
Block
<
T
,
Registered
>>
{
let
mut
matches
=
Vec
::
with_capacity
(
hashes
.len
());
for
hash
in
hashes
{
let
mut
found
=
false
;
for
pool
in
&
mut
self
.priority_pools
{
if
let
Some
(
block
)
=
pool
.pop
(
hash
)
{
matches
.push
(
block
);
if
touch
{
self
.frequency_tracker
.touch
(
hash
.as_u128
());
}
found
=
true
;
break
;
}
}
if
!
found
{
break
;
}
}
matches
}
fn
scan_matches
(
&
mut
self
,
hashes
:
&
[
SequenceHash
],
touch
:
bool
,
)
->
Vec
<
(
SequenceHash
,
Block
<
T
,
Registered
>
)
>
{
let
mut
matches
=
Vec
::
new
();
for
hash
in
hashes
{
for
pool
in
&
mut
self
.priority_pools
{
if
let
Some
(
block
)
=
pool
.pop
(
hash
)
{
if
touch
{
self
.frequency_tracker
.touch
(
hash
.as_u128
());
}
matches
.push
((
*
hash
,
block
));
break
;
// Found in this pool, move to next hash
}
}
// Unlike find_matches: NO break on miss - continue scanning
}
matches
}
fn
allocate
(
&
mut
self
,
count
:
usize
)
->
Vec
<
Block
<
T
,
Registered
>>
{
let
mut
allocated
=
Vec
::
with_capacity
(
count
);
for
_
in
0
..
count
{
let
mut
found
=
false
;
for
pool
in
&
mut
self
.priority_pools
{
if
let
Some
((
_
seq_hash
,
block
))
=
pool
.pop_lru
()
{
allocated
.push
(
block
);
found
=
true
;
break
;
}
}
if
!
found
{
break
;
}
}
allocated
}
fn
insert
(
&
mut
self
,
block
:
Block
<
T
,
Registered
>
)
{
let
seq_hash
=
block
.sequence_hash
();
let
level
=
self
.calculate_priority_level
(
seq_hash
);
// Assert the target pool isn't full (would cause eviction)
debug_assert!
(
self
.priority_pools
[
level
]
.len
()
<
self
.priority_pools
[
level
]
.cap
()
.get
(),
"MultiLRU level {} insert would cause eviction! len={}, cap={}.
\
This indicates insufficient capacity for all blocks."
,
level
,
self
.priority_pools
[
level
]
.len
(),
self
.priority_pools
[
level
]
.cap
()
.get
()
);
self
.priority_pools
[
level
]
.put
(
seq_hash
,
block
);
}
fn
len
(
&
self
)
->
usize
{
self
.priority_pools
.iter
()
.map
(|
pool
|
pool
.len
())
.sum
()
}
fn
has_block
(
&
self
,
seq_hash
:
SequenceHash
)
->
bool
{
self
.priority_pools
.iter
()
.any
(|
pool
|
pool
.peek
(
&
seq_hash
)
.is_some
())
}
fn
allocate_all
(
&
mut
self
)
->
Vec
<
Block
<
T
,
Registered
>>
{
let
total_len
:
usize
=
self
.priority_pools
.iter
()
.map
(|
p
|
p
.len
())
.sum
();
let
mut
allocated
=
Vec
::
with_capacity
(
total_len
);
for
pool
in
&
mut
self
.priority_pools
{
while
let
Some
((
_
seq_hash
,
block
))
=
pool
.pop_lru
()
{
allocated
.push
(
block
);
}
}
allocated
}
}
#[cfg(test)]
mod
tests
{
use
super
::
*
;
use
crate
::
pools
::
tests
::
fixtures
::
*
;
use
crate
::
tinylfu
::
TinyLFUTracker
;
impl
<
T
:
BlockMetadata
>
MultiLruBackend
<
T
>
{
pub
fn
new
(
capacity
:
NonZeroUsize
,
frequency_tracker
:
Arc
<
dyn
FrequencyTracker
<
u128
>>
,
)
->
Self
{
Self
::
new_with_thresholds
(
capacity
,
&
[
2
,
6
,
15
],
frequency_tracker
)
.unwrap
()
}
}
#[test]
fn
test_multi_lru_priority_levels
()
{
let
frequency_tracker
=
Arc
::
new
(
TinyLFUTracker
::
new
(
100
));
let
mut
backend
=
MultiLruBackend
::
new
(
NonZeroUsize
::
new
(
12
)
.unwrap
(),
frequency_tracker
.clone
());
let
(
block1
,
hash1
)
=
create_registered_block
(
1
,
&
tokens_for_id
(
1
));
let
(
block2
,
hash2
)
=
create_registered_block
(
2
,
&
tokens_for_id
(
2
));
let
(
block3
,
hash3
)
=
create_registered_block
(
3
,
&
tokens_for_id
(
3
));
let
(
block4
,
hash4
)
=
create_registered_block
(
4
,
&
tokens_for_id
(
4
));
frequency_tracker
.touch
(
hash2
.as_u128
());
frequency_tracker
.touch
(
hash2
.as_u128
());
for
_
in
0
..
6
{
frequency_tracker
.touch
(
hash3
.as_u128
());
}
for
_
in
0
..
16
{
frequency_tracker
.touch
(
hash4
.as_u128
());
}
let
_
freq1
=
frequency_tracker
.count
(
hash1
.as_u128
());
let
_
freq2
=
frequency_tracker
.count
(
hash2
.as_u128
());
let
_
freq3
=
frequency_tracker
.count
(
hash3
.as_u128
());
let
_
freq4
=
frequency_tracker
.count
(
hash4
.as_u128
());
assert_eq!
(
backend
.calculate_priority_level
(
hash1
),
0
);
// Cold
assert_eq!
(
backend
.calculate_priority_level
(
hash2
),
1
);
// Warm
assert_eq!
(
backend
.calculate_priority_level
(
hash3
),
2
);
// Hot
assert_eq!
(
backend
.calculate_priority_level
(
hash4
),
3
);
// Very hot (15)
backend
.insert
(
block1
);
backend
.insert
(
block2
);
backend
.insert
(
block3
);
backend
.insert
(
block4
);
assert_eq!
(
backend
.len
(),
4
);
assert
!
(
backend
.has_block
(
hash1
));
assert
!
(
backend
.has_block
(
hash2
));
assert
!
(
backend
.has_block
(
hash3
));
assert
!
(
backend
.has_block
(
hash4
));
}
#[test]
fn
test_multi_lru_eviction_order
()
{
let
frequency_tracker
=
Arc
::
new
(
TinyLFUTracker
::
new
(
100
));
let
mut
backend
=
MultiLruBackend
::
new
(
NonZeroUsize
::
new
(
8
)
.unwrap
(),
frequency_tracker
.clone
());
let
(
block1
,
hash1
)
=
create_registered_block
(
1
,
&
tokens_for_id
(
1
));
let
(
block2
,
hash2
)
=
create_registered_block
(
2
,
&
tokens_for_id
(
2
));
let
(
block3
,
hash3
)
=
create_registered_block
(
3
,
&
tokens_for_id
(
3
));
for
_
in
0
..
6
{
frequency_tracker
.touch
(
hash3
.as_u128
());
}
backend
.insert
(
block1
);
backend
.insert
(
block2
);
backend
.insert
(
block3
);
let
allocated
=
backend
.allocate
(
2
);
assert_eq!
(
allocated
.len
(),
2
);
assert_eq!
(
allocated
[
0
]
.block_id
(),
1
);
assert_eq!
(
allocated
[
1
]
.block_id
(),
2
);
assert
!
(
!
backend
.has_block
(
hash1
));
assert
!
(
!
backend
.has_block
(
hash2
));
assert
!
(
backend
.has_block
(
hash3
));
}
#[test]
fn
test_multi_lru_find_matches
()
{
let
frequency_tracker
=
Arc
::
new
(
TinyLFUTracker
::
new
(
100
));
let
mut
backend
=
MultiLruBackend
::
new_with_thresholds
(
NonZeroUsize
::
new
(
8
)
.unwrap
(),
&
[
2
,
4
,
8
],
frequency_tracker
.clone
(),
)
.unwrap
();
let
(
block1
,
hash1
)
=
create_registered_block
(
1
,
&
tokens_for_id
(
1
));
let
(
block2
,
hash2
)
=
create_registered_block
(
2
,
&
tokens_for_id
(
2
));
let
(
block3
,
hash3
)
=
create_registered_block
(
3
,
&
tokens_for_id
(
3
));
for
_
in
0
..
3
{
frequency_tracker
.touch
(
hash2
.as_u128
());
}
for
_
in
0
..
10
{
frequency_tracker
.touch
(
hash3
.as_u128
());
}
backend
.insert
(
block1
);
backend
.insert
(
block2
);
backend
.insert
(
block3
);
let
matches
=
backend
.find_matches
(
&
[
hash1
,
hash2
,
hash3
],
true
);
assert_eq!
(
matches
.len
(),
3
);
assert_eq!
(
backend
.len
(),
0
);
}
#[test]
fn
test_multi_lru_capacity_distribution
()
{
let
frequency_tracker
=
Arc
::
new
(
TinyLFUTracker
::
new
(
100
));
let
mut
backend
=
MultiLruBackend
::
new_with_thresholds
(
NonZeroUsize
::
new
(
16
)
.unwrap
(),
&
[
2
,
6
,
15
],
frequency_tracker
.clone
(),
)
.unwrap
();
// Create blocks with different frequencies to test distribution across levels
let
(
block1
,
hash1
)
=
create_registered_block
(
1
,
&
tokens_for_id
(
1
));
// Level 0 (cold, freq=0)
let
(
block2
,
hash2
)
=
create_registered_block
(
2
,
&
tokens_for_id
(
2
));
// Level 1 (warm, freq=3)
let
(
block3
,
hash3
)
=
create_registered_block
(
3
,
&
tokens_for_id
(
3
));
// Level 2 (hot, freq=7)
let
(
block4
,
hash4
)
=
create_registered_block
(
4
,
&
tokens_for_id
(
4
));
// Level 3 (very hot, freq=15)
// Set up frequency tracking
for
_
in
0
..
3
{
frequency_tracker
.touch
(
hash2
.as_u128
());
// Warm: frequency 3
}
for
_
in
0
..
7
{
frequency_tracker
.touch
(
hash3
.as_u128
());
// Hot: frequency 7
}
for
_
in
0
..
15
{
frequency_tracker
.touch
(
hash4
.as_u128
());
// Very hot: frequency 15
}
// Verify priority level calculation
assert_eq!
(
backend
.calculate_priority_level
(
hash1
),
0
);
// Cold (freq=0)
assert_eq!
(
backend
.calculate_priority_level
(
hash2
),
1
);
// Warm (freq=3)
assert_eq!
(
backend
.calculate_priority_level
(
hash3
),
2
);
// Hot (freq=7)
assert_eq!
(
backend
.calculate_priority_level
(
hash4
),
3
);
// Very hot (freq=15)
// Insert blocks - should be distributed across different levels
backend
.insert
(
block1
);
backend
.insert
(
block2
);
backend
.insert
(
block3
);
backend
.insert
(
block4
);
assert_eq!
(
backend
.len
(),
4
);
assert
!
(
backend
.has_block
(
hash1
));
assert
!
(
backend
.has_block
(
hash2
));
assert
!
(
backend
.has_block
(
hash3
));
assert
!
(
backend
.has_block
(
hash4
));
// Test that we can allocate from all levels
let
allocated
=
backend
.allocate
(
4
);
assert_eq!
(
allocated
.len
(),
4
);
assert_eq!
(
backend
.len
(),
0
);
}
}
lib/kvbm-logical/src/pools/inactive/backends/reuse_policy.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Reuse policies for determining block allocation priority.
//!
//! Different policies (FIFO, LRU, etc.) control which inactive registered
//! block should be allocated next when the reset pool is exhausted.
use
super
::{
BlockId
,
InactiveBlock
};
#[derive(Debug,
thiserror::Error)]
pub
enum
ReusePolicyError
{
#[error(
"Block {0} already exists in free list"
)]
BlockAlreadyExists
(
BlockId
),
#[error(
"Block {0} not found in free list"
)]
BlockNotFound
(
BlockId
),
}
/// Trait for managing a free list of blocks
///
/// Different implementations can provide different priority strategies
/// for selecting which block to allocate next.
pub
trait
ReusePolicy
:
Send
+
Sync
+
std
::
fmt
::
Debug
{
/// Insert a block into the free list
///
/// The implementation will compute the priority key and manage the free list
/// based on its specific strategy.
fn
insert
(
&
mut
self
,
inactive_block
:
InactiveBlock
)
->
Result
<
(),
ReusePolicyError
>
;
/// Remove a specific block from the free list
fn
remove
(
&
mut
self
,
block_id
:
BlockId
)
->
Result
<
(),
ReusePolicyError
>
;
/// Get the next free block based on the implementation's priority strategy
///
/// Returns None if the free list is empty.
/// The returned InactiveBlock contains both the block_id and seq_hash needed
/// to look up the block in the InactivePool's HashMap.
fn
next_free
(
&
mut
self
)
->
Option
<
InactiveBlock
>
;
/// Check if the free list is empty
fn
is_empty
(
&
self
)
->
bool
;
/// Get the number of free blocks
fn
len
(
&
self
)
->
usize
;
}
lib/kvbm-logical/src/pools/inactive/backends/tests.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
#[cfg(test)]
mod
backend_tests
{
use
std
::
num
::
NonZeroUsize
;
use
rstest
::
rstest
;
use
crate
::{
BlockId
,
pools
::
tests
::
fixtures
::
*
,
tinylfu
::
TinyLFUTracker
};
// Import TestData from fixtures (which re-exports from testing)
use
crate
::
pools
::
tests
::
TestData
;
use
super
::
super
::
*
;
#[derive(Clone,
Copy,
Debug)]
enum
BackendType
{
HashMap
,
Lru
,
MultiLru
,
Lineage
,
}
fn
create_backend
(
backend_type
:
BackendType
)
->
Box
<
dyn
InactivePoolBackend
<
TestData
>>
{
match
backend_type
{
BackendType
::
HashMap
=>
Box
::
new
(
HashMapBackend
::
new
(
Box
::
new
(
FifoReusePolicy
::
new
()))),
BackendType
::
Lru
=>
Box
::
new
(
LruBackend
::
new
(
NonZeroUsize
::
new
(
10
)
.unwrap
())),
BackendType
::
MultiLru
=>
Box
::
new
(
MultiLruBackend
::
new
(
NonZeroUsize
::
new
(
10
)
.unwrap
(),
Arc
::
new
(
TinyLFUTracker
::
new
(
100
)),
)),
BackendType
::
Lineage
=>
Box
::
new
(
LineageBackend
::
new
()),
}
}
#[rstest]
#[case::hashmap(BackendType::HashMap)]
#[case::lru(BackendType::Lru)]
#[case::lineage(BackendType::Lineage)]
fn
test_insert_and_len
(
#[case]
backend_type
:
BackendType
)
{
let
mut
backend
=
create_backend
(
backend_type
);
assert_eq!
(
backend
.len
(),
0
);
assert
!
(
backend
.is_empty
());
let
(
block
,
_
)
=
create_registered_block
(
1
,
&
tokens_for_id
(
1
));
backend
.insert
(
block
);
assert_eq!
(
backend
.len
(),
1
);
assert
!
(
!
backend
.is_empty
());
}
#[rstest]
#[case::hashmap(BackendType::HashMap)]
#[case::lru(BackendType::Lru)]
#[case::lineage(BackendType::Lineage)]
fn
test_has_block
(
#[case]
backend_type
:
BackendType
)
{
let
mut
backend
=
create_backend
(
backend_type
);
let
(
block
,
hash
)
=
create_registered_block
(
1
,
&
tokens_for_id
(
1
));
backend
.insert
(
block
);
assert
!
(
backend
.has_block
(
hash
));
// assert!(!backend.has_block(999));
}
#[rstest]
#[case::hashmap(BackendType::HashMap)]
#[case::lru(BackendType::Lru)]
#[case::lineage(BackendType::Lineage)]
fn
test_find_matches
(
#[case]
backend_type
:
BackendType
)
{
let
mut
backend
=
create_backend
(
backend_type
);
let
(
block1
,
hash1
)
=
create_registered_block
(
1
,
&
tokens_for_id
(
1
));
let
(
block2
,
hash2
)
=
create_registered_block
(
2
,
&
tokens_for_id
(
2
));
backend
.insert
(
block1
);
// For HashMap backend with FIFO, we need a sleep to ensure different timestamps
if
matches!
(
backend_type
,
BackendType
::
HashMap
)
{
std
::
thread
::
sleep
(
std
::
time
::
Duration
::
from_millis
(
2
));
}
backend
.insert
(
block2
);
let
matches
=
backend
.find_matches
(
&
[
hash1
,
hash2
],
true
);
assert_eq!
(
matches
.len
(),
2
);
assert_eq!
(
backend
.len
(),
0
);
}
#[rstest]
#[case::hashmap(BackendType::HashMap)]
#[case::lru(BackendType::Lru)]
#[case::lineage(BackendType::Lineage)]
fn
test_find_matches_stops_on_miss
(
#[case]
backend_type
:
BackendType
)
{
let
mut
backend
=
create_backend
(
backend_type
);
let
(
block1
,
hash1
)
=
create_registered_block
(
1
,
&
tokens_for_id
(
1
));
let
(
block2
,
hash2
)
=
create_registered_block
(
2
,
&
tokens_for_id
(
2
));
let
block_na
=
create_complete_block
(
3
,
&
tokens_for_id
(
3
));
backend
.insert
(
block1
);
backend
.insert
(
block2
);
let
matches
=
backend
.find_matches
(
&
[
hash1
,
block_na
.sequence_hash
(),
hash2
],
true
);
assert_eq!
(
matches
.len
(),
1
);
assert_eq!
(
backend
.len
(),
1
);
}
#[rstest]
#[case::hashmap(BackendType::HashMap)]
#[case::lru(BackendType::Lru)]
#[case::lineage(BackendType::Lineage)]
fn
test_allocate
(
#[case]
backend_type
:
BackendType
)
{
let
mut
backend
=
create_backend
(
backend_type
);
let
(
block1
,
_
)
=
create_registered_block
(
1
,
&
tokens_for_id
(
1
));
backend
.insert
(
block1
);
// For HashMap backend with FIFO, we need a sleep to ensure different timestamps
if
matches!
(
backend_type
,
BackendType
::
HashMap
)
{
std
::
thread
::
sleep
(
std
::
time
::
Duration
::
from_millis
(
2
));
}
let
(
block2
,
_
)
=
create_registered_block
(
2
,
&
tokens_for_id
(
2
));
backend
.insert
(
block2
);
let
allocated
=
backend
.allocate
(
1
);
assert_eq!
(
allocated
.len
(),
1
);
assert_eq!
(
backend
.len
(),
1
);
}
#[rstest]
#[case::hashmap(BackendType::HashMap)]
#[case::lru(BackendType::Lru)]
#[case::multi_lru(BackendType::MultiLru)]
#[case::lineage(BackendType::Lineage)]
fn
test_allocate_more_than_available
(
#[case]
backend_type
:
BackendType
)
{
let
mut
backend
=
create_backend
(
backend_type
);
let
(
block1
,
_
)
=
create_registered_block
(
1
,
&
tokens_for_id
(
1
));
backend
.insert
(
block1
);
// For HashMap backend with FIFO, we need a sleep to ensure different timestamps
if
matches!
(
backend_type
,
BackendType
::
HashMap
)
{
std
::
thread
::
sleep
(
std
::
time
::
Duration
::
from_millis
(
2
));
}
let
(
block2
,
_
)
=
create_registered_block
(
2
,
&
tokens_for_id
(
2
));
backend
.insert
(
block2
);
let
allocated
=
backend
.allocate
(
5
);
assert_eq!
(
allocated
.len
(),
2
);
assert_eq!
(
backend
.len
(),
0
);
}
#[rstest]
#[case::hashmap(BackendType::HashMap)]
#[case::lru(BackendType::Lru)]
#[case::multi_lru(BackendType::MultiLru)]
#[case::lineage(BackendType::Lineage)]
fn
test_allocate_all
(
#[case]
backend_type
:
BackendType
)
{
let
mut
backend
=
create_backend
(
backend_type
);
// Insert several blocks
let
block_ids
:
Vec
<
u64
>
=
vec!
[
1
,
2
,
3
,
4
,
5
];
for
&
i
in
&
block_ids
{
let
(
block
,
_
)
=
create_registered_block
(
i
as
BlockId
,
&
tokens_for_id
(
i
));
backend
.insert
(
block
);
// For HashMap backend with FIFO, we need a sleep to ensure different timestamps
if
matches!
(
backend_type
,
BackendType
::
HashMap
)
{
std
::
thread
::
sleep
(
std
::
time
::
Duration
::
from_millis
(
2
));
}
}
assert_eq!
(
backend
.len
(),
5
);
// Allocate all blocks
let
allocated
=
backend
.allocate_all
();
assert_eq!
(
allocated
.len
(),
5
);
assert_eq!
(
backend
.len
(),
0
);
assert
!
(
backend
.is_empty
());
}
#[rstest]
#[case::hashmap(BackendType::HashMap)]
#[case::lru(BackendType::Lru)]
#[case::multi_lru(BackendType::MultiLru)]
#[case::lineage(BackendType::Lineage)]
fn
test_allocate_all_empty_pool
(
#[case]
backend_type
:
BackendType
)
{
let
mut
backend
=
create_backend
(
backend_type
);
assert_eq!
(
backend
.len
(),
0
);
// Allocate all from empty pool
let
allocated
=
backend
.allocate_all
();
assert_eq!
(
allocated
.len
(),
0
);
assert
!
(
backend
.is_empty
());
}
#[rstest]
#[case::hashmap(BackendType::HashMap)]
#[case::lru(BackendType::Lru)]
#[case::multi_lru(BackendType::MultiLru)]
#[case::lineage(BackendType::Lineage)]
fn
test_scan_matches
(
#[case]
backend_type
:
BackendType
)
{
let
mut
backend
=
create_backend
(
backend_type
);
let
(
block1
,
hash1
)
=
create_registered_block
(
1
,
&
tokens_for_id
(
1
));
let
(
block2
,
_
hash2
)
=
create_registered_block
(
2
,
&
tokens_for_id
(
2
));
let
(
block3
,
hash3
)
=
create_registered_block
(
3
,
&
tokens_for_id
(
3
));
let
missing_block
=
create_complete_block
(
4
,
&
tokens_for_id
(
4
));
let
missing_hash
=
missing_block
.sequence_hash
();
backend
.insert
(
block1
);
// For HashMap backend with FIFO, we need a sleep to ensure different timestamps
if
matches!
(
backend_type
,
BackendType
::
HashMap
)
{
std
::
thread
::
sleep
(
std
::
time
::
Duration
::
from_millis
(
2
));
}
backend
.insert
(
block2
);
if
matches!
(
backend_type
,
BackendType
::
HashMap
)
{
std
::
thread
::
sleep
(
std
::
time
::
Duration
::
from_millis
(
2
));
}
backend
.insert
(
block3
);
assert_eq!
(
backend
.len
(),
3
);
// Scan for [hash1, missing_hash, hash3] - should continue past the miss
let
matches
=
backend
.scan_matches
(
&
[
hash1
,
missing_hash
,
hash3
],
true
);
assert_eq!
(
matches
.len
(),
2
,
"scan_matches should find 2 blocks, skipping the miss"
);
// Verify correct hashes were returned
let
found_hashes
:
Vec
<
_
>
=
matches
.iter
()
.map
(|(
h
,
_
)|
*
h
)
.collect
();
assert
!
(
found_hashes
.contains
(
&
hash1
));
assert
!
(
found_hashes
.contains
(
&
hash3
));
// block2 should still be in the backend (was not scanned for)
assert_eq!
(
backend
.len
(),
1
,
"Only block2 should remain"
);
}
}
lib/kvbm-logical/src/pools/inactive/mod.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Thread-safe pool for registered immutable blocks with automatic RAII return.
//!
//! Manages blocks in the Registered state, providing:
//! - Finding blocks by sequence hash with O(1) lookup
//! - Conversion of registered blocks back to mutable blocks for reuse
//! - Thread-safe access via interior mutability
//! - Automatic block return via RAII ImmutableBlock guards
pub
mod
backends
;
use
parking_lot
::
RwLock
;
use
std
::
sync
::
Arc
;
use
crate
::
metrics
::
BlockPoolMetrics
;
use
super
::{
Block
,
BlockId
,
BlockMetadata
,
InactiveBlock
,
MutableBlock
,
PrimaryBlock
,
Registered
,
RegisteredBlock
,
SequenceHash
,
reset
::
ResetPool
,
};
// pub(crate) use backends::*;
/// Backend trait for InactivePool storage strategies
pub
(
crate
)
trait
InactivePoolBackend
<
T
:
BlockMetadata
>
:
Send
+
Sync
{
/// Find blocks matching the given hashes in order, stopping on first miss.
fn
find_matches
(
&
mut
self
,
hashes
:
&
[
SequenceHash
],
touch
:
bool
)
->
Vec
<
Block
<
T
,
Registered
>>
;
/// Scan for blocks matching any of the given hashes (full scan, doesn't stop on miss).
/// Unlike find_matches, continues scanning even when a hash is not found.
/// Acquires/removes found blocks from pool (caller owns until dropped).
fn
scan_matches
(
&
mut
self
,
hashes
:
&
[
SequenceHash
],
touch
:
bool
,
)
->
Vec
<
(
SequenceHash
,
Block
<
T
,
Registered
>
)
>
;
fn
allocate
(
&
mut
self
,
count
:
usize
)
->
Vec
<
Block
<
T
,
Registered
>>
;
fn
insert
(
&
mut
self
,
block
:
Block
<
T
,
Registered
>
);
fn
len
(
&
self
)
->
usize
;
#[allow(dead_code)]
fn
is_empty
(
&
self
)
->
bool
{
self
.len
()
==
0
}
#[allow(dead_code)]
fn
has_block
(
&
self
,
seq_hash
:
SequenceHash
)
->
bool
;
/// Allocate all blocks from the pool, removing them from the backend.
/// Default implementation calls len() then allocate(), which is atomic
/// since the caller holds the lock.
fn
allocate_all
(
&
mut
self
)
->
Vec
<
Block
<
T
,
Registered
>>
{
let
count
=
self
.len
();
self
.allocate
(
count
)
}
}
use
crate
::
blocks
::{
RegisteredReturnFn
,
ResetReturnFn
};
/// Pool for managing registered (immutable) blocks
///
/// This pool handles blocks in the Registered state and provides them as
/// RegisteredBlock RAII guards that automatically return to the pool on drop.
#[derive(Clone)]
pub
(
crate
)
struct
InactivePool
<
T
:
BlockMetadata
>
{
// Inner state protected by RwLock for thread-safe access from guards
inner
:
Arc
<
RwLock
<
InactivePoolInner
<
T
>>>
,
// Return function for MutableBlocks to return to ResetPool
reset_return_fn
:
ResetReturnFn
<
T
>
,
return_fn
:
RegisteredReturnFn
<
T
>
,
#[expect(dead_code)]
block_size
:
usize
,
metrics
:
Option
<
Arc
<
BlockPoolMetrics
>>
,
}
struct
InactivePoolInner
<
T
:
BlockMetadata
>
{
backend
:
Box
<
dyn
InactivePoolBackend
<
T
>>
,
}
impl
<
T
:
BlockMetadata
+
Sync
>
InactivePool
<
T
>
{
/// Create a new InactivePool with the given backend and reset pool
pub
(
crate
)
fn
new
(
backend
:
Box
<
dyn
InactivePoolBackend
<
T
>>
,
reset_pool
:
&
ResetPool
<
T
>
,
metrics
:
Option
<
Arc
<
BlockPoolMetrics
>>
,
)
->
Self
{
let
inner
=
Arc
::
new
(
RwLock
::
new
(
InactivePoolInner
{
backend
}));
let
inner_clone
=
inner
.clone
();
let
metrics_clone
=
metrics
.clone
();
let
return_fn
=
Arc
::
new
(
move
|
block
:
Arc
<
Block
<
T
,
Registered
>>
|
{
let
seq_hash
=
block
.sequence_hash
();
let
mut
inner
=
inner_clone
.write
();
match
Arc
::
try_unwrap
(
block
)
{
Ok
(
block
)
=>
{
let
block_id
=
block
.block_id
();
inner
.backend
.insert
(
block
);
if
let
Some
(
ref
m
)
=
metrics_clone
{
m
.inc_inactive_pool_size
();
}
tracing
::
trace!
(
?
seq_hash
,
block_id
,
"Block stored in inactive pool"
);
}
Err
(
block
)
=>
{
let
block_id
=
block
.block_id
();
let
weak
=
Arc
::
downgrade
(
&
block
);
drop
(
block
);
if
weak
.strong_count
()
==
0
{
tracing
::
warn!
(
?
seq_hash
,
block_id
,
"Possible KV Block leak detected"
);
}
}
}
})
as
Arc
<
dyn
Fn
(
Arc
<
Block
<
T
,
Registered
>>
)
+
Send
+
Sync
>
;
Self
{
inner
,
reset_return_fn
:
reset_pool
.return_fn
(),
return_fn
,
block_size
:
reset_pool
.block_size
(),
metrics
,
}
}
/// Find blocks by sequence hashes and return them as RegisteredBlock guards.
/// Stops on first miss.
pub
(
crate
)
fn
find_blocks
(
&
self
,
hashes
:
&
[
SequenceHash
],
touch
:
bool
,
)
->
Vec
<
Arc
<
dyn
RegisteredBlock
<
T
>>>
{
let
mut
inner
=
self
.inner
.write
();
let
matched_blocks
=
inner
.backend
.find_matches
(
hashes
,
touch
);
let
count
=
matched_blocks
.len
();
if
let
Some
(
ref
m
)
=
self
.metrics
{
for
_
in
0
..
count
{
m
.dec_inactive_pool_size
();
}
}
matched_blocks
.into_iter
()
.map
(|
block
|
{
PrimaryBlock
::
new_attached
(
Arc
::
new
(
block
),
self
.return_fn
.clone
())
as
Arc
<
dyn
RegisteredBlock
<
T
>>
})
.collect
()
}
/// Scan for all blocks matching the given hashes (doesn't stop on miss).
/// Acquires/removes found blocks from pool - caller owns until dropped.
/// Returns RAII guards (PrimaryBlocks) for found blocks.
pub
(
crate
)
fn
scan_blocks
(
&
self
,
hashes
:
&
[
SequenceHash
],
touch
:
bool
,
)
->
Vec
<
(
SequenceHash
,
Arc
<
dyn
RegisteredBlock
<
T
>>
)
>
{
let
mut
inner
=
self
.inner
.write
();
let
found
=
inner
.backend
.scan_matches
(
hashes
,
touch
);
let
count
=
found
.len
();
if
let
Some
(
ref
m
)
=
self
.metrics
{
for
_
in
0
..
count
{
m
.dec_inactive_pool_size
();
}
}
found
.into_iter
()
.map
(|(
hash
,
block
)|
{
let
registered
=
PrimaryBlock
::
new_attached
(
Arc
::
new
(
block
),
self
.return_fn
.clone
())
as
Arc
<
dyn
RegisteredBlock
<
T
>>
;
(
hash
,
registered
)
})
.collect
()
}
/// Allocate blocks from registered pool, converting them to MutableBlocks for ResetPool
pub
(
crate
)
fn
allocate_blocks
(
&
self
,
count
:
usize
)
->
Option
<
Vec
<
MutableBlock
<
T
>>>
{
if
count
==
0
{
return
Some
(
Vec
::
new
());
}
let
mut
inner
=
self
.inner
.write
();
if
inner
.backend
.len
()
<
count
{
return
None
;
}
let
allocated_blocks
=
inner
.backend
.allocate
(
count
);
if
allocated_blocks
.len
()
==
count
{
if
let
Some
(
ref
m
)
=
self
.metrics
{
for
_
in
0
..
count
{
m
.dec_inactive_pool_size
();
}
}
let
mut
mutable_blocks
=
Vec
::
with_capacity
(
count
);
mutable_blocks
.extend
(
allocated_blocks
.into_iter
()
.map
(|
registered_block
|
{
let
reset_block
=
registered_block
.reset
();
MutableBlock
::
new
(
reset_block
,
self
.reset_return_fn
.clone
(),
self
.metrics
.clone
(),
)
}));
Some
(
mutable_blocks
)
}
else
{
for
block
in
allocated_blocks
{
inner
.backend
.insert
(
block
);
}
None
}
}
/// Check if a block exists in the pool
#[allow(dead_code)]
pub
(
crate
)
fn
has_block
(
&
self
,
hash
:
SequenceHash
)
->
bool
{
let
inner
=
self
.inner
.read
();
inner
.backend
.has_block
(
hash
)
}
/// Find and promote a single block from inactive to active by sequence hash.
/// Returns the concrete `Arc<PrimaryBlock<T>>` for duplicate referencing.
///
/// This differs from `find_blocks()` which returns trait objects. This method
/// returns the concrete type needed when creating `DuplicateBlock` references.
///
/// Uses `new_unattached` because this is called from `try_find_existing_block`
/// while the attachments lock is held. The caller MUST call
/// `PrimaryBlock::store_weak_refs()` after dropping the attachments lock.
pub
(
crate
)
fn
find_block_as_primary
(
&
self
,
hash
:
SequenceHash
,
touch
:
bool
,
)
->
Option
<
Arc
<
PrimaryBlock
<
T
>>>
{
let
mut
inner
=
self
.inner
.write
();
let
matched
=
inner
.backend
.find_matches
(
&
[
hash
],
touch
);
matched
.into_iter
()
.next
()
.map
(|
block
|
{
if
let
Some
(
ref
m
)
=
self
.metrics
{
m
.dec_inactive_pool_size
();
}
PrimaryBlock
::
new_unattached
(
Arc
::
new
(
block
),
self
.return_fn
.clone
())
})
}
/// Get the number of blocks in the pool
pub
(
crate
)
fn
len
(
&
self
)
->
usize
{
let
inner
=
self
.inner
.read
();
inner
.backend
.len
()
}
/// Check if the pool is empty
#[allow(dead_code)]
pub
(
crate
)
fn
is_empty
(
&
self
)
->
bool
{
let
inner
=
self
.inner
.read
();
inner
.backend
.is_empty
()
}
pub
(
crate
)
fn
return_fn
(
&
self
)
->
RegisteredReturnFn
<
T
>
{
self
.return_fn
.clone
()
}
/// Allocate all blocks from the pool, converting them to MutableBlocks.
/// The MutableBlocks will return to the ResetPool when dropped via RAII.
pub
(
crate
)
fn
allocate_all_blocks
(
&
self
)
->
Vec
<
MutableBlock
<
T
>>
{
let
mut
inner
=
self
.inner
.write
();
let
blocks
=
inner
.backend
.allocate_all
();
let
count
=
blocks
.len
();
if
let
Some
(
ref
m
)
=
self
.metrics
{
for
_
in
0
..
count
{
m
.dec_inactive_pool_size
();
}
}
blocks
.into_iter
()
.map
(|
registered_block
|
{
let
reset_block
=
registered_block
.reset
();
MutableBlock
::
new
(
reset_block
,
self
.reset_return_fn
.clone
(),
self
.metrics
.clone
(),
)
})
.collect
()
}
}
#[cfg(test)]
mod
tests
{
use
super
::
backends
::
FifoReusePolicy
;
use
super
::
*
;
use
crate
::
testing
::{
TestMeta
,
create_registered_block
,
tokens_for_id
};
impl
<
T
:
BlockMetadata
>
InactivePool
<
T
>
{
fn
insert
(
&
self
,
block
:
Block
<
T
,
Registered
>
)
{
let
mut
inner
=
self
.inner
.write
();
inner
.backend
.insert
(
block
);
}
}
fn
create_test_pool
()
->
(
InactivePool
<
TestMeta
>
,
ResetPool
<
TestMeta
>
)
{
use
super
::
backends
::
HashMapBackend
;
let
reuse_policy
=
Box
::
new
(
FifoReusePolicy
::
new
());
let
backend
=
Box
::
new
(
HashMapBackend
::
new
(
reuse_policy
));
let
reset_blocks
:
Vec
<
_
>
=
(
0
..
10_u
size
)
.map
(|
i
|
Block
::
new
(
i
,
4
))
.collect
();
let
reset_pool
=
ResetPool
::
new
(
reset_blocks
,
4
,
None
);
let
inactive_pool
=
InactivePool
::
new
(
backend
,
&
reset_pool
,
None
);
(
inactive_pool
,
reset_pool
)
}
/// Create a sequence hash for a block that doesn't exist in any pool.
fn
nonexistent_hash
()
->
SequenceHash
{
// Create a registered block just to get its sequence hash, then drop it
let
(
_
,
seq_hash
)
=
create_registered_block
::
<
TestMeta
>
(
999
,
&
[
9999
,
9998
,
9997
,
9996
]);
seq_hash
}
#[test]
fn
test_new_pool_starts_empty
()
{
let
(
pool
,
_
reset_pool
)
=
create_test_pool
();
assert_eq!
(
pool
.len
(),
0
);
assert
!
(
pool
.is_empty
());
assert
!
(
!
pool
.has_block
(
nonexistent_hash
()));
}
#[test]
fn
test_return_and_find_single_block
()
{
let
(
pool
,
_
reset_pool
)
=
create_test_pool
();
let
(
block
,
seq_hash
)
=
create_registered_block
::
<
TestMeta
>
(
1
,
&
tokens_for_id
(
1
));
pool
.insert
(
block
);
assert_eq!
(
pool
.len
(),
1
);
assert
!
(
pool
.has_block
(
seq_hash
));
let
found_blocks
=
pool
.find_blocks
(
&
[
seq_hash
],
true
);
assert_eq!
(
found_blocks
.len
(),
1
);
assert_eq!
(
found_blocks
[
0
]
.block_id
(),
1
);
assert_eq!
(
found_blocks
[
0
]
.sequence_hash
(),
seq_hash
);
// Block should be removed from pool after finding
assert_eq!
(
pool
.len
(),
0
);
assert
!
(
!
pool
.has_block
(
seq_hash
));
}
#[test]
fn
test_find_blocks_stops_on_first_miss
()
{
let
(
pool
,
_
reset_pool
)
=
create_test_pool
();
let
(
block1
,
seq_hash1
)
=
create_registered_block
::
<
TestMeta
>
(
1
,
&
tokens_for_id
(
1
));
let
(
block3
,
seq_hash3
)
=
create_registered_block
::
<
TestMeta
>
(
3
,
&
tokens_for_id
(
3
));
pool
.insert
(
block1
);
pool
.insert
(
block3
);
assert_eq!
(
pool
.len
(),
2
);
let
missing
=
nonexistent_hash
();
let
found_blocks
=
pool
.find_blocks
(
&
[
seq_hash1
,
missing
,
seq_hash3
],
true
);
assert_eq!
(
found_blocks
.len
(),
1
);
assert_eq!
(
found_blocks
[
0
]
.sequence_hash
(),
seq_hash1
);
// Block 3 should still be in pool since search stopped at first miss
assert_eq!
(
pool
.len
(),
1
);
assert
!
(
pool
.has_block
(
seq_hash3
));
}
#[test]
fn
test_raii_auto_return
()
{
let
(
pool
,
_
reset_pool
)
=
create_test_pool
();
let
(
block
,
seq_hash
)
=
create_registered_block
::
<
TestMeta
>
(
1
,
&
tokens_for_id
(
1
));
pool
.insert
(
block
);
assert_eq!
(
pool
.len
(),
1
);
{
let
_
found_blocks
=
pool
.find_blocks
(
&
[
seq_hash
],
true
);
assert_eq!
(
pool
.len
(),
0
);
}
assert_eq!
(
pool
.len
(),
1
);
assert
!
(
pool
.has_block
(
seq_hash
));
}
#[test]
fn
test_allocate_blocks
()
{
let
(
pool
,
reset_pool
)
=
create_test_pool
();
let
(
block1
,
_
)
=
create_registered_block
::
<
TestMeta
>
(
1
,
&
tokens_for_id
(
1
));
let
(
block2
,
_
)
=
create_registered_block
::
<
TestMeta
>
(
2
,
&
tokens_for_id
(
2
));
let
(
block3
,
_
)
=
create_registered_block
::
<
TestMeta
>
(
3
,
&
tokens_for_id
(
3
));
pool
.insert
(
block1
);
pool
.insert
(
block2
);
pool
.insert
(
block3
);
assert_eq!
(
pool
.len
(),
3
);
let
mutable_blocks
=
pool
.allocate_blocks
(
1
)
.expect
(
"Should allocate 1 block"
);
assert_eq!
(
mutable_blocks
.len
(),
1
);
assert_eq!
(
pool
.len
(),
2
);
drop
(
mutable_blocks
);
assert_eq!
(
pool
.len
(),
2
);
assert_eq!
(
reset_pool
.available_blocks
(),
11
);
}
#[test]
fn
test_allocate_more_than_available_fails
()
{
let
(
pool
,
_
reset_pool
)
=
create_test_pool
();
let
(
block1
,
_
)
=
create_registered_block
::
<
TestMeta
>
(
1
,
&
tokens_for_id
(
1
));
let
(
block2
,
_
)
=
create_registered_block
::
<
TestMeta
>
(
2
,
&
tokens_for_id
(
2
));
pool
.insert
(
block1
);
pool
.insert
(
block2
);
assert_eq!
(
pool
.len
(),
2
);
let
result
=
pool
.allocate_blocks
(
3
);
assert
!
(
result
.is_none
());
assert_eq!
(
pool
.len
(),
2
);
}
#[test]
fn
test_scan_blocks
()
{
let
(
pool
,
_
reset_pool
)
=
create_test_pool
();
let
(
block1
,
seq_hash1
)
=
create_registered_block
::
<
TestMeta
>
(
1
,
&
tokens_for_id
(
1
));
let
(
block3
,
seq_hash3
)
=
create_registered_block
::
<
TestMeta
>
(
3
,
&
tokens_for_id
(
3
));
pool
.insert
(
block1
);
// Sleep for FIFO timestamp uniqueness (HashMap backend)
std
::
thread
::
sleep
(
std
::
time
::
Duration
::
from_millis
(
2
));
pool
.insert
(
block3
);
assert_eq!
(
pool
.len
(),
2
);
let
missing
=
nonexistent_hash
();
// scan_blocks should NOT stop on miss — should find both hash1 and hash3
let
found
=
pool
.scan_blocks
(
&
[
seq_hash1
,
missing
,
seq_hash3
],
true
);
assert_eq!
(
found
.len
(),
2
,
"scan_blocks should find both blocks, skipping the miss"
);
let
found_hashes
:
Vec
<
_
>
=
found
.iter
()
.map
(|(
h
,
_
)|
*
h
)
.collect
();
assert
!
(
found_hashes
.contains
(
&
seq_hash1
));
assert
!
(
found_hashes
.contains
(
&
seq_hash3
));
// Both blocks were removed from the pool
assert_eq!
(
pool
.len
(),
0
);
// RAII return: dropping the found blocks should return them
drop
(
found
);
assert_eq!
(
pool
.len
(),
2
);
}
#[test]
fn
test_allocate_all_blocks
()
{
let
(
pool
,
reset_pool
)
=
create_test_pool
();
let
(
block1
,
_
)
=
create_registered_block
::
<
TestMeta
>
(
1
,
&
tokens_for_id
(
1
));
let
(
block2
,
_
)
=
create_registered_block
::
<
TestMeta
>
(
2
,
&
tokens_for_id
(
2
));
let
(
block3
,
_
)
=
create_registered_block
::
<
TestMeta
>
(
3
,
&
tokens_for_id
(
3
));
pool
.insert
(
block1
);
pool
.insert
(
block2
);
pool
.insert
(
block3
);
assert_eq!
(
pool
.len
(),
3
);
let
mutable_blocks
=
pool
.allocate_all_blocks
();
assert_eq!
(
mutable_blocks
.len
(),
3
);
assert_eq!
(
pool
.len
(),
0
);
// Verify they are MutableBlocks by checking block_id
for
block
in
&
mutable_blocks
{
let
_
id
=
block
.block_id
();
}
// Drop them — they should return to the reset pool
drop
(
mutable_blocks
);
// 10 original reset blocks + 3 returned = 13
assert_eq!
(
reset_pool
.available_blocks
(),
13
);
}
}
lib/kvbm-logical/src/pools/mod.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Block pool RAII guards and allocation traits for thread-safe block management.
//!
//! This module provides:
//! - Type-safe RAII guards (MutableBlock, CompleteBlock, ImmutableBlock) for automatic resource cleanup
//! - ResetPool: Pool for mutable blocks in reset state
//! - InactivePool: Pool for inactive immutable registered blocks
//! - BlockRegistry: Global registry for block deduplication via weak references
//! - Pluggable allocation and reuse policies
mod
active
;
mod
inactive
;
mod
reset
;
#[cfg(test)]
pub
mod
tests
;
#[cfg(test)]
mod
block_proptest
;
pub
(
crate
)
use
active
::
ActivePool
;
pub
(
crate
)
use
inactive
::
backends
;
pub
(
crate
)
use
inactive
::{
InactivePool
,
InactivePoolBackend
};
pub
(
crate
)
use
reset
::
ResetPool
;
// Re-export RAII guards from guards module
use
crate
::
blocks
::{
Block
,
BlockId
,
BlockMetadata
,
ImmutableBlock
,
MutableBlock
,
PrimaryBlock
,
RegisteredBlock
,
state
::{
Registered
,
Reset
},
};
pub
(
crate
)
use
crate
::
SequenceHash
;
pub
(
crate
)
trait
BlockAllocator
<
T
:
BlockMetadata
>
{
// fn new(blocks: Vec<Block<T, Reset>>) -> Arc<Self>
// where
// Self: Sized;
/// Insert a block into the pool
fn
insert
(
&
mut
self
,
block
:
Block
<
T
,
Reset
>
);
/// Acquire the first block to be reused
fn
pop
(
&
mut
self
)
->
Option
<
Block
<
T
,
Reset
>>
;
/// Get the number of available blocks
fn
len
(
&
self
)
->
usize
;
/// Check if the pool is empty
fn
is_empty
(
&
self
)
->
bool
{
self
.len
()
==
0
}
}
#[expect(dead_code)]
pub
(
crate
)
trait
BlockMatcher
<
T
:
BlockMetadata
>
{
fn
find_match
(
&
self
,
seq_hash
:
SequenceHash
)
->
Option
<
ImmutableBlock
<
T
>>
;
}
// Re-export block duplication policy
pub
use
crate
::
blocks
::
BlockDuplicationPolicy
;
// Re-export reuse policy from inactive backends
pub
use
inactive
::
backends
::{
ReusePolicy
,
ReusePolicyError
};
// Re-export the new RAII guard types - no need to re-export here since they're defined in this module
/// A block that is free and available for allocation
/// This block must be in a Registered state and have a valid sequence hash
#[derive(Debug,
Clone,
Copy,
PartialEq,
Eq)]
pub
struct
InactiveBlock
{
pub
block_id
:
BlockId
,
pub
seq_hash
:
SequenceHash
,
}
// RegisteredPool implementation moved to registered.rs
lib/kvbm-logical/src/pools/reset.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Thread-safe pool for mutable blocks in reset state with pluggable allocation strategies.
//!
//! The ResetPool manages blocks available for allocation, using:
//! - Pluggable BlockAllocator for flexible allocation strategies
//! - RAII MutableBlock guards for automatic return
//! - Thread-safe access via parking_lot::Mutex
use
crate
::
BlockId
;
use
crate
::
metrics
::
BlockPoolMetrics
;
use
super
::{
Block
,
BlockAllocator
,
BlockMetadata
,
MutableBlock
,
Reset
};
use
parking_lot
::
Mutex
;
use
std
::{
collections
::
VecDeque
,
sync
::
Arc
};
pub
(
crate
)
struct
ResetPool
<
T
>
{
block_allocator
:
Arc
<
Mutex
<
dyn
BlockAllocator
<
T
>
+
Send
+
Sync
>>
,
return_fn
:
Arc
<
dyn
Fn
(
Block
<
T
,
Reset
>
)
+
Send
+
Sync
>
,
block_size
:
usize
,
metrics
:
Option
<
Arc
<
BlockPoolMetrics
>>
,
}
impl
<
T
:
BlockMetadata
>
ResetPool
<
T
>
{
pub
(
crate
)
fn
new
(
blocks
:
Vec
<
Block
<
T
,
Reset
>>
,
block_size
:
usize
,
metrics
:
Option
<
Arc
<
BlockPoolMetrics
>>
,
)
->
Self
{
let
allocator
=
DequeBlockAllocator
::
new
();
Self
::
from_block_allocator
(
allocator
,
blocks
,
block_size
,
metrics
)
}
pub
(
crate
)
fn
from_block_allocator
(
mut
allocator
:
impl
BlockAllocator
<
T
>
+
Send
+
Sync
+
'static
,
blocks
:
Vec
<
Block
<
T
,
Reset
>>
,
block_size
:
usize
,
metrics
:
Option
<
Arc
<
BlockPoolMetrics
>>
,
)
->
Self
{
for
(
i
,
block
)
in
blocks
.iter
()
.enumerate
()
{
if
block
.block_id
()
!=
i
as
BlockId
{
panic!
(
"Block ids must be monotonically increasing starting at 0"
);
}
}
for
block
in
blocks
{
allocator
.insert
(
block
);
}
let
block_allocator
=
Arc
::
new
(
Mutex
::
new
(
allocator
));
let
allocator_clone
=
block_allocator
.clone
();
let
metrics_clone
=
metrics
.clone
();
let
return_fn
=
Arc
::
new
(
move
|
block
:
Block
<
T
,
Reset
>
|
{
allocator_clone
.lock
()
.insert
(
block
);
if
let
Some
(
ref
m
)
=
metrics_clone
{
m
.inc_reset_pool_size
();
}
});
Self
{
block_allocator
,
return_fn
,
block_size
,
metrics
,
}
}
/// Tries to allocate upto `count` blocks from the pool.
/// Will return less than `count` blocks if the pool has less than `count` blocks available.
pub
(
crate
)
fn
allocate_blocks
(
&
self
,
count
:
usize
)
->
Vec
<
MutableBlock
<
T
>>
{
let
mut
blocks
=
Vec
::
with_capacity
(
count
);
let
mut
allocator
=
self
.block_allocator
.lock
();
let
available_count
=
std
::
cmp
::
min
(
count
,
allocator
.len
());
for
_
in
0
..
available_count
{
if
let
Some
(
ref
m
)
=
self
.metrics
{
m
.dec_reset_pool_size
();
}
blocks
.push
(
MutableBlock
::
new
(
allocator
.pop
()
.unwrap
(),
self
.return_fn
.clone
(),
self
.metrics
.clone
(),
));
}
blocks
}
/// Get the number of available blocks
#[allow(dead_code)]
pub
(
crate
)
fn
available_blocks
(
&
self
)
->
usize
{
self
.block_allocator
.lock
()
.len
()
}
pub
(
crate
)
fn
len
(
&
self
)
->
usize
{
self
.block_allocator
.lock
()
.len
()
}
/// Check if the pool is empty
#[allow(dead_code)]
pub
(
crate
)
fn
is_empty
(
&
self
)
->
bool
{
self
.block_allocator
.lock
()
.is_empty
()
}
/// Create a return function for blocks to return to this pool
/// This allows other pools to create MutableBlocks that return here
pub
(
crate
)
fn
return_fn
(
&
self
)
->
Arc
<
dyn
Fn
(
Block
<
T
,
Reset
>
)
+
Send
+
Sync
>
{
self
.return_fn
.clone
()
}
/// Get the expected block size for this pool
pub
(
crate
)
fn
block_size
(
&
self
)
->
usize
{
self
.block_size
}
}
#[derive(Debug)]
pub
(
crate
)
struct
DequeBlockAllocator
<
T
:
BlockMetadata
>
{
blocks
:
VecDeque
<
Block
<
T
,
Reset
>>
,
}
impl
<
T
:
BlockMetadata
>
Default
for
DequeBlockAllocator
<
T
>
{
fn
default
()
->
Self
{
Self
::
new
()
}
}
impl
<
T
:
BlockMetadata
>
DequeBlockAllocator
<
T
>
{
pub
(
crate
)
fn
new
()
->
Self
{
Self
{
blocks
:
VecDeque
::
new
(),
}
}
}
impl
<
T
:
BlockMetadata
>
BlockAllocator
<
T
>
for
DequeBlockAllocator
<
T
>
{
fn
insert
(
&
mut
self
,
block
:
Block
<
T
,
Reset
>
)
{
self
.blocks
.push_back
(
block
);
}
fn
pop
(
&
mut
self
)
->
Option
<
Block
<
T
,
Reset
>>
{
self
.blocks
.pop_front
()
}
fn
len
(
&
self
)
->
usize
{
self
.blocks
.len
()
}
}
#[cfg(test)]
mod
tests
{
use
super
::
*
;
use
crate
::
testing
::
TestMeta
;
fn
create_test_blocks
(
count
:
usize
)
->
Vec
<
Block
<
TestMeta
,
Reset
>>
{
(
0
..
count
as
BlockId
)
.map
(|
id
|
Block
::
new
(
id
,
4
))
.collect
()
}
#[test]
fn
test_mutable_block_raii_return
()
{
let
blocks
=
create_test_blocks
(
3
);
let
pool
=
ResetPool
::
new
(
blocks
,
4
,
None
);
assert_eq!
(
pool
.len
(),
3
);
{
let
allocated
=
pool
.allocate_blocks
(
2
);
assert_eq!
(
allocated
.len
(),
2
);
assert_eq!
(
pool
.len
(),
1
);
}
assert_eq!
(
pool
.len
(),
3
);
}
#[test]
fn
test_pool_allocation_and_return_cycle
()
{
let
blocks
=
create_test_blocks
(
5
);
let
pool
=
ResetPool
::
new
(
blocks
,
4
,
None
);
for
_
in
0
..
3
{
assert_eq!
(
pool
.len
(),
5
);
{
let
allocated
=
pool
.allocate_blocks
(
2
);
assert_eq!
(
allocated
.len
(),
2
);
assert_eq!
(
pool
.len
(),
3
);
}
assert_eq!
(
pool
.len
(),
5
);
}
}
}
lib/kvbm-logical/src/pools/tests.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Shared test utilities and fixtures for block pool testing.
use
super
::
super
::{
blocks
::{
state
::
*
,
*
},
pools
::
*
,
testing
::{
self
,
TestMeta
},
};
/// Re-export TestMeta as TestData for backward compatibility
pub
type
TestData
=
TestMeta
;
#[cfg(test)]
#[allow(unused,
dead_code)]
pub
(
crate
)
mod
fixtures
{
use
super
::
*
;
use
dynamo_tokens
::
TokenBlock
;
use
std
::
sync
::
Arc
;
// Re-export from testing module with TestData specialization
pub
use
super
::
testing
::
tokens_for_id
;
pub
fn
create_reset_block
(
id
:
BlockId
)
->
Block
<
TestData
,
Reset
>
{
testing
::
create_reset_block
::
<
TestData
>
(
id
,
4
)
}
pub
fn
create_reset_blocks
(
count
:
usize
)
->
Vec
<
Block
<
TestData
,
Reset
>>
{
testing
::
create_reset_blocks
::
<
TestData
>
(
count
,
4
)
}
pub
fn
create_token_block
(
tokens
:
&
[
u32
],
block_size
:
u32
)
->
TokenBlock
{
testing
::
create_test_token_block
(
tokens
,
block_size
)
}
pub
fn
create_complete_block
(
id
:
BlockId
,
tokens
:
&
[
u32
])
->
Block
<
TestData
,
Staged
>
{
testing
::
create_staged_block
::
<
TestData
>
(
id
,
tokens
)
}
pub
fn
create_registered_block
(
id
:
BlockId
,
tokens
:
&
[
u32
],
)
->
(
Block
<
TestData
,
Registered
>
,
SequenceHash
)
{
testing
::
create_registered_block
::
<
TestData
>
(
id
,
tokens
)
}
pub
fn
create_test_reset_pool
(
count
:
usize
)
->
ResetPool
<
TestData
>
{
testing
::
TestPoolSetupBuilder
::
default
()
.block_count
(
count
)
.build
()
.unwrap
()
.build_reset_pool
::
<
TestData
>
()
}
pub
fn
create_test_registered_pool
()
->
(
InactivePool
<
TestData
>
,
ResetPool
<
TestData
>
)
{
testing
::
TestPoolSetupBuilder
::
default
()
.build
()
.unwrap
()
.build_pools
::
<
TestData
>
()
}
/// Type alias for TestBlockBuilder specialized to TestData
pub
type
TestBlockBuilder
=
testing
::
TestBlockBuilder
<
TestData
>
;
/// Type alias for BlockSequenceBuilder specialized to TestData
pub
type
BlockSequenceBuilder
=
testing
::
BlockSequenceBuilder
<
TestData
>
;
}
#[cfg(test)]
use
fixtures
::
*
;
#[test]
fn
test_fill_iota_default_block_size
()
{
let
block
=
TestBlockBuilder
::
new
(
1
)
.fill_iota
(
100
)
.build_staged
();
assert_eq!
(
block
.block_id
(),
1
);
assert_eq!
(
block
.block_size
(),
4
);
}
#[test]
fn
test_fill_iota_custom_block_size
()
{
let
block
=
TestBlockBuilder
::
new
(
2
)
.with_block_size
(
8
)
.fill_iota
(
200
)
.build_staged
();
assert_eq!
(
block
.block_id
(),
2
);
assert_eq!
(
block
.block_size
(),
8
);
}
#[test]
fn
test_with_tokens_overrides_fill_iota
()
{
let
custom_tokens
=
vec!
[
99
,
98
,
97
,
96
];
let
block
=
TestBlockBuilder
::
new
(
3
)
.fill_iota
(
100
)
// This should be overridden
.with_tokens
(
custom_tokens
)
.build_staged
();
assert_eq!
(
block
.block_id
(),
3
);
assert_eq!
(
block
.block_size
(),
4
);
}
#[test]
fn
test_fill_iota_overrides_with_tokens
()
{
let
block
=
TestBlockBuilder
::
new
(
4
)
.with_tokens
(
vec!
[
1
,
2
,
3
,
4
])
// This should be overridden
.fill_iota
(
500
)
.build_staged
();
assert_eq!
(
block
.block_id
(),
4
);
assert_eq!
(
block
.block_size
(),
4
);
}
#[test]
fn
test_block_sequence_from_tokens
()
{
let
tokens
=
vec!
[
100
,
101
,
102
,
103
,
104
,
105
,
106
,
107
];
// 2 blocks of size 4
let
blocks
=
BlockSequenceBuilder
::
from_tokens
(
tokens
)
.with_block_size
(
4
)
.with_salt
(
42
)
.build
();
assert_eq!
(
blocks
.len
(),
2
);
assert_eq!
(
blocks
[
0
]
.0
.block_id
(),
0
);
assert_eq!
(
blocks
[
1
]
.0
.block_id
(),
1
);
assert_eq!
(
blocks
[
0
]
.0
.block_size
(),
4
);
assert_eq!
(
blocks
[
1
]
.0
.block_size
(),
4
);
}
#[test]
fn
test_block_sequence_individual_mode
()
{
let
blocks
=
BlockSequenceBuilder
::
new
()
.add_block_with
(
1
,
|
b
|
b
.fill_iota
(
100
))
.add_block_with
(
2
,
|
b
|
b
.fill_iota
(
200
))
.add_block
(
3
)
.build
();
assert_eq!
(
blocks
.len
(),
3
);
assert_eq!
(
blocks
[
0
]
.0
.block_id
(),
1
);
assert_eq!
(
blocks
[
1
]
.0
.block_id
(),
2
);
assert_eq!
(
blocks
[
2
]
.0
.block_id
(),
3
);
}
#[test]
#[should_panic(expected
=
"Token count 7 must be divisible by block size 4"
)]
fn
test_block_sequence_invalid_token_count
()
{
let
tokens
=
vec!
[
1
,
2
,
3
,
4
,
5
,
6
,
7
];
// 7 tokens, not divisible by 4
BlockSequenceBuilder
::
from_tokens
(
tokens
)
.with_block_size
(
4
)
.build
();
}
#[test]
fn
test_block_sequence_custom_block_size
()
{
let
tokens
:
Vec
<
u32
>
=
(
0
..
16
)
.collect
();
// 2 blocks of size 8
let
blocks
=
BlockSequenceBuilder
::
from_tokens
(
tokens
)
.with_block_size
(
8
)
.build
();
assert_eq!
(
blocks
.len
(),
2
);
assert_eq!
(
blocks
[
0
]
.0
.block_size
(),
8
);
assert_eq!
(
blocks
[
1
]
.0
.block_size
(),
8
);
}
#[test]
fn
test_mutable_block_complete_error_returns_block
()
{
use
crate
::
blocks
::
BlockError
;
let
reset_pool
=
create_test_reset_pool
(
5
);
let
mut
mutable_blocks
=
reset_pool
.allocate_blocks
(
1
);
let
mutable_block
=
mutable_blocks
.pop
()
.unwrap
();
let
original_block_id
=
mutable_block
.block_id
();
// block_size is 4, but token block has 8 tokens
let
big_token_block
=
testing
::
create_test_token_block
(
&
[
1
,
2
,
3
,
4
,
5
,
6
,
7
,
8
],
8
);
let
result
=
mutable_block
.complete
(
&
big_token_block
);
assert
!
(
result
.is_err
());
match
result
{
Err
(
BlockError
::
BlockSizeMismatch
{
expected
,
actual
,
block
:
recovered_block
,
})
=>
{
assert_eq!
(
expected
,
4
);
assert_eq!
(
actual
,
8
);
// Block is recoverable from the error
assert_eq!
(
recovered_block
.block_id
(),
original_block_id
);
}
_
=>
panic!
(
"Expected BlockSizeMismatch error"
),
}
}
#[test]
fn
test_mutable_block_stage_and_debug
()
{
let
reset_pool
=
create_test_reset_pool
(
5
);
let
mut
mutable_blocks
=
reset_pool
.allocate_blocks
(
1
);
let
mutable_block
=
mutable_blocks
.pop
()
.unwrap
();
// Exercise Debug for MutableBlock
let
debug_str
=
format!
(
"{:?}"
,
mutable_block
);
assert
!
(
debug_str
.contains
(
"MutableBlock"
));
// Exercise the `stage` method (bypass block_size check)
let
seq_hash
=
crate
::
KvbmSequenceHashProvider
::
kvbm_sequence_hash
(
&
testing
::
create_test_token_block
(
&
[
10
,
11
,
12
,
13
],
4
),
);
let
complete_block
=
mutable_block
.stage
(
seq_hash
,
4
)
.expect
(
"block size should match"
);
assert_eq!
(
complete_block
.sequence_hash
(),
seq_hash
);
}
#[test]
fn
test_complete_block_reset
()
{
let
reset_pool
=
create_test_reset_pool
(
5
);
let
mut
mutable_blocks
=
reset_pool
.allocate_blocks
(
1
);
let
mutable_block
=
mutable_blocks
.pop
()
.unwrap
();
let
original_block_id
=
mutable_block
.block_id
();
let
token_block
=
create_token_block
(
&
[
10
,
11
,
12
,
13
],
4
);
let
complete_block
=
mutable_block
.complete
(
&
token_block
)
.expect
(
"Should complete"
);
assert_eq!
(
complete_block
.block_id
(),
original_block_id
);
// Reset the complete block back to a mutable block
let
reset_mutable
=
complete_block
.reset
();
assert_eq!
(
reset_mutable
.block_id
(),
original_block_id
);
}
#[test]
fn
test_immutable_block_downgrade_and_upgrade
()
{
let
manager
=
testing
::
create_test_manager
::
<
TestData
>
(
10
);
let
token_block
=
testing
::
create_iota_token_block
(
100
,
4
);
let
seq_hash
=
crate
::
KvbmSequenceHashProvider
::
kvbm_sequence_hash
(
&
token_block
);
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete"
);
let
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
let
immutable_block
=
immutable_blocks
.into_iter
()
.next
()
.unwrap
();
// Check accessors
assert_eq!
(
immutable_block
.sequence_hash
(),
seq_hash
);
let
_
block_id
=
immutable_block
.block_id
();
let
_
handle
=
immutable_block
.registration_handle
();
assert
!
(
immutable_block
.use_count
()
>=
1
);
// Downgrade to WeakBlock
let
weak_block
=
immutable_block
.downgrade
();
assert_eq!
(
weak_block
.sequence_hash
(),
seq_hash
);
// Upgrade while original is alive — should succeed via direct Weak path
let
upgraded
=
weak_block
.upgrade
()
.expect
(
"Should upgrade while original alive"
);
assert_eq!
(
upgraded
.sequence_hash
(),
seq_hash
);
assert_eq!
(
upgraded
.block_id
(),
immutable_block
.block_id
());
}
#[test]
fn
test_weak_block_upgrade_via_upgrade_fn
()
{
let
manager
=
testing
::
create_test_manager
::
<
TestData
>
(
10
);
let
token_block
=
testing
::
create_iota_token_block
(
200
,
4
);
let
seq_hash
=
crate
::
KvbmSequenceHashProvider
::
kvbm_sequence_hash
(
&
token_block
);
// Create weak block, then drop the original
let
weak_block
=
{
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete"
);
let
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
let
immutable_block
=
immutable_blocks
.into_iter
()
.next
()
.unwrap
();
immutable_block
.downgrade
()
};
// original dropped — block returns to inactive pool
// Upgrade should succeed via the upgrade_fn path (finds block in inactive pool)
let
upgraded
=
weak_block
.upgrade
()
.expect
(
"upgrade should succeed via upgrade_fn"
);
assert_eq!
(
upgraded
.sequence_hash
(),
seq_hash
);
}
#[test]
fn
test_immutable_and_weak_block_debug
()
{
let
manager
=
testing
::
create_test_manager
::
<
TestData
>
(
10
);
let
token_block
=
testing
::
create_iota_token_block
(
300
,
4
);
let
_
seq_hash
=
crate
::
KvbmSequenceHashProvider
::
kvbm_sequence_hash
(
&
token_block
);
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete"
);
let
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
let
immutable_block
=
immutable_blocks
.into_iter
()
.next
()
.unwrap
();
// Exercise Debug for ImmutableBlock
let
debug_str
=
format!
(
"{:?}"
,
immutable_block
);
assert
!
(
debug_str
.contains
(
"ImmutableBlock"
));
// Exercise Debug for WeakBlock
let
weak_block
=
immutable_block
.downgrade
();
let
weak_debug_str
=
format!
(
"{:?}"
,
weak_block
);
assert
!
(
weak_debug_str
.contains
(
"WeakBlock"
));
}
#[test]
fn
test_weak_block_upgrade_fails_when_evicted
()
{
let
manager
=
testing
::
create_test_manager
::
<
TestData
>
(
10
);
let
token_block
=
testing
::
create_test_token_block
(
&
[
999
,
998
,
997
,
996
],
4
);
let
weak_block
=
{
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
token_block
)
.expect
(
"Should complete"
);
let
immutable_blocks
=
manager
.register_blocks
(
vec!
[
complete_block
]);
let
immutable_block
=
immutable_blocks
.into_iter
()
.next
()
.unwrap
();
immutable_block
.downgrade
()
};
// Fill up the pool with other blocks to force eviction of original
for
i
in
0
..
10
{
let
tokens
=
vec!
[
1000
+
i
,
1001
+
i
,
1002
+
i
,
1003
+
i
];
let
tb
=
testing
::
create_test_token_block
(
&
tokens
,
4
);
let
mutable_blocks
=
manager
.allocate_blocks
(
1
)
.expect
(
"Should allocate"
);
let
complete_block
=
mutable_blocks
.into_iter
()
.next
()
.unwrap
()
.complete
(
&
tb
)
.expect
(
"Should complete"
);
let
_
immutable
=
manager
.register_blocks
(
vec!
[
complete_block
]);
}
// Upgrade should fail since original was evicted
let
result
=
weak_block
.upgrade
();
assert
!
(
result
.is_none
(),
"Upgrade should fail after eviction"
);
}
lib/kvbm-logical/src/pubsub.rs
0 → 100644
View file @
94fad479
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Publisher trait for distributed messaging.
//!
//! This module provides the `Publisher` trait used by the event pipeline.
//! Concrete implementations (NATS, Stub) remain in `dynamo-kvbm`.
use
anyhow
::
Result
;
use
bytes
::
Bytes
;
use
futures
::
future
::
BoxFuture
;
/// Publisher trait for sending messages to subjects.
///
/// Publishers are responsible for sending messages to named subjects.
/// Messages are delivered to all subscribers matching the subject pattern.
pub
trait
Publisher
:
Send
+
Sync
{
/// Publish a message to a subject.
///
/// This queues the message for delivery and returns immediately.
/// Use [`flush`](Publisher::flush) to ensure delivery.
fn
publish
(
&
self
,
subject
:
&
str
,
payload
:
Bytes
)
->
Result
<
()
>
;
/// Flush pending messages to ensure delivery.
///
/// Returns when all previously published messages have been acknowledged
/// by the messaging system.
fn
flush
(
&
self
)
->
BoxFuture
<
'static
,
Result
<
()
>>
;
}
Prev
1
2
3
Next
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment