dynamo-llm 1.0.2

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use std::sync::atomic::AtomicU64;

use crate::block_manager::block::{BlockState, locality::LocalityProvider};

use super::*;
use priority_key::PriorityKey;

use tracing::instrument;

#[derive(Default)]
pub struct InactiveBlockPool<S: Storage, L: LocalityProvider, M: BlockMetadata> {
    // Direct lookup by sequence_hash.
    lookup_map: HashMap<SequenceHash, Block<S, L, M>>,

    // Ordered by timestamp (oldest first)
    priority_set: BTreeSet<PriorityKey<M>>,

    // Fully Uninitialized
    uninitialized_set: VecDeque<Block<S, L, M>>,

    // Return Tick
    return_tick: u64,

    // Total blocks counter
    total_blocks: Arc<AtomicU64>,

    // Inactive blocks
    available_blocks: Arc<AtomicU64>,
}

impl<S: Storage, L: LocalityProvider, M: BlockMetadata> InactiveBlockPool<S, L, M> {
    /// Creates a new, empty [`InactiveBlockPool`].
    ///
    /// # Returns
    ///
    /// A new instance of [`InactiveBlockPool`].
    pub(crate) fn new() -> Self {
        Self {
            lookup_map: HashMap::new(),
            priority_set: BTreeSet::new(),
            uninitialized_set: VecDeque::new(),
            return_tick: 0,
            total_blocks: Arc::new(AtomicU64::new(0)),
            available_blocks: Arc::new(AtomicU64::new(0)),
        }
    }

    /// Returns a counter for the number of available blocks.
    ///
    /// # Returns
    ///
    /// A counter for the number of available blocks as an [`Arc<AtomicU64>`].
    pub fn available_blocks_counter(&self) -> Arc<AtomicU64> {
        self.available_blocks.clone()
    }

    /// Returns a counter for the total number of blocks.
    ///
    /// # Returns
    ///
    /// A counter for the total number of blocks as an [`Arc<AtomicU64>`].
    pub fn total_blocks_counter(&self) -> Arc<AtomicU64> {
        self.total_blocks.clone()
    }

    /// Returns the total number of blocks managed by this pool (both available and acquired).
    ///
    /// # Returns
    ///
    /// The total block count as a [`u64`].
    pub fn total_blocks(&self) -> u64 {
        self.total_blocks.load(Ordering::Relaxed)
    }

    /// Returns the number of blocks currently available in the pool.
    ///
    /// This is calculated dynamically based on the blocks in the [`uninitialized_set`]
    /// and the [`lookup_map`].
    ///
    /// # Returns
    ///
    /// The available block count as a [`u64`].
    pub fn available_blocks(&self) -> u64 {
        self.uninitialized_set.len() as u64 + self.lookup_map.len() as u64
    }

    /// Inserts a block into the pool using its sequence hash for potential reuse.
    ///
    /// If an entry with the same sequence hash already exists in the [`lookup_map`]
    /// the block is reset and moved to the [`uninitialized_set`].
    /// Otherwise, the block is added to the [`lookup_map`].
    ///
    /// # Arguments
    ///
    /// * `block` - The block to insert ([`Block<T, M>`]).
    /// * `sequence_hash` - The sequence hash associated with the block's content ([`SequenceHash`]).
    #[instrument(level = "trace", skip(self, block), fields(sequence_hash = ?sequence_hash))]
    fn insert_with_sequence_hash(&mut self, block: Block<S, L, M>, sequence_hash: SequenceHash) {
        let priority_key = PriorityKey::new(block.metadata().clone(), sequence_hash);
        if self.priority_set.contains(&priority_key) {
            tracing::trace!(
                "multiple entries with the same sequence hash, resetting block and inserting into uninitialized set"
            );
            let mut block = block;
            block.reset();
            self.uninitialized_set.push_back(block);
        } else {
            tracing::trace!("inserting block to map and priority set");

            self.priority_set.insert(priority_key);
            self.lookup_map.insert(sequence_hash, block);
        }
    }

