heapdump_vmo/

allocations_table_v1.rs

// Copyright 2023 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use memory_mapped_vmo::{
    ReadOnlyMemoryMappedVmo, ReadWriteMemoryMappedArcVmo, ReadWriteMemoryMappedVmo,
};
use static_assertions::const_assert;
use std::collections::HashSet;
use std::mem::{align_of, size_of};
use std::ops::Deref;
use std::sync::Mutex;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering::{Relaxed, SeqCst};
use zerocopy::FromBytes;

pub use crate::resources_table_v1::ResourceKey;

type NodeIndex = u32;
type BucketHeads = [AtomicNodeIndex; NUM_BUCKETS];
const NUM_BUCKETS: usize = 1 << 16;
const NODE_INVALID: NodeIndex = NodeIndex::MAX;

/// Minimum memory alignment of an allocation table.
pub const MIN_ALIGNMENT: usize = align_of::<Node>();
const_assert!(MIN_ALIGNMENT % align_of::<BucketHeads>() == 0);

/// Maximum number of free nodes maintained locally in each shard. They can be
/// used without locking the centralized SharedAllocator.
const LOCAL_CACHE_CAPACITY: usize = 24;

/// Number of free nodes transferred at once to/from the SharedAllocator.
const REFILL_SPILL_BATCH: usize = 12;

// Define AtomicNodeIndex as a newtype so we can implement traits on it.
#[repr(transparent)]
#[derive(Debug, FromBytes)]
struct AtomicNodeIndex(AtomicU32);

impl AtomicNodeIndex {
    const fn new(value: u32) -> AtomicNodeIndex {
        AtomicNodeIndex(AtomicU32::new(value))
    }
}

impl Deref for AtomicNodeIndex {
    type Target = AtomicU32;

    fn deref(&self) -> &AtomicU32 {
        &self.0
    }
}

#[repr(C)]
#[derive(Debug, FromBytes)]
pub struct Node {
    next: AtomicNodeIndex,
    pub address: u64,
    pub size: u64,
    pub timestamp: zx::MonotonicInstant,
    pub thread_info_key: ResourceKey,
    pub stack_trace_key: ResourceKey,
}

/// Computes the capacity (number of nodes) of a memory region given its size (bytes).
fn compute_nodes_count(num_bytes: usize) -> Result<usize, crate::Error> {
    let Some(nodes_size) = num_bytes.checked_sub(size_of::<BucketHeads>()) else {
        return Err(crate::Error::BufferTooSmall);
    };

    let num_nodes = nodes_size / size_of::<Node>();
    if num_nodes > NODE_INVALID as usize {
        return Err(crate::Error::BufferTooBig);
    }

    Ok(num_nodes)
}

/// Precomputed shard and bucket indices for a given allocation address.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct AllocationKey<const NUM_SHARDS: usize> {
    shard_index: usize,
    bucket_index: usize,
    address: u64,
}

impl<const NUM_SHARDS: usize> AllocationKey<NUM_SHARDS> {
    /// Computes the key for `address` when the table is partitioned into `NUM_SHARDS` shards.
    pub const fn make(address: u64) -> AllocationKey<NUM_SHARDS> {
        let key = address ^ (address >> 33);
        let hash = key.wrapping_mul(0x9e3779b97f4a7c15);

        // Take the top N bits.
        const { assert!(NUM_BUCKETS.is_power_of_two()) };
        const { assert!(NUM_SHARDS.is_power_of_two()) };
        const { assert!(NUM_BUCKETS >= NUM_SHARDS) };
        let bucket_index = (hash >> (64 - NUM_BUCKETS.ilog2())) as usize;
        let shard_index = bucket_index / (NUM_BUCKETS / NUM_SHARDS);

        AllocationKey { shard_index, bucket_index, address }
    }

    pub const fn shard_index(&self) -> usize {
        self.shard_index
    }

    pub const fn address(&self) -> u64 {
        self.address
    }
}

