Make chunk map as global side metadata, and each space only lists its own chunks using chunk map (#1304)

The chunk map records the lowest chunk and the highest chunk used by a
space, and accesses all the chunk states (local side metadata) in the
range. There are multiple issues when we use discontiguous spaces.
1. Chunks in the range may be ~~used~~ unused, thus the local side
metadata may be unmapped.
2. Chunks in the range may belong to other spaces. If other spaces also
use the chunk map, they will list all the chunks even when the chunks do
not belong to the space.

This PR makes the chunk state side metadata as a global side metadata.
We need the chunk map for plans that use the chunk map. As we will use
mark sweep or nonmoving immix as the non moving policy, they both use
the chunk map, and the non moving policy is included in every plan, we
just include the side metadata for every plan.

This PR fixes #1197.

Author: qinsoon

src/policy/immix/immixspace.rs

@@ -254,7 +254,6 @@ impl<VM: VMBinding> ImmixSpace<VM> {
             vec![
                 MetadataSpec::OnSide(Block::DEFRAG_STATE_TABLE),
                 MetadataSpec::OnSide(Block::MARK_TABLE),
-                MetadataSpec::OnSide(ChunkMap::ALLOC_TABLE),
                 *VM::VMObjectModel::LOCAL_MARK_BIT_SPEC,
                 *VM::VMObjectModel::LOCAL_FORWARDING_BITS_SPEC,
                 *VM::VMObjectModel::LOCAL_FORWARDING_POINTER_SPEC,
@@ -266,7 +265,6 @@ impl<VM: VMBinding> ImmixSpace<VM> {
                 MetadataSpec::OnSide(Line::MARK_TABLE),
                 MetadataSpec::OnSide(Block::DEFRAG_STATE_TABLE),
                 MetadataSpec::OnSide(Block::MARK_TABLE),
-                MetadataSpec::OnSide(ChunkMap::ALLOC_TABLE),
                 *VM::VMObjectModel::LOCAL_MARK_BIT_SPEC,
                 *VM::VMObjectModel::LOCAL_FORWARDING_BITS_SPEC,
                 *VM::VMObjectModel::LOCAL_FORWARDING_POINTER_SPEC,
@@ -315,6 +313,7 @@ impl<VM: VMBinding> ImmixSpace<VM> {
         let scheduler = args.scheduler.clone();
         let common =
             CommonSpace::new(args.into_policy_args(true, false, Self::side_metadata_specs()));
+        let space_index = common.descriptor.get_index();
         ImmixSpace {
             pr: if common.vmrequest.is_discontiguous() {
                 BlockPageResource::new_discontiguous(
@@ -332,7 +331,7 @@ impl<VM: VMBinding> ImmixSpace<VM> {
                 )
             },
             common,
-            chunk_map: ChunkMap::new(),
+            chunk_map: ChunkMap::new(space_index),
             line_mark_state: AtomicU8::new(Line::RESET_MARK_STATE),
             line_unavail_state: AtomicU8::new(Line::RESET_MARK_STATE),
             lines_consumed: AtomicUsize::new(0),
@@ -541,7 +540,7 @@ impl<VM: VMBinding> ImmixSpace<VM> {
         self.defrag.notify_new_clean_block(copy);
         let block = Block::from_aligned_address(block_address);
         block.init(copy);
-        self.chunk_map.set(block.chunk(), ChunkState::Allocated);
+        self.chunk_map.set_allocated(block.chunk(), true);
         self.lines_consumed
             .fetch_add(Block::LINES, Ordering::SeqCst);
         Some(block)
@@ -928,7 +927,7 @@ struct SweepChunk<VM: VMBinding> {
 
 impl<VM: VMBinding> GCWork<VM> for SweepChunk<VM> {
     fn do_work(&mut self, _worker: &mut GCWorker<VM>, mmtk: &'static MMTK<VM>) {
-        assert_eq!(self.space.chunk_map.get(self.chunk), ChunkState::Allocated);
+        assert!(self.space.chunk_map.get(self.chunk).unwrap().is_allocated());
 
         let mut histogram = self.space.defrag.new_histogram();
         let line_mark_state = if super::BLOCK_ONLY {
@@ -979,7 +978,7 @@ impl<VM: VMBinding> GCWork<VM> for SweepChunk<VM> {
         probe!(mmtk, sweep_chunk, allocated_blocks);
         // Set this chunk as free if there is not live blocks.
         if allocated_blocks == 0 {
-            self.space.chunk_map.set(self.chunk, ChunkState::Free)
+            self.space.chunk_map.set_allocated(self.chunk, false)
         }
         self.space.defrag.add_completed_mark_histogram(histogram);
         self.epilogue.finish_one_work_packet();

src/policy/marksweepspace/native_ms/global.rs

@@ -288,7 +288,7 @@ impl<VM: VMBinding> MarkSweepSpace<VM> {
         let vm_map = args.vm_map;
         let is_discontiguous = args.vmrequest.is_discontiguous();
         let local_specs = {
-            metadata::extract_side_metadata(&vec![
+            metadata::extract_side_metadata(&[
                 MetadataSpec::OnSide(Block::NEXT_BLOCK_TABLE),
                 MetadataSpec::OnSide(Block::PREV_BLOCK_TABLE),
                 MetadataSpec::OnSide(Block::FREE_LIST_TABLE),
@@ -300,11 +300,11 @@ impl<VM: VMBinding> MarkSweepSpace<VM> {
                 MetadataSpec::OnSide(Block::BLOCK_LIST_TABLE),
                 MetadataSpec::OnSide(Block::TLS_TABLE),
                 MetadataSpec::OnSide(Block::MARK_TABLE),
-                MetadataSpec::OnSide(ChunkMap::ALLOC_TABLE),
                 *VM::VMObjectModel::LOCAL_MARK_BIT_SPEC,
             ])
         };
         let common = CommonSpace::new(args.into_policy_args(false, false, local_specs));
+        let space_index = common.descriptor.get_index();
         MarkSweepSpace {
             pr: if is_discontiguous {
                 BlockPageResource::new_discontiguous(
@@ -322,7 +322,7 @@ impl<VM: VMBinding> MarkSweepSpace<VM> {
                 )
             },
             common,
-            chunk_map: ChunkMap::new(),
+            chunk_map: ChunkMap::new(space_index),
             scheduler,
             abandoned: Mutex::new(AbandonedBlockLists::new()),
             abandoned_in_gc: Mutex::new(AbandonedBlockLists::new()),
@@ -402,7 +402,7 @@ impl<VM: VMBinding> MarkSweepSpace<VM> {
 
     pub fn record_new_block(&self, block: Block) {
         block.init();
-        self.chunk_map.set(block.chunk(), ChunkState::Allocated);
+        self.chunk_map.set_allocated(block.chunk(), true);
     }
 
     pub fn prepare(&mut self, full_heap: bool) {
@@ -567,7 +567,7 @@ struct PrepareChunkMap<VM: VMBinding> {
 
 impl<VM: VMBinding> GCWork<VM> for PrepareChunkMap<VM> {
     fn do_work(&mut self, _worker: &mut GCWorker<VM>, _mmtk: &'static MMTK<VM>) {
-        debug_assert!(self.space.chunk_map.get(self.chunk) == ChunkState::Allocated);
+        debug_assert!(self.space.chunk_map.get(self.chunk).unwrap().is_allocated());
         // number of allocated blocks.
         let mut n_occupied_blocks = 0;
         self.chunk
@@ -581,7 +581,7 @@ impl<VM: VMBinding> GCWork<VM> for PrepareChunkMap<VM> {
             });
         if n_occupied_blocks == 0 {
             // Set this chunk as free if there is no live blocks.
-            self.space.chunk_map.set(self.chunk, ChunkState::Free)
+            self.space.chunk_map.set_allocated(self.chunk, false)
         } else {
             // Otherwise this chunk is occupied, and we reset the mark bit if it is on the side.
             if let MetadataSpec::OnSide(side) = *VM::VMObjectModel::LOCAL_MARK_BIT_SPEC {
@@ -617,7 +617,7 @@ struct SweepChunk<VM: VMBinding> {
 
 impl<VM: VMBinding> GCWork<VM> for SweepChunk<VM> {
     fn do_work(&mut self, _worker: &mut GCWorker<VM>, _mmtk: &'static MMTK<VM>) {
-        assert_eq!(self.space.chunk_map.get(self.chunk), ChunkState::Allocated);
+        assert!(self.space.chunk_map.get(self.chunk).unwrap().is_allocated());
 
         // number of allocated blocks.
         let mut allocated_blocks = 0;
@@ -636,7 +636,7 @@ impl<VM: VMBinding> GCWork<VM> for SweepChunk<VM> {
         probe!(mmtk, sweep_chunk, allocated_blocks);
         // Set this chunk as free if there is not live blocks.
         if allocated_blocks == 0 {
-            self.space.chunk_map.set(self.chunk, ChunkState::Free)
+            self.space.chunk_map.set_allocated(self.chunk, false);
         }
         self.epilogue.finish_one_work_packet();
     }

src/util/heap/chunk_map.rs

@@ -44,20 +44,68 @@ impl Chunk {
     }
 }
 
-/// Chunk allocation state
-#[repr(u8)]
-#[derive(Debug, PartialEq, Clone, Copy)]
-pub enum ChunkState {
-    /// The chunk is not allocated.
-    Free = 0,
-    /// The chunk is allocated.
-    Allocated = 1,
+/// The allocation state for a chunk in the chunk map. It includes whether each chunk is allocated or free, and the space the chunk belongs to.
+/// Highest bit: 0 = free, 1 = allocated
+/// Lower 4 bits: Space index (0-15) if the chunk is allocated.
+#[repr(transparent)]
+#[derive(PartialEq, Clone, Copy)]
+pub struct ChunkState(u8);
+
+impl ChunkState {
+    const ALLOC_BIT_MASK: u8 = 0x80;
+    const SPACE_INDEX_MASK: u8 = 0x0F;
+
+    /// Create a new ChunkState that represents being allocated in the given space
+    pub fn allocated(space_index: usize) -> ChunkState {
+        debug_assert!(space_index < crate::util::heap::layout::heap_parameters::MAX_SPACES);
+        let mut encode = space_index as u8;
+        encode |= Self::ALLOC_BIT_MASK;
+        ChunkState(encode)
+    }
+    /// Create a new ChunkState that represents being free
+    pub fn free() -> ChunkState {
+        ChunkState(0u8)
+    }
+    /// Is the chunk free?
+    pub fn is_free(&self) -> bool {
+        self.0 == 0
+    }
+    /// Is the chunk allocated?
+    pub fn is_allocated(&self) -> bool {
+        !self.is_free()
+    }
+    /// Get the space index of the chunk
+    pub fn get_space_index(&self) -> usize {
+        debug_assert!(self.is_allocated());
+        let index = (self.0 & Self::SPACE_INDEX_MASK) as usize;
+        debug_assert!(index < crate::util::heap::layout::heap_parameters::MAX_SPACES);
+        index
+    }
+}
+
+impl std::fmt::Debug for ChunkState {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if self.is_free() {
+            write!(f, "Free")
+        } else {
+            write!(f, "Allocated({})", self.get_space_index())
+        }
+    }
 }
 
 /// A byte-map to record all the allocated chunks.
 /// A plan can use this to maintain records for the chunks that they used, and the states of the chunks.
-/// Any plan that uses the chunk map should include the `ALLOC_TABLE` spec in their local sidemetadata specs
+/// Any plan that uses the chunk map should include the `ALLOC_TABLE` spec in their global sidemetadata specs.
+///
+/// A chunk map is created for a space (identified by the space index), and will only update or list chunks for that space.
 pub struct ChunkMap {
+    /// The space that uses this chunk map.
+    space_index: usize,
+    /// The range of chunks that are used by the space. The range only records the lowest chunk and the highest chunk.
+    /// All the chunks that are used for the space are within the range, but not necessarily that all the chunks in the range
+    /// are used for the space. Spaces may be discontiguous, thus the range may include chunks that do not belong to the space.
+    /// We need to use the space index in the chunk map and the space index encoded with the chunk state to know if
+    /// the chunk belongs to the current space.
     chunk_range: Mutex<Range<Chunk>>,
 }
 
@@ -66,22 +114,40 @@ impl ChunkMap {
     pub const ALLOC_TABLE: SideMetadataSpec =
         crate::util::metadata::side_metadata::spec_defs::CHUNK_MARK;
 
-    pub fn new() -> Self {
+    pub fn new(space_index: usize) -> Self {
         Self {
+            space_index,
             chunk_range: Mutex::new(Chunk::ZERO..Chunk::ZERO),
         }
     }
 
-    /// Set chunk state
-    pub fn set(&self, chunk: Chunk, state: ChunkState) {
+    /// Set a chunk as allocated, or as free.
+    pub fn set_allocated(&self, chunk: Chunk, allocated: bool) {
+        let state = if allocated {
+            ChunkState::allocated(self.space_index)
+        } else {
+            ChunkState::free()
+        };
         // Do nothing if the chunk is already in the expected state.
-        if self.get(chunk) == state {
+        if self.get_internal(chunk) == state {
             return;
         }
+        #[cfg(debug_assertions)]
+        {
+            let old_state = self.get_internal(chunk);
+            // If a chunk is free, any space may use it. If a chunk is not free, only the current space may update its state.
+            assert!(
+                old_state.is_free() || old_state.get_space_index() == self.space_index,
+                "Chunk {:?}: old state {:?}, new state {:?}. Cannot set to new state.",
+                chunk,
+                old_state,
+                state
+            );
+        }
         // Update alloc byte
-        unsafe { Self::ALLOC_TABLE.store::<u8>(chunk.start(), state as u8) };
+        unsafe { Self::ALLOC_TABLE.store::<u8>(chunk.start(), state.0) };
         // If this is a newly allcoated chunk, then expand the chunk range.
-        if state == ChunkState::Allocated {
+        if allocated {
             debug_assert!(!chunk.start().is_zero());
             let mut range = self.chunk_range.lock();
             if range.start == Chunk::ZERO {
@@ -96,20 +162,23 @@ impl ChunkMap {
         }
     }
 
-    /// Get chunk state
-    pub fn get(&self, chunk: Chunk) -> ChunkState {
+    /// Get chunk state. Return None if the chunk does not belong to the space.
+    pub fn get(&self, chunk: Chunk) -> Option<ChunkState> {
+        let state = self.get_internal(chunk);
+        (state.is_allocated() && state.get_space_index() == self.space_index).then_some(state)
+    }
+
+    /// Get chunk state, regardless of the space. This should always be private.
+    fn get_internal(&self, chunk: Chunk) -> ChunkState {
         let byte = unsafe { Self::ALLOC_TABLE.load::<u8>(chunk.start()) };
-        match byte {
-            0 => ChunkState::Free,
-            1 => ChunkState::Allocated,
-            _ => unreachable!(),
-        }
+        ChunkState(byte)
     }
 
-    /// A range of all chunks in the heap.
-    pub fn all_chunks(&self) -> RegionIterator<Chunk> {
+    /// A range of all allocated chunks by this space in the heap.
+    pub fn all_chunks(&self) -> impl Iterator<Item = Chunk> + '_ {
         let chunk_range = self.chunk_range.lock();
         RegionIterator::<Chunk>::new(chunk_range.start, chunk_range.end)
+            .filter(|c| self.get(*c).is_some())
     }
 
     /// Helper function to create per-chunk processing work packets for each allocated chunks.
@@ -118,18 +187,9 @@ impl ChunkMap {
         func: impl Fn(Chunk) -> Box<dyn GCWork<VM>>,
     ) -> Vec<Box<dyn GCWork<VM>>> {
         let mut work_packets: Vec<Box<dyn GCWork<VM>>> = vec![];
-        for chunk in self
-            .all_chunks()
-            .filter(|c| self.get(*c) == ChunkState::Allocated)
-        {
+        for chunk in self.all_chunks() {
             work_packets.push(func(chunk));
         }
         work_packets
     }
 }
-
-impl Default for ChunkMap {
-    fn default() -> Self {
-        Self::new()
-    }
-}

src/util/metadata/side_metadata/global.rs

@@ -1345,6 +1345,11 @@ impl SideMetadataContext {
             }
         }
 
+        // Any plan that uses the chunk map needs to reserve the chunk map table.
+        // As we use either the mark sweep or (non moving) immix as the non moving space,
+        // and both policies use the chunk map, we just add the chunk map table globally.
+        ret.push(crate::util::heap::chunk_map::ChunkMap::ALLOC_TABLE);
+
         ret.extend_from_slice(specs);
         ret
     }

src/util/metadata/side_metadata/spec_defs.rs

@@ -60,6 +60,8 @@ define_side_metadata_specs!(
     MS_ACTIVE_CHUNK = (global: true, log_num_of_bits: 3, log_bytes_in_region: LOG_BYTES_IN_CHUNK),
     // Track the index in SFT map for a chunk (only used for SFT sparse chunk map)
     SFT_DENSE_CHUNK_MAP_INDEX   = (global: true, log_num_of_bits: 3, log_bytes_in_region: LOG_BYTES_IN_CHUNK),
+    // Mark chunks (any plan that uses the chunk map should include this spec in their global sidemetadata specs)
+    CHUNK_MARK   = (global: true, log_num_of_bits: 3, log_bytes_in_region: crate::util::heap::chunk_map::Chunk::LOG_BYTES),
 );
 
 // This defines all LOCAL side metadata used by mmtk-core.
@@ -75,8 +77,6 @@ define_side_metadata_specs!(
     IX_BLOCK_DEFRAG = (global: false, log_num_of_bits: 3, log_bytes_in_region: crate::policy::immix::block::Block::LOG_BYTES),
     // Mark blocks by immix
     IX_BLOCK_MARK   = (global: false, log_num_of_bits: 3, log_bytes_in_region: crate::policy::immix::block::Block::LOG_BYTES),
-    // Mark chunks (any plan that uses the chunk map should include this spec in their local sidemetadata specs)
-    CHUNK_MARK   = (global: false, log_num_of_bits: 3, log_bytes_in_region: crate::util::heap::chunk_map::Chunk::LOG_BYTES),
     // Mark blocks by (native mimalloc) marksweep
     MS_BLOCK_MARK   = (global: false, log_num_of_bits: 3, log_bytes_in_region: crate::policy::marksweepspace::native_ms::Block::LOG_BYTES),
     // Next block in list for native mimalloc

src/util/object_enum.rs

@@ -5,10 +5,7 @@ use std::marker::PhantomData;
 use crate::vm::VMBinding;
 
 use super::{
-    heap::{
-        chunk_map::{ChunkMap, ChunkState},
-        MonotonePageResource,
-    },
+    heap::{chunk_map::ChunkMap, MonotonePageResource},
     linear_scan::Region,
     metadata::{side_metadata::spec_defs::VO_BIT, vo_bit},
     Address, ObjectReference,
@@ -85,11 +82,9 @@ pub(crate) fn enumerate_blocks_from_chunk_map<B>(
     B: BlockMayHaveObjects,
 {
     for chunk in chunk_map.all_chunks() {
-        if chunk_map.get(chunk) == ChunkState::Allocated {
-            for block in chunk.iter_region::<B>() {
-                if block.may_have_objects() {
-                    enumerator.visit_address_range(block.start(), block.end());
-                }
+        for block in chunk.iter_region::<B>() {
+            if block.may_have_objects() {
+                enumerator.visit_address_range(block.start(), block.end());
             }
         }
     }