    /// Internal helper to insert a block into the appropriate internal collection
    /// based on its current state.
    ///
    /// - [`BlockState::Reset`], [`BlockState::Partial`], [`BlockState::Complete`] states result in the block being reset and added
    ///   to the `uninitialized_set`.
    /// - [`BlockState::Registered`] state results in the block being added via [`insert_with_sequence_hash`].
    ///
    /// # Arguments
    ///
    /// * `block` - The block to insert ([`Block<S, M>`]).
    #[instrument(level = "trace", skip(self, block), fields(block_state = ?block.state()))]
    fn insert(&mut self, block: Block<S, L, M>) {
        tracing::trace!("Inserting block into available pool");

        // If we already have an entry for this sequence hash or the block is reset,
        // we need to move it to the uninitialized set
        match block.state() {
            BlockState::Reset | BlockState::Partial(_) | BlockState::Complete(_) => {
                let mut block = block;
                block.reset();
                self.uninitialized_set.push_back(block);
            }
            BlockState::Registered(state, _) => {
                let sequence_hash = state.sequence_hash();
                self.insert_with_sequence_hash(block, sequence_hash);
            }
        }

        self.available_blocks.fetch_add(1, Ordering::Relaxed);
    }

    /// Adds multiple blocks to the pool.
    ///
    /// Each block is reset before being inserted. The total block count is updated.
    ///
    /// # Arguments
    ///
    /// * `blocks` - A vector of blocks ([`Block<T, M>`]) to add.
    #[instrument(level = "debug", skip(self, blocks))]
    pub fn add_blocks(&mut self, blocks: Vec<Block<S, L, M>>) {
        let count = blocks.len();
        tracing::debug!(count, "Adding blocks to pool");

        for (i, mut block) in blocks.into_iter().enumerate() {
            tracing::trace!(current = i + 1, total = count, "Processing block");
            block.reset();
            self.insert(block);
        }

        self.total_blocks.fetch_add(count as u64, Ordering::Relaxed);
    }

    /// Adds multiple blocks to the pool.
    ///
    /// The state of the blocks are not reset.
    ///
    /// # Arguments
    ///
    /// * `blocks` - A vector of blocks ([`Block<T, M>`]) to add.
    #[instrument(level = "debug", skip(self, blocks))]
    pub fn add_blocks_with_state(&mut self, blocks: Vec<Block<S, L, M>>) {
        let count = blocks.len();
        tracing::debug!(count, "Adding blocks to pool");
        self.total_blocks.fetch_add(count as u64, Ordering::Relaxed);
        // self.available_blocks += count as u64;
        self.return_blocks(blocks);
    }

    /// Returns a single block to the pool.
    ///
    /// Increments the internal return tick, updates the block's metadata,
    /// and inserts the block back into the appropriate internal collection.
    ///
    /// # Arguments
    ///
    /// * `block` - The block ([`Block<S, M>`]) to return.
    #[instrument(level = "debug", skip(self, block))]
    pub fn return_block(&mut self, mut block: Block<S, L, M>) {
        // increment the return tick
        self.return_tick += 1;

        // update the metadata
        block.metadata_on_returned(self.return_tick);

        // insert the block into the pool
        self.insert(block);

        // self.available_blocks += 1;
    }

    /// Returns multiple blocks to the pool.
    ///
    /// Iterates through the blocks in order and calls
    /// `return_block` for each one.
    ///
    /// # Arguments
    ///
    /// * `blocks` - A vector of blocks ([`Block<T, M>`]) to return.
    #[instrument(level = "debug", skip(self, blocks))]
    pub fn return_blocks(&mut self, blocks: Vec<Block<S, L, M>>) {
        let count = blocks.len();
        tracing::debug!(count, "Returning blocks to pool");
        // return the block to the pool from tail to head
        for (i, block) in blocks.into_iter().enumerate() {
            tracing::trace!(current = i + 1, total = count, "Returning block");
            // Note: return_block has its own instrumentation
            self.return_block(block);
        }
    }

    /// Attempts to remove and return a block associated with the given sequence hash
    /// from the [`lookup_map`] and [`priority_set`].
    ///
    /// # Arguments
    ///
    /// * `sequence_hash` - The sequence hash ([`SequenceHash`]) of the block to take.
    ///
    /// # Returns
    ///
    /// An [`Option<Block<S, M>>`] containing the block if found, otherwise `None`.
    #[instrument(level = "trace", skip(self), fields(sequence_hash = ?sequence_hash))]
    fn take_with_sequence_hash(&mut self, sequence_hash: SequenceHash) -> Option<Block<S, L, M>> {
        match self.lookup_map.remove(&sequence_hash) {
            Some(block) => {
                // Remove from priority set.
                let priority_key = PriorityKey::new(block.metadata().clone(), sequence_hash);
                // Remove from priority set, if it exists.
                self.priority_set.remove(&priority_key);

                self.available_blocks.fetch_sub(1, Ordering::Relaxed);
                Some(block)
            }
            None => None,
        }
    }