/// Mediates write access to a VMO containing an hash table of allocated blocks indexed by block
/// address.
///
/// Operations can be "sharded" according to a hash function of the block address (see
/// `AllocationKey`). All the shards operate on the same VMO concurrently, but their implementation
/// is such that they never end up accessing the same data. Hence, callers do not need to worry
/// about race conditions.
///
/// From the point of view of an external observer, all the updates happen atomically so that, at
/// any time, snapshotting the VMO always results in a coherent snapshot of the table. The value of
/// `NUM_SHARDS` does not affect the "wire format" of data stored in the VMO.
///
/// Hash collisions are handled by maintaining per-bucket linked lists. Specifically, an array of
/// list heads, one for each bucket, is stored at the beginning of the VMO. The remaining part of
/// the VMO contains the linked list nodes.
pub struct AllocationsTableWriter<const NUM_SHARDS: usize> {
    storage: ReadWriteMemoryMappedArcVmo<Mutex<AllocationsTableWriterSharedState>>,
    shard_index: usize,
    local_free_list_head: NodeIndex,
    local_free_count: usize,
}

/// Writer state shared by all the shards, stored in the ReadWriteMemoryMappedArcVmo's metadata
/// and protected by a mutex.
///
/// Each shard maintains a local cache of up to LOCAL_CACHE_CAPACITY free nodes. When this cache
/// overflows/underflows, the centralized state is used to drain/refill it.
///
/// Information in this struct is not necessary for reading a snapshot. Therefore, we do not need
/// to store them in the main VMO or to offer any snapshot guarantee about them.
struct AllocationsTableWriterSharedState {
    // New free nodes are minted by a simple watermark allocator and the capacity of the hash table
    // (i.e. the maximum number of nodes) is fixed.
    watermark: NodeIndex,
    max_num_nodes: usize,

    // In order to make it possible to reuse nodes that have been freed, we also keep a linked list
    // of free nodes in addition to the watermark. When it is not empty, nodes are allocated by
    // popping the head of this list instead of incrementing the watermark.
    free_list_head: NodeIndex,
}

impl<const NUM_SHARDS: usize> AllocationsTableWriter<NUM_SHARDS> {
    /// Initializes a VMO as an empty table and creates `NUM_SHARDS` writers that can operate on
    /// disjoint shards of the table concurrently.
    ///
    /// Each address must be routed to the shard identified by `key.shard_index()`, where `key` is
    /// computed with `AllocationKey::make(address)`.
    ///
    /// # Safety
    /// The caller must guarantee that the `vmo` is not accessed by others while the returned
    /// instances are alive. However, it is always safe to take a snapshot and read that instead.
    pub unsafe fn new(
        vmo: &zx::Vmo,
        control_vmo_name: &zx::Name,
    ) -> Result<[AllocationsTableWriter<NUM_SHARDS>; NUM_SHARDS], crate::Error> {
        let mut raw_storage = unsafe { ReadWriteMemoryMappedVmo::new(vmo)? };
        let max_num_nodes = compute_nodes_count(raw_storage.vmo_size())?;

        // Clear the hash table.
        let bucket_heads = raw_storage.get_object_mut::<BucketHeads>(0).unwrap();
        for bucket_head in bucket_heads.iter() {
            bucket_head.store(NODE_INVALID, SeqCst);
        }

        let shared_allocator = Mutex::new(AllocationsTableWriterSharedState {
            max_num_nodes,
            watermark: 0,
            free_list_head: NODE_INVALID,
        });
        let storage = unsafe {
            ReadWriteMemoryMappedArcVmo::wrap_with_name(
                raw_storage,
                shared_allocator,
                control_vmo_name,
            )?
        };

        Ok(std::array::from_fn(|shard_index| AllocationsTableWriter {
            storage: storage.clone(),
            shard_index,
            local_free_list_head: NODE_INVALID,
            local_free_count: 0,
        }))
    }

    /// Returns a reference to the head of a given bucket's linked list.
    fn bucket_head_at(&self, bucket_index: usize) -> &AtomicNodeIndex {
        // The bucket heads are stored at the beginning of the VMO.
        let bucket_heads = self.storage.get_object::<BucketHeads>(0).unwrap();
        &bucket_heads[bucket_index]
    }

    /// Returns a reference to a given node.
    fn node_at(&self, node_index: NodeIndex) -> &Node {
        // The nodes are stored consecutively immediately after the bucket heads.
        let byte_offset = size_of::<BucketHeads>() + node_index as usize * size_of::<Node>();
        self.storage.get_object::<Node>(byte_offset).unwrap()
    }

    /// Returns a mutable reference to a given node.
    fn node_at_mut(&mut self, node_index: NodeIndex) -> &mut Node {
        // The nodes are stored consecutively immediately after the bucket heads.
        let byte_offset = size_of::<BucketHeads>() + node_index as usize * size_of::<Node>();
        self.storage.get_object_mut::<Node>(byte_offset).unwrap()
    }