    /// Attempts to find and take a block matching the given sequence hash.
    ///
    /// This is a convenience wrapper around `take_with_sequence_hash`.
    ///
    /// # Arguments
    ///
    /// * `sequence_hash` - The sequence hash ([`SequenceHash`]) to match.
    ///
    /// # Returns
    ///
    /// An [`Option<Block<S, M>>`] containing the block if found, otherwise `None`.
    #[instrument(level = "debug", skip(self), fields(sequence_hash = ?sequence_hash))]
    pub fn match_sequence_hash(&mut self, sequence_hash: SequenceHash) -> Option<Block<S, L, M>> {
        self.take_with_sequence_hash(sequence_hash)
    }

    /// Attempts to find and take multiple blocks matching a sequence of hashes.
    ///
    /// Iterates through the provided hashes and takes blocks using `take_with_sequence_hash`.
    /// Stops if a hash is not found.
    ///
    /// # Arguments
    ///
    /// * `sequence_hashes` - A vector of sequence hashes ([`SequenceHash`]) to match.
    ///
    /// # Returns
    ///
    /// A vector containing the blocks ([`Block<T, M>`]) that were successfully matched and taken.
    /// The vector may be shorter than `sequence_hashes` if not all hashes were found.
    #[instrument(level = "debug", skip(self, sequence_hashes), fields(num_hashes = sequence_hashes.len()))]
    pub fn match_sequence_hashes(
        &mut self,
        sequence_hashes: Vec<SequenceHash>,
    ) -> Vec<Block<S, L, M>> {
        let total_hashes = sequence_hashes.len();
        let mut matched_blocks = Vec::with_capacity(total_hashes);

        for (i, hash) in sequence_hashes.into_iter().enumerate() {
            tracing::trace!(current = i + 1, total = total_hashes, sequence_hash = ?hash, "Attempting to match sequence hash");
            // Note: take_with_sequence_hash has its own instrumentation
            if let Some(block) = self.take_with_sequence_hash(hash) {
                tracing::trace!(current = i + 1, total = total_hashes, sequence_hash = ?hash, "Matched sequence hash");
                matched_blocks.push(block);
            } else {
                tracing::trace!(current = i + 1, total = total_hashes, sequence_hash = ?hash, "Sequence hash not found, stopping match");
                break;
            }
        }

        matched_blocks
    }

    /// Attempts to find and take multiple blocks matching a sequence of `TokenBlock`s.
    ///
    /// Extracts sequence hashes from the [`TokenBlock`]s and calls [`take_with_sequence_hash`].
    /// Stops if a hash is not found.
    ///
    /// # Arguments
    ///
    /// * `token_blocks` - A slice of [`TokenBlock`]s to match.
    ///
    /// # Returns
    ///
    /// A vector containing the blocks ([`Block<T, M>`]) that were successfully matched and taken.
    /// The vector may be shorter than `token_blocks` if not all corresponding hashes were found.
    #[instrument(level = "debug", skip(self, token_blocks), fields(num_token_blocks = token_blocks.len()))]
    pub fn match_token_blocks(&mut self, token_blocks: &[TokenBlock]) -> Vec<Block<S, L, M>> {
        let total_blocks = token_blocks.len();
        let mut matched_blocks = Vec::with_capacity(total_blocks);

        tracing::debug!("Attempting to match {} token blocks", total_blocks);

        for (i, token_block) in token_blocks.iter().enumerate() {
            let sequence_hash = token_block.sequence_hash();
            tracing::trace!(sequence_hash = ?sequence_hash, "Attempting to match token block hash {}/{}", i + 1, total_blocks);
            if let Some(block) = self.take_with_sequence_hash(sequence_hash) {
                tracing::trace!(sequence_hash = ?sequence_hash, "Matched token block hash");
                matched_blocks.push(block);
            } else {
                tracing::trace!(sequence_hash = ?sequence_hash, "Token block hash not found, stopping match");
                break;
            }
        }

        tracing::debug!(
            "Matched {} of {} token blocks",
            matched_blocks.len(),
            total_blocks
        );

        matched_blocks
    }

    /// Acquires a single free block from the pool.
    ///
    /// Prioritizes blocks from the [`uninitialized_set`] first, then takes the
    /// lowest priority block from the [`priority_set`] (and [`lookup_map`]).
    /// If a block is taken from the priority set, it is reset.
    ///
    /// # Returns
    ///
    /// An [`Option<Block<T, M>>`] containing a free block if available, otherwise `None`.
    ///
    /// # Panics
    ///
    /// This function can panic if there is an inconsistency between the [`priority_set`]
    /// and [`lookup_map`] (i.e., a key exists in the set but not the map). This indicates
    /// a bug in the pool's internal logic.
    #[instrument(level = "debug", skip(self))]
    pub fn acquire_free_block(&mut self) -> Option<Block<S, L, M>> {
        // First try uninitialized blocks - these are often part of sequences
        // that have been arranged in the correct order
        if let Some(mut block) = self.uninitialized_set.pop_front() {
            tracing::trace!("Acquired uninitialized block");
            self.return_tick += 1;
            block.metadata_on_acquired(self.return_tick);
            self.available_blocks.fetch_sub(1, Ordering::Relaxed);
            return Some(block);
        }

        // if we have blocks in the priority set, pop the first (it's sorted by priority)
        // a fatal error will occur if the block is not found in the lookup map
        if let Some(key) = self.priority_set.pop_first() {
            tracing::trace!("Acquired priority/registered block map; resetting block");
            match self.lookup_map.remove(&key.sequence_hash()) {
                Some(mut block) => {
                    block.reset();
                    self.return_tick += 1;
                    block.metadata_on_acquired(self.return_tick);
                    self.available_blocks.fetch_sub(1, Ordering::Relaxed);
                    Some(block)
                }
                None => {
                    panic!(
                        "Block from priority set not found in lookup map! Inconsistency detected."
                    );
                }
            }
        } else {
            // No blocks available in either set
            None
        }
    }

    /// Acquires a specified number of free blocks from the pool.
    ///
    /// Checks if enough blocks are available and then calls [`acquire_free_block`] repeatedly.
    ///
    /// # Arguments
    ///
    /// * `count` - The number of free blocks to acquire.
    ///
    /// # Returns
    ///
    /// A [`Result`] containing:
    /// - `Ok(Vec<Block<T, M>>)`: A vector of the acquired blocks if successful.
    /// - `Err(BlockPoolError::InsufficientBlocksAvailable)`: If the requested number
    ///   of blocks is not available, or if an inconsistency occurred during acquisition.
    ///
    /// # Panics
    ///
    /// This function can panic if [`acquire_free_block`] panics due to internal inconsistencies.
    #[instrument(level = "debug", skip(self))]
    pub fn acquire_free_blocks(
        &mut self,
        count: usize,
    ) -> Result<Vec<Block<S, L, M>>, BlockPoolError> {
        if count == 0 {
            return Ok(Vec::new());
        }

        let mut blocks = Vec::with_capacity(count);

        let available_now = self.uninitialized_set.len() + self.lookup_map.len();
        tracing::debug!(
            available_now,
            requested = count,
            "Attempting to acquire free blocks"
        );

        if count > available_now {
            tracing::debug!(
                available_now,
                requested = count,
                "Insufficient blocks available"
            );
            return Err(BlockPoolError::NotEnoughBlocksAvailable(
                count,
                available_now,
            ));
        }

        for i in 0..count {
            tracing::trace!(current = i + 1, total = count, "Acquiring free block");
            // Directly call the logic in acquire_free_block
            // Note: acquire_free_block has its own instrumentation
            if let Some(block) = self.acquire_free_block() {
                blocks.push(block);
            } else {
                // This should not happen if the initial check passed and there are no concurrent modifications.
                // If it does, it indicates an inconsistency or a logic error.
                tracing::error!(
                    requested = count,
                    acquired = blocks.len(),
                    available_at_start = available_now,
                    current_available = self.uninitialized_set.len() + self.lookup_map.len(),
                    "Insufficient blocks during acquisition loop despite initial check."
                );
                // Return the blocks acquired so far, or handle as an error.
                // For now, we break and return what we have, but decrementing 'available_blocks'
                // needs to account for the actual number acquired.
                // Consider returning an error or panicking in debug.
                break;
            }
        }

        let acquired_count = blocks.len();
        tracing::debug!(
            acquired_count,
            requested = count,
            "Finished acquiring blocks"
        );

        // Check if we got the requested number of blocks
        if acquired_count != count {
            // This path is taken if the loop broke early due to unexpected `None` from acquire_free_block
            // Return an error indicating partial success or failure
            // Depending on the desired behavior, you might return the partial list
            // or a more specific error.
            // For consistency with the original check, let's return an error if count wasn't met.
            return Err(BlockPoolError::NotEnoughBlocksAvailable(
                count,
                blocks.len(),
            ));
        }

        Ok(blocks)
    }