    /// Inserts a new entry in the hash table.
    ///
    /// Returns Ok(true) if it has been inserted, Ok(false) if a previous entry with the same
    /// address already existed, or an error if no free nodes are available.
    pub fn insert_allocation(
        &mut self,
        key: AllocationKey<NUM_SHARDS>,
        size: u64,
        thread_info_key: ResourceKey,
        stack_trace_key: ResourceKey,
        timestamp: zx::MonotonicInstant,
    ) -> Result<bool, crate::Error> {
        assert_eq!(key.shard_index, self.shard_index);

        let old_head = self.bucket_head_at(key.bucket_index).load(Relaxed);

        // Verify that no entry with the same address already exists.
        let mut curr_index = old_head;
        while curr_index != NODE_INVALID {
            let curr_data = self.node_at(curr_index);
            if curr_data.address == key.address {
                return Ok(false);
            }
            curr_index = curr_data.next.load(Relaxed);
        }

        // Insert a new entry at the head of the list.
        let new_index = self.pop_free_node()?;
        *self.node_at_mut(new_index) = Node {
            address: key.address,
            size,
            timestamp,
            thread_info_key,
            stack_trace_key,
            next: AtomicNodeIndex::new(old_head),
        };
        self.bucket_head_at(key.bucket_index).store(new_index, SeqCst);
        Ok(true)
    }

    /// Removes an entry from the hash table and returns the value of the removed entry's size
    /// field.
    pub fn erase_allocation(&mut self, key: AllocationKey<NUM_SHARDS>) -> Option<u64> {
        assert_eq!(key.shard_index, self.shard_index);

        // Search the entry to be removed.
        let mut prev_index = None;
        let mut curr_index = self.bucket_head_at(key.bucket_index).load(Relaxed);
        while curr_index != NODE_INVALID {
            let curr_data = self.node_at(curr_index);
            let curr_data_size = curr_data.size;
            let curr_data_next = curr_data.next.load(Relaxed);

            // Is this the entry we were looking for?
            if curr_data.address == key.address {
                if let Some(prev_index) = prev_index {
                    self.node_at(prev_index).next.store(curr_data_next, SeqCst);
                } else {
                    self.bucket_head_at(key.bucket_index).store(curr_data_next, SeqCst);
                }
                self.push_free_node(curr_index);
                return Some(curr_data_size);
            }

            prev_index = Some(curr_index);
            curr_index = curr_data.next.load(Relaxed);
        }

        // Not found.
        None
    }

    /// Atomically updates metadata for the given address in the hash table.
    ///
    /// Returns Ok(Some(old_size)) if it has been updated, Ok(None) if no entry with the given
    /// address exists, or an error if no free nodes are available.
    pub fn replace_allocation(
        &mut self,
        key: AllocationKey<NUM_SHARDS>,
        size: u64,
        thread_info_key: ResourceKey,
        stack_trace_key: ResourceKey,
        timestamp: zx::MonotonicInstant,
    ) -> Result<Option<u64>, crate::Error> {
        assert_eq!(key.shard_index, self.shard_index);

        // Updates are implemented as an atomic insertion of a new node at the head followed by an
        // atomic deletion of the old node. Therefore, while an update is in progress, two nodes
        // with the same address may exist. `AllocationsTableReader`` is aware of this, and it
        // only returns the first one (i.e. the newer one) in such a case.
        // Note: the body of this function has been split in two parts to test reads of the
        // intermediate state.

        // Check preconditions and insert the new node.
        let Some((new_node, old_node, old_size)) =
            self.replace_allocation_begin(key, size, thread_info_key, stack_trace_key, timestamp)?
        else {
            return Ok(None);
        };

        // Remove the old node.
        self.replace_allocation_end(new_node, old_node);
        Ok(Some(old_size))
    }

    /// First part of `replace_allocation`. Locates the old node and inserts the new one at the head
    fn replace_allocation_begin(
        &mut self,
        key: AllocationKey<NUM_SHARDS>,
        size: u64,
        thread_info_key: ResourceKey,
        stack_trace_key: ResourceKey,
        timestamp: zx::MonotonicInstant,
    ) -> Result<Option<(NodeIndex, NodeIndex, u64)>, crate::Error> {
        let old_head = self.bucket_head_at(key.bucket_index).load(Relaxed);

        // Search the existing entry with the requested address.
        let (old_index, old_size) = 'search: {
            let mut curr_index = old_head;
            while curr_index != NODE_INVALID {
                let curr_data = self.node_at(curr_index);
                if curr_data.address == key.address {
                    break 'search (curr_index, curr_data.size); // Found.
                }
                curr_index = curr_data.next.load(Relaxed);
            }

            return Ok(None); // Not found.
        };