    /// Resets the pool to its initial state.
    ///
    /// This function will acquire all blocks, which will reset their state, then return them.
    ///
    /// A [`Result`] containing `Ok(())` if the reset was successful, otherwise an error.
    pub fn reset(&mut self) -> Result<(), BlockPoolError> {
        let total_blocks = self.total_blocks.load(Ordering::Relaxed);
        let available_blocks = self.available_blocks.load(Ordering::Relaxed);

        if total_blocks != available_blocks {
            return Err(BlockPoolError::ResetError(format!(
                "total blocks: {}, available blocks: {}",
                total_blocks, available_blocks
            )));
        }

        let blocks = self.acquire_free_blocks(total_blocks as usize)?;

        for block in blocks.into_iter() {
            self.return_block(block);
        }

        Ok(())
    }

    /// Returns the [`PoolStatus`] of the pool.
    pub fn status(&self) -> (usize, usize) {
        let inactive_blocks = self.priority_set.len();
        let empty_blocks = self.uninitialized_set.len();
        (inactive_blocks, empty_blocks)
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use crate::{
        block_manager::{
            block::{
                Blocks, PrivateBlockExt, locality::Local, registry::BlockRegistry,
                state::CompleteState,
            },
            events::NullEventManager,
            layout::{BlockLayout, FullyContiguous, LayoutConfigBuilder},
            storage::tests::{NullDeviceAllocator, NullDeviceStorage},
        },
        tokens::{Token, Tokens},
    };

    use super::*;

    #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Ord, PartialOrd)]
    pub struct TestMetadata {
        priority: u32,
        returned_tick: u64,
        acquired_tick: u64,
    }

    impl BlockMetadata for TestMetadata {
        fn on_acquired(&mut self, tick: u64) {
            self.acquired_tick = tick;
        }

        fn on_returned(&mut self, tick: u64) {
            self.returned_tick = tick;
        }

        fn reset_metadata(&mut self) {
            self.priority = 0;
        }

        fn offload_priority(&self) -> Option<u64> {
            Some(self.priority as u64)
        }

        fn with_priority(&self, priority: u32) -> Self {
            Self {
                priority,
                returned_tick: self.returned_tick,
                acquired_tick: self.acquired_tick,
            }
        }
    }

    type TestPriorityKey = PriorityKey<TestMetadata>;

    fn make_priority_key(
        priority: u32,
        returned_tick: u64,
        sequence_hash: SequenceHash,
    ) -> TestPriorityKey {
        TestPriorityKey::new(
            TestMetadata {
                priority,
                returned_tick,
                acquired_tick: 0,
            },
            sequence_hash,
        )
    }

    #[test]
    fn test_priority_key_ord() {
        let mut map = BTreeSet::new();

        let hash1 = SequenceHash::from(1u64);
        let hash2 = SequenceHash::from(2u64);
        let hash3 = SequenceHash::from(3u64);

        map.insert(make_priority_key(0, 2, hash1));
        map.insert(make_priority_key(1, 1, hash2));
        map.insert(make_priority_key(0, 3, hash3));

        // Test popping from the map to verify ordering
        let first_key = map.pop_first().unwrap();
        assert_eq!(first_key.metadata().priority, 0);
        assert_eq!(first_key.metadata().returned_tick, 2);
        assert_eq!(first_key.sequence_hash(), hash1);

        let second_key = map.pop_first().unwrap();
        assert_eq!(second_key.metadata().priority, 0);
        assert_eq!(second_key.metadata().returned_tick, 3);
        assert_eq!(second_key.sequence_hash(), hash3);

        let third_key = map.pop_first().unwrap();
        assert_eq!(third_key.metadata().priority, 1);
        assert_eq!(third_key.metadata().returned_tick, 1);
        assert_eq!(third_key.sequence_hash(), hash2);

        // Map should now be empty
        assert!(map.is_empty());
    }

    #[test]
    fn test_with_priority_updates_priority() {
        let metadata = TestMetadata {
            priority: 10,
            returned_tick: 100,
            acquired_tick: 50,
        };
        let updated = metadata.with_priority(80);

        assert_eq!(updated.priority, 80);
        assert_eq!(updated.returned_tick, 100); // preserved
        assert_eq!(updated.acquired_tick, 50); // preserved
    }

    #[test]
    fn test_with_priority_immutability() {
        let original = TestMetadata {
            priority: 35,
            returned_tick: 10,
            acquired_tick: 5,
        };
        let updated = original.with_priority(100);

        assert_eq!(original.priority, 35); // unchanged
        assert_eq!(updated.priority, 100);
    }

    #[test]
    fn test_with_priority_boundary_values() {
        let metadata = TestMetadata::default();

        assert_eq!(metadata.with_priority(0).priority, 0);
        assert_eq!(metadata.with_priority(100).priority, 100);
        assert_eq!(metadata.with_priority(u32::MAX).priority, u32::MAX);
    }

    // Helper function to create a sequence of tokens
    pub fn create_token_sequence(values: &[u32]) -> Tokens {
        let tokens: Vec<Token> = values.iter().map(|&v| Token::from(v)).collect();
        Tokens::from(tokens)
    }

    /// Creates a block collection with the given number of blocks.
    pub fn create_block_collection(
        num_blocks: usize,
    ) -> Blocks<impl BlockLayout<StorageType = NullDeviceStorage>, TestMetadata> {
        let config = LayoutConfigBuilder::default()
            .num_blocks(num_blocks)
            .num_layers(61)
            .outer_dim(1)
            .page_size(16)
            .inner_dim(576)
            .build()
            .unwrap();

        let layout = FullyContiguous::allocate(config, &NullDeviceAllocator)
            .expect("Failed to allocate layout/storage");

        Blocks::<_, TestMetadata>::new(layout, 42, 0).unwrap()
    }

    /// Creates a vector of Blocks from a token sequence and block size.
    /// Each block is initialized to the Complete state and then Registered.
    pub fn create_blocks(
        tokens: Tokens,
        block_size: u32,
        async_runtime: Handle,
    ) -> Vec<Block<NullDeviceStorage, Local, TestMetadata>> {
        let (token_blocks, _partial_token_block) =
            tokens.into_sequence(block_size, None).into_parts();
        let num_blocks = token_blocks.len();

        if num_blocks == 0 {
            return Vec::new();
        }

        let mut blocks = create_block_collection(num_blocks).into_blocks().unwrap();

        let event_manager = NullEventManager::new();
        let mut registry =
            BlockRegistry::new(event_manager, GlobalRegistry::default(), async_runtime);

        // Iterate through the generated TokenBlocks and the template Blocks,
        // setting the state and registering each one.
        for (block, token_block) in blocks.iter_mut().zip(token_blocks.into_iter()) {
            assert!(block.state().is_reset()); // Start with empty blocks
            block.update_state(BlockState::Complete(CompleteState::new(token_block)));
            block
                .register(&mut registry)
                .expect("Failed to register block in test helper");
            assert!(block.state().is_registered()); // Ensure registration worked
        }

        blocks
    }

    pub fn create_block_pool(
        num_blocks: usize,
    ) -> InactiveBlockPool<NullDeviceStorage, Local, TestMetadata> {
        let mut pool = InactiveBlockPool::new();
        let blocks = create_block_collection(num_blocks).into_blocks().unwrap();
        pool.add_blocks(blocks);

        pool
    }

    pub fn acquire_blocks(
        tokens: Tokens,
        block_size: u32,
        pool: &mut InactiveBlockPool<NullDeviceStorage, Local, TestMetadata>,
        async_runtime: Handle,
    ) -> (Vec<Block<NullDeviceStorage, Local, TestMetadata>>, usize) {
        let (mut token_blocks, _partial_token_block) =
            tokens.into_sequence(block_size, None).into_parts();

        let total_complete_blocks = token_blocks.len();

        // this will match the token_blocks to any matching blocks in the inactive pool
        // these blocks have the same sequence hash as the token_blocks, thus no updates are needed
        let mut matched_blocks = pool.match_token_blocks(&token_blocks);
        let matched_block_count = matched_blocks.len();

        let event_manager = NullEventManager::new();
        let mut registry =
            BlockRegistry::new(event_manager, GlobalRegistry::default(), async_runtime);

        // all matched blocks should be in the complete or registered state
        for block in &mut matched_blocks {
            assert!(block.state().is_registered());
        }

        // drain the matched blocks from the token_blocks
        token_blocks.drain(0..matched_block_count);

        assert_eq!(
            token_blocks.len() + matched_blocks.len(),
            total_complete_blocks
        );

        // try to acquire the remaining blocks
        let mut unmatched_blocks = pool.acquire_free_blocks(token_blocks.len()).unwrap();

        assert_eq!(unmatched_blocks.len(), token_blocks.len());

        for unmatched in &unmatched_blocks {
            assert!(unmatched.state().is_reset());
        }

        for (unmatched, token_block) in unmatched_blocks.iter_mut().zip(token_blocks.into_iter()) {
            assert!(unmatched.state().is_reset());
            unmatched.update_state(BlockState::Complete(CompleteState::new(token_block)));
            unmatched.register(&mut registry).unwrap();
            assert!(unmatched.state().is_registered());
        }

        let mut blocks = matched_blocks;
        blocks.extend(unmatched_blocks);
        (blocks, matched_block_count)
    }

    #[test]
    fn test_block_pool_lifecycle() {
        dynamo_runtime::logging::init();

        let async_runtime = tokio::runtime::Runtime::new().unwrap();

        const PAGE_SIZE: u32 = 2;

        let mut pool = create_block_pool(10);
        assert_eq!(pool.total_blocks(), 10);
        assert_eq!(pool.available_blocks(), 10);

        let blocks = pool.acquire_free_blocks(10).unwrap();
        assert_eq!(blocks.len(), 10);
        assert_eq!(pool.total_blocks(), 10);
        assert_eq!(pool.available_blocks(), 0);

        pool.return_blocks(blocks);

        assert_eq!(pool.total_blocks(), 10);
        assert_eq!(pool.available_blocks(), 10);
        assert_eq!(
            pool.available_blocks_counter().load(Ordering::Relaxed),
            pool.available_blocks()
        );

        let tokens = create_token_sequence(&[1, 2, 3, 4]);

        let (blocks, matched_block_count) = acquire_blocks(
            tokens.clone(),
            PAGE_SIZE,
            &mut pool,
            async_runtime.handle().clone(),
        );
        assert_eq!(blocks.len(), 2);
        assert_eq!(matched_block_count, 0);
        assert_eq!(pool.available_blocks(), 8);
        assert_eq!(
            pool.available_blocks_counter().load(Ordering::Relaxed),
            pool.available_blocks()
        );

        pool.return_blocks(blocks);

        assert_eq!(pool.total_blocks(), 10);
        assert_eq!(pool.available_blocks(), 10);
        assert_eq!(
            pool.available_blocks_counter().load(Ordering::Relaxed),
            pool.available_blocks()
        );

        let (blocks, matched_block_count) = acquire_blocks(
            tokens.clone(),
            PAGE_SIZE,
            &mut pool,
            async_runtime.handle().clone(),
        );
        assert_eq!(blocks.len(), 2);
        assert_eq!(matched_block_count, 2);
        assert_eq!(pool.available_blocks(), 8);
        assert_eq!(
            pool.available_blocks_counter().load(Ordering::Relaxed),
            pool.available_blocks()
        );

        pool.return_blocks(blocks);

        assert_eq!(pool.total_blocks(), 10);
        assert_eq!(pool.available_blocks(), 10);
        assert_eq!(
            pool.available_blocks_counter().load(Ordering::Relaxed),
            pool.available_blocks()
        );

        let blocks = pool.acquire_free_blocks(10).unwrap();
        for block in &blocks {
            assert!(block.state().is_reset());
        }
    }

    #[test]
    fn test_basic_sequence_matching() {
        let mut pool = InactiveBlockPool::new();

        let async_runtime = tokio::runtime::Runtime::new().unwrap();

        // Create a sequence of 4 tokens split into blocks of 2
        let sequence = create_token_sequence(&[1, 2, 3, 4]);
        let blocks = create_blocks(sequence, 2, async_runtime.handle().clone());
        assert_eq!(blocks.len(), 2);

        // Match the blocks in sequence
        let hashes: Vec<_> = blocks
            .iter()
            .map(|b| {
                b.sequence_hash()
                    .expect("Block should have a sequence hash in this test")
            })
            .collect();

        // Insert blocks into pool
        pool.add_blocks_with_state(blocks);

        assert_eq!(pool.total_blocks(), 2);
        assert_eq!(pool.available_blocks(), 2);
        assert_eq!(
            pool.available_blocks_counter().load(Ordering::Relaxed),
            pool.available_blocks()
        );

        // Match the blocks in sequence
        let matched = pool.match_sequence_hashes(hashes.clone());
        assert_eq!(matched.len(), 2);

        assert_eq!(pool.total_blocks(), 2);
        assert_eq!(pool.available_blocks(), 0);
        assert_eq!(
            pool.available_blocks_counter().load(Ordering::Relaxed),
            pool.available_blocks()
        );

        // Validate the blocks are in the correct order and match the sequence hashes
        assert_eq!(matched[0].sequence_hash().unwrap(), hashes[0]);
        assert_eq!(matched[1].sequence_hash().unwrap(), hashes[1]);

        // Return blocks in reverse order (tail to root)
        pool.return_blocks(matched);

        assert_eq!(pool.total_blocks(), 2);
        assert_eq!(pool.available_blocks(), 2);
        assert_eq!(
            pool.available_blocks_counter().load(Ordering::Relaxed),
            pool.available_blocks()
        );
    }

    /// Test that validates blocks allocated from the pool always have the default
    /// priority (0), regardless of what priority they had in a previous allocation.
    ///
    /// This test exposes a bug where blocks in Reset state that are returned to
    /// the pool retain their non-default priority when re-acquired, because the
    /// uninitialized_set path in acquire_free_block() does not call block.reset().
    #[test]
    fn test_allocated_blocks_have_default_priority() {
        let mut pool = create_block_pool(3);

        // Step 1: Acquire blocks (they come from uninitialized_set in Reset state)
        let mut blocks = pool.acquire_free_blocks(3).unwrap();
        assert_eq!(blocks.len(), 3);

        // Verify initial priority is 0 (default)
        for block in &blocks {
            assert_eq!(
                block.metadata().offload_priority(),
                Some(0),
                "Newly acquired block should have default priority"
            );
        }

        // Step 2: Set non-default priority on blocks (keep them in Reset state)
        for block in &mut blocks {
            let updated_metadata = block.metadata().with_priority(100);
            block.update_metadata(updated_metadata);
            assert_eq!(block.metadata().offload_priority(), Some(100));
        }

        // Step 3: Return blocks to inactive pool
        // Since blocks are in Reset state, insert() will put them in uninitialized_set
        // WITHOUT calling reset()
        pool.return_blocks(blocks);
        assert_eq!(pool.available_blocks(), 3);

        // Step 4: Acquire blocks again
        let reacquired_blocks = pool.acquire_free_blocks(3).unwrap();

        // Step 5: Verify priority is reset to default (0)
        for (i, block) in reacquired_blocks.iter().enumerate() {
            assert_eq!(
                block.metadata().offload_priority(),
                Some(0),
                "Block {} should have default priority after reallocation, but has {:?}",
                i,
                block.metadata().offload_priority()
            );
        }
    }

    /// Validates that after pool.reset(), all blocks have default priority
    /// regardless of what priority they had when registered.
    ///
    /// This test follows the exact flow described:
    /// 1. Create a tokens sequence
    /// 2. Allocate mutable blocks
    /// 3. Apply the tokens sequence and some non-default priority
    /// 4. Release them to the inactive pool (they go to priority_set as Registered)
    /// 5. Reset the inactive pool
    /// 6. Validate all blocks have default priority
    ///
    /// This test should PASS because blocks evicted from priority_set go through
    /// block.reset() which clears the priority.
    #[test]
    fn test_pool_reset_clears_priority_on_registered_blocks() {
        let async_runtime = tokio::runtime::Runtime::new().unwrap();
        const BLOCK_SIZE: u32 = 4;

        let mut pool = create_block_pool(3);
        assert_eq!(pool.available_blocks(), 3);

        // Step 1 & 2: Create tokens and allocate blocks
        let tokens = create_token_sequence(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]);
        let (mut blocks, _matched) = acquire_blocks(
            tokens,
            BLOCK_SIZE,
            &mut pool,
            async_runtime.handle().clone(),
        );
        assert_eq!(blocks.len(), 3);

        // Verify blocks are in Registered state
        for block in &blocks {
            assert!(
                block.state().is_registered(),
                "Block should be in Registered state after acquire_blocks"
            );
        }

        // Step 3: Set non-default priority on blocks
        for block in &mut blocks {
            let updated_metadata = block.metadata().with_priority(100);
            block.update_metadata(updated_metadata);
            assert_eq!(
                block.metadata().offload_priority(),
                Some(100),
                "Priority should be set to 100"
            );
        }

        // Step 4: Release blocks to inactive pool
        // Since blocks are Registered, they go to priority_set
        pool.return_blocks(blocks);
        assert_eq!(pool.available_blocks(), 3);

        // Verify blocks are in priority_set (not uninitialized_set)
        let (priority_count, uninit_count) = pool.status();
        assert_eq!(priority_count, 3, "All blocks should be in priority_set");
        assert_eq!(uninit_count, 0, "No blocks should be in uninitialized_set");

        // Step 5: Reset the pool
        // This calls acquire_free_blocks() which evicts from priority_set
        // and calls block.reset() on each, then returns them
        pool.reset().expect("Pool reset should succeed");

        // After reset, all blocks should be in uninitialized_set
        let (priority_count, uninit_count) = pool.status();
        assert_eq!(
            priority_count, 0,
            "priority_set should be empty after reset"
        );
        assert_eq!(
            uninit_count, 3,
            "All blocks should be in uninitialized_set after reset"
        );

        // Step 6: Acquire all blocks and verify priority is default (0)
        let reset_blocks = pool.acquire_free_blocks(3).unwrap();
        for (i, block) in reset_blocks.iter().enumerate() {
            assert_eq!(
                block.metadata().offload_priority(),
                Some(0),
                "Block {} should have default priority after pool reset, but has {:?}",
                i,
                block.metadata().offload_priority()
            );
        }
    }
}