        // Insert a new entry at the head of the list.
        let new_index = self.pop_free_node()?;
        *self.node_at_mut(new_index) = Node {
            address: key.address,
            size,
            timestamp,
            thread_info_key,
            stack_trace_key,
            next: AtomicNodeIndex::new(old_head),
        };
        self.bucket_head_at(key.bucket_index).store(new_index, SeqCst);
        Ok(Some((new_index, old_index, old_size)))
    }

    // Second part of `replace_allocation`. Removes the old node.
    fn replace_allocation_end(&mut self, new_node: NodeIndex, old_node: NodeIndex) {
        let tail = self.node_at(old_node).next.load(Relaxed);

        // Note: `replace_allocation_begin` always inserts the new node at the head; therefore, the
        // old node will always be one of its successors.
        let mut scan_node = new_node;
        loop {
            assert_ne!(scan_node, NODE_INVALID, "the old node must be a successor of new node");
            let scan_data = self.node_at(scan_node);
            if scan_data.next.load(Relaxed) == old_node {
                scan_data.next.store(tail, SeqCst);
                self.push_free_node(old_node);
                return;
            }
            scan_node = scan_data.next.load(Relaxed);
        }
    }

    /// Inserts a node into the free list, spilling half the local cache to the shared free list if
    /// the local cache exceeds `LOCAL_CACHE_CAPACITY`.
    fn push_free_node(&mut self, index: NodeIndex) {
        self.node_at(index).next.store(self.local_free_list_head, Relaxed);
        self.local_free_list_head = index;
        self.local_free_count += 1;

        if self.local_free_count > LOCAL_CACHE_CAPACITY {
            let mut shared = self.storage.metadata().lock().unwrap();

            // Transfer REFILL_SPILL_BATCH nodes from the local free list to the shared free list.
            for _ in 0..REFILL_SPILL_BATCH {
                let node_index = self.local_free_list_head;
                let node_data = self.node_at(node_index);
                let next_local = node_data.next.load(Relaxed);

                node_data.next.store(shared.free_list_head, Relaxed);
                shared.free_list_head = node_index;

                self.local_free_list_head = next_local;
                self.local_free_count -= 1;
            }
        }
    }

    /// Takes a node from the local free list, refilling `REFILL_SPILL_BATCH` nodes from the shared
    /// free list or watermark allocator when the local cache is empty.
    fn pop_free_node(&mut self) -> Result<NodeIndex, crate::Error> {
        if self.local_free_list_head == NODE_INVALID {
            let mut shared = self.storage.metadata().lock().unwrap();

            // Obtain up to REFILL_SPILL_BATCH nodes from the shared allocator.
            while self.local_free_count < REFILL_SPILL_BATCH {
                let new_node = if shared.free_list_head != NODE_INVALID {
                    // Pop a node from the free list.
                    let index = shared.free_list_head;
                    shared.free_list_head = self.node_at(index).next.load(Relaxed);
                    index
                } else if (shared.watermark as usize) < shared.max_num_nodes {
                    // Allocate one node with the watermark allocator.
                    let index = shared.watermark;
                    shared.watermark += 1;
                    index
                } else {
                    break;
                };

                self.node_at(new_node).next.store(self.local_free_list_head, Relaxed);
                self.local_free_list_head = new_node;
                self.local_free_count += 1;
            }
        }

        if self.local_free_list_head != NODE_INVALID {
            let result = self.local_free_list_head;
            self.local_free_list_head = self.node_at(result).next.load(Relaxed);
            self.local_free_count -= 1;
            Ok(result)
        } else {
            // We are out of space.
            Err(crate::Error::OutOfSpace)
        }
    }
}

/// Mediates read access to a VMO written by AllocationsTableWriter.
pub struct AllocationsTableReader {
    storage: ReadOnlyMemoryMappedVmo,
    max_num_nodes: usize,
}

impl AllocationsTableReader {
    /// # Safety
    /// The caller must guarantee that the `vmo` is not accessed by others while the returned
    /// instance is alive, usually by taking a snapshot of the VMO that AllocationsTableWriter
    /// operates on and then reading the snapshot instead.
    pub unsafe fn new(vmo: &zx::Vmo) -> Result<AllocationsTableReader, crate::Error> {
        let storage = unsafe { ReadOnlyMemoryMappedVmo::new(vmo)? };
        let max_num_nodes = compute_nodes_count(storage.vmo_size())?;

        Ok(AllocationsTableReader { storage, max_num_nodes })
    }

    /// Iterates all the allocated blocks that are present in the hash table.
    pub fn iter(&self) -> impl Iterator<Item = Result<&Node, crate::Error>> {
        // The bucket heads are stored at the beginning of the VMO.
        let bucket_heads = self.storage.get_object::<BucketHeads>(0).unwrap();
        bucket_heads.iter().map(|head| self.iterate_bucket(head.load(Relaxed))).flatten()
    }

    fn iterate_bucket(&self, head: NodeIndex) -> impl Iterator<Item = Result<&Node, crate::Error>> {
        let mut curr_index = head;
        let mut seen_addresses = HashSet::new();
        std::iter::from_fn(move || {
            while curr_index != NODE_INVALID {
                if (curr_index as usize) < self.max_num_nodes {
                    // The nodes are stored consecutively immediately after the bucket heads.
                    let byte_offset =
                        size_of::<BucketHeads>() + curr_index as usize * size_of::<Node>();
                    let curr_data = self.storage.get_object::<Node>(byte_offset).unwrap();
                    curr_index = curr_data.next.load(Relaxed);

                    // Only return the first occurrence of each address. This ensures that only the
                    // newer record is returned if an update is in progress.
                    if seen_addresses.insert(curr_data.address) {
                        return Some(Ok(curr_data));
                    }
                } else {
                    return Some(Err(crate::Error::InvalidInput));
                };
            }

            None
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use assert_matches::assert_matches;
    use std::alloc::Layout;
    use std::collections::HashMap;

    // Some tests below use this constant to ensure that each bucket is been used at least 10 times.
    const NUM_ITERATIONS: usize = NUM_BUCKETS * 10;

    // Ensure we allocate enough nodes to store all the blocks the tests need plus some buffer.
    const NUM_NODES: usize = NUM_ITERATIONS + 100;

    // Placeholder values used in the tests below:
    const THREAD_INFO_RESOURCE_KEY_1: ResourceKey = ResourceKey::from_raw(0x1122);
    const THREAD_INFO_RESOURCE_KEY_2: ResourceKey = ResourceKey::from_raw(0x3344);
    const STACK_TRACE_RESOURCE_KEY_1: ResourceKey = ResourceKey::from_raw(0x1212);
    const STACK_TRACE_RESOURCE_KEY_2: ResourceKey = ResourceKey::from_raw(0x3434);

    struct TestStorage {
        vmo: zx::Vmo,
    }

    impl TestStorage {
        pub fn new(num_nodes: usize) -> (TestStorage, AllocationsTableWriter<1>) {
            let (storage, [writer]) = Self::new_sharded::<1>(num_nodes);
            (storage, writer)
        }

        pub fn new_sharded<const NUM_SHARDS: usize>(
            num_nodes: usize,
        ) -> (TestStorage, [AllocationsTableWriter<NUM_SHARDS>; NUM_SHARDS]) {
            let nodes_layout = Layout::array::<Node>(num_nodes).unwrap();
            let (layout, nodes_offset) = Layout::new::<BucketHeads>().extend(nodes_layout).unwrap();
            assert_eq!(nodes_offset, size_of::<BucketHeads>());

            let vmo = zx::Vmo::create(layout.size() as u64).unwrap();

            // SAFETY: Each shard writer in the returned array operates on a disjoint shard.
            let writers =
                unsafe { AllocationsTableWriter::new(&vmo, &zx::Name::default()) }.unwrap();

            (TestStorage { vmo }, writers)
        }

        fn create_reader(&self) -> AllocationsTableReader {
            let snapshot = self
                .vmo
                .create_child(
                    zx::VmoChildOptions::SNAPSHOT | zx::VmoChildOptions::NO_WRITE,
                    0,
                    self.vmo.get_size().unwrap(),
                )
                .unwrap();

            // SAFETY: Readers operate on the read-only VMO snapshot we just created.
            unsafe { AllocationsTableReader::new(&snapshot) }.unwrap()
        }
    }

    #[test]
    fn test_cannot_insert_twice() {
        let (_storage, mut writer) = TestStorage::new(NUM_NODES);

        let result = writer.insert_allocation(
            AllocationKey::make(0x1234),
            0x5678,
            THREAD_INFO_RESOURCE_KEY_1,
            STACK_TRACE_RESOURCE_KEY_1,
            zx::MonotonicInstant::ZERO,
        );
        assert_eq!(result, Ok(true));

        let result = writer.insert_allocation(
            AllocationKey::make(0x1234),
            0x5678,
            THREAD_INFO_RESOURCE_KEY_1,
            STACK_TRACE_RESOURCE_KEY_1,
            zx::MonotonicInstant::ZERO,
        );
        assert_eq!(result, Ok(false));
    }

    #[test]
    fn test_cannot_erase_twice() {
        let (_storage, mut writer) = TestStorage::new(NUM_NODES);

        let result = writer.insert_allocation(
            AllocationKey::make(0x1234),
            0x5678,
            THREAD_INFO_RESOURCE_KEY_1,
            STACK_TRACE_RESOURCE_KEY_1,
            zx::MonotonicInstant::ZERO,
        );
        assert_eq!(result, Ok(true));

        let result = writer.erase_allocation(AllocationKey::make(0x1234));
        assert_eq!(result, Some(0x5678));

        let result = writer.erase_allocation(AllocationKey::make(0x1234));
        assert_eq!(result, None);
    }

    #[test]
    fn test_out_of_space() {
        let (_storage, mut writer) = TestStorage::new(NUM_NODES);

        // Test that inserting up to `NUM_NODES` works.
        for i in 0..NUM_NODES {
            let result = writer.insert_allocation(
                AllocationKey::make(i as u64),
                1,
                THREAD_INFO_RESOURCE_KEY_1,
                STACK_TRACE_RESOURCE_KEY_1,
                zx::MonotonicInstant::ZERO,
            );
            assert_eq!(result, Ok(true));
        }

        // Test that inserting an extra node fails.
        let result = writer.insert_allocation(
            AllocationKey::make(NUM_NODES as u64),
            1,
            THREAD_INFO_RESOURCE_KEY_1,
            STACK_TRACE_RESOURCE_KEY_1,
            zx::MonotonicInstant::ZERO,
        );
        assert_eq!(result, Err(crate::Error::OutOfSpace));

        // Test that removing an element and then inserting again succeeds.
        let result = writer.erase_allocation(AllocationKey::make(0));
        assert_eq!(result, Some(1));
        let result = writer.insert_allocation(
            AllocationKey::make(NUM_NODES as u64),
            1,
            THREAD_INFO_RESOURCE_KEY_1,
            STACK_TRACE_RESOURCE_KEY_1,
            zx::MonotonicInstant::ZERO,
        );
        assert_eq!(result, Ok(true));
    }

    #[test]
    fn test_loop_insert_then_erase() {
        let (_storage, mut writer) = TestStorage::new(NUM_NODES);

        for i in 0..NUM_ITERATIONS {
            let result = writer.insert_allocation(
                AllocationKey::make(i as u64),
                1,
                THREAD_INFO_RESOURCE_KEY_1,
                STACK_TRACE_RESOURCE_KEY_1,
                zx::MonotonicInstant::ZERO,
            );
            assert_eq!(result, Ok(true), "failed to insert 0x{:x}", i);

            let result = writer.erase_allocation(AllocationKey::make(i as u64));
            assert_eq!(result, Some(1), "failed to erase 0x{:x}", i);
        }
    }

    #[test]
    fn test_bulk_insert_then_erase_same_order() {
        let (_storage, mut writer) = TestStorage::new(NUM_NODES);

        for i in 0..NUM_ITERATIONS {
            let result = writer.insert_allocation(
                AllocationKey::make(i as u64),
                1,
                THREAD_INFO_RESOURCE_KEY_1,
                STACK_TRACE_RESOURCE_KEY_1,
                zx::MonotonicInstant::ZERO,
            );
            assert_eq!(result, Ok(true), "failed to insert 0x{:x}", i);
        }
        for i in 0..NUM_ITERATIONS {
            let result = writer.erase_allocation(AllocationKey::make(i as u64));
            assert_eq!(result, Some(1), "failed to erase 0x{:x}", i);
        }
    }

    #[test]
    fn test_bulk_insert_then_erase_reverse_order() {
        let (_storage, mut writer) = TestStorage::new(NUM_NODES);

        for i in 0..NUM_ITERATIONS {
            let result = writer.insert_allocation(
                AllocationKey::make(i as u64),
                1,
                THREAD_INFO_RESOURCE_KEY_1,
                STACK_TRACE_RESOURCE_KEY_1,
                zx::MonotonicInstant::ZERO,
            );
            assert_eq!(result, Ok(true), "failed to insert 0x{:x}", i);
        }
        for i in (0..NUM_ITERATIONS).rev() {
            let result = writer.erase_allocation(AllocationKey::make(i as u64));
            assert_eq!(result, Some(1), "failed to erase 0x{:x}", i);
        }
    }

    #[test]
    fn test_read_empty() {
        let (storage, _writer) = TestStorage::new(NUM_NODES);

        // Verify that it is empty.
        let reader = storage.create_reader();
        assert_eq!(reader.iter().count(), 0);
    }

    #[test]
    fn test_read_populated() {
        let (storage, mut writer) = TestStorage::new(NUM_NODES);

        // Fill the hash table.
        let mut expected_map = HashMap::new();
        for i in 0..NUM_ITERATIONS as u64 {
            let thread_info_key =
                if i % 4 >= 2 { THREAD_INFO_RESOURCE_KEY_1 } else { THREAD_INFO_RESOURCE_KEY_2 };
            let stack_trace_key =
                if i % 2 == 0 { STACK_TRACE_RESOURCE_KEY_1 } else { STACK_TRACE_RESOURCE_KEY_2 };
            let timestamp =
                zx::MonotonicInstant::from_nanos((NUM_ITERATIONS as i64 / 2) - (i as i64)); // test negative values too
            let result = writer.insert_allocation(
                AllocationKey::make(i),
                1,
                thread_info_key,
                stack_trace_key,
                timestamp,
            );
            assert_eq!(result, Ok(true), "failed to insert 0x{:x}", i);

            expected_map.insert(i, (1, thread_info_key, stack_trace_key, timestamp));
        }

        // Read it back and verify.
        let reader = storage.create_reader();
        let mut actual_map = HashMap::new();
        for node in reader.iter() {
            let Node { address, size, thread_info_key, stack_trace_key, timestamp, .. } =
                node.unwrap();
            assert!(
                actual_map
                    .insert(*address, (*size, *thread_info_key, *stack_trace_key, *timestamp))
                    .is_none(),
                "address 0x{:x} was read more than once",
                address
            );
        }

        assert_eq!(actual_map, expected_map);
    }

    #[test]
    fn test_read_bad_bucket_head() {
        let (storage, writer) = TestStorage::new(NUM_NODES);

        // Initialize the hash table and corrupt one of the heads.
        writer.bucket_head_at(NUM_BUCKETS / 2).store(NODE_INVALID - 1, SeqCst);

        // Try to read it back and verify that the iterator returns an error.
        let reader = storage.create_reader();
        let contains_error = reader.iter().any(|e| e.is_err());

        assert!(contains_error);
    }

    #[test]
    fn test_read_bad_node_next() {
        let (storage, mut writer) = TestStorage::new(NUM_NODES);

        // Initialize the hash table and insert a node with a bad next pointer.
        writer.bucket_head_at(NUM_BUCKETS / 2).store(0, SeqCst);
        *writer.node_at_mut(0) = Node {
            next: AtomicNodeIndex::new(NODE_INVALID - 1),
            address: 0x1234,
            size: 0x5678,
            thread_info_key: THREAD_INFO_RESOURCE_KEY_1,
            stack_trace_key: STACK_TRACE_RESOURCE_KEY_1,
            timestamp: zx::MonotonicInstant::from_nanos(99999999),
        };

        // Try to read it back and verify that the iterator returns an error.
        let reader = storage.create_reader();
        let contains_error = reader.iter().any(|e| e.is_err());

        assert!(contains_error);
    }

    #[test]
    fn test_replace() {
        let (storage, mut writer) = TestStorage::new(NUM_NODES);

        let result = writer.insert_allocation(
            AllocationKey::make(0x1234),
            0x1111,
            THREAD_INFO_RESOURCE_KEY_1,
            STACK_TRACE_RESOURCE_KEY_1,
            zx::MonotonicInstant::ZERO,
        );
        assert_eq!(result, Ok(true));

        // Verify initial contents.
        let reader = storage.create_reader();
        let mut iter = reader.iter();
        assert_matches!(iter.next(), Some(Ok(Node { address: 0x1234, size: 0x1111, .. })));
        assert_matches!(iter.next(), None);

        let result = writer.replace_allocation_begin(
            AllocationKey::make(0x1234),
            0x2222,
            THREAD_INFO_RESOURCE_KEY_2,
            STACK_TRACE_RESOURCE_KEY_2,
            zx::MonotonicInstant::ZERO,
        );
        let Ok(Some((new_node, old_node, old_size))) = result else {
            panic!("Update begin is supposed to succeed in this test, got {:?} instead", result)
        };
        assert_eq!(old_size, 0x1111);

        // Verify that intermediate contents already reflect the update.
        let reader = storage.create_reader();
        let mut iter = reader.iter();
        assert_matches!(iter.next(), Some(Ok(Node { address: 0x1234, size: 0x2222, .. })));
        assert_matches!(iter.next(), None);

        writer.replace_allocation_end(new_node, old_node);

        // Verify that final contents reflect the update.
        let reader = storage.create_reader();
        let mut iter = reader.iter();
        assert_matches!(iter.next(), Some(Ok(Node { address: 0x1234, size: 0x2222, .. })));
        assert_matches!(iter.next(), None);
    }

    #[test]
    fn test_cannot_replace_nonexisting() {
        let (_storage, mut writer) = TestStorage::new(NUM_NODES);

        let result = writer.replace_allocation(
            AllocationKey::make(0x1234),
            0x5678,
            THREAD_INFO_RESOURCE_KEY_1,
            STACK_TRACE_RESOURCE_KEY_1,
            zx::MonotonicInstant::ZERO,
        );
        assert_eq!(result, Ok(None));
    }

    #[test]
    fn test_sharded_spill_and_refill() {
        let (storage, [mut writer0, mut writer1]) = TestStorage::new_sharded::<2>(NUM_NODES);

        // Helper to generate keys belonging to shard 0 or shard 1.
        let key_for_shard = |shard: usize, i: usize| -> AllocationKey<2> {
            AllocationKey {
                shard_index: shard,
                bucket_index: (shard * NUM_BUCKETS + i % NUM_BUCKETS) / 2,
                address: ((i * 2 + shard) as u64) << 4,
            }
        };

        // Exhaust all `NUM_NODES` nodes on shard 0.
        for i in 0..NUM_NODES {
            assert_eq!(
                writer0.insert_allocation(
                    key_for_shard(0, i),
                    1,
                    THREAD_INFO_RESOURCE_KEY_1,
                    STACK_TRACE_RESOURCE_KEY_1,
                    zx::MonotonicInstant::ZERO,
                ),
                Ok(true)
            );
        }

        // Shard 1 cannot allocate yet because watermark is at NUM_NODES and nothing was spilled.
        assert_eq!(
            writer1.insert_allocation(
                key_for_shard(1, 0),
                1,
                THREAD_INFO_RESOURCE_KEY_1,
                STACK_TRACE_RESOURCE_KEY_1,
                zx::MonotonicInstant::ZERO,
            ),
            Err(crate::Error::OutOfSpace)
        );

        // Free `LOCAL_CACHE_CAPACITY + 1` nodes on shard 0 to trigger a spill of
        // `REFILL_SPILL_BATCH` nodes to the shared free list.
        for i in 0..(LOCAL_CACHE_CAPACITY + 1) {
            assert_eq!(writer0.erase_allocation(key_for_shard(0, i)), Some(1));
        }

        // Shard 1 can now allocate `REFILL_SPILL_BATCH` nodes from the shared free list.
        for i in 0..REFILL_SPILL_BATCH {
            assert_eq!(
                writer1.insert_allocation(
                    key_for_shard(1, i),
                    2,
                    THREAD_INFO_RESOURCE_KEY_2,
                    STACK_TRACE_RESOURCE_KEY_2,
                    zx::MonotonicInstant::ZERO,
                ),
                Ok(true)
            );
        }

        let reader = storage.create_reader();
        assert_eq!(
            reader.iter().count(),
            NUM_NODES - (LOCAL_CACHE_CAPACITY + 1) + REFILL_SPILL_BATCH
        );
    }
}