Add lock-free BlockPageResource for ImmixSpace (#506)

This PR adds a lock-free block queue and uses it for Immix's clean block allocation (as `BlockPageResource`) and recyclable block allocation (as a reusable block queue).

The lock-free block queue has per-thread producer endpoints. Only GC workers can add blocks to the queue. Workers will add the blocks to their thread-local queues first, and flush them to the global pool if the queue is full.

Block allocation can only happen at block granularity. The consumers still share a single endpoint. But on block allocation fast-path, the cost is just a simple atomic increment to update the allocation cursor and pop a block.  _However, for `BlockPageResource`, considering the current design of `Space.acquire`, it still acquires a lock before calling `pr.get_new_pages`._

Author: wenyuzhao

Cargo.toml

@@ -33,7 +33,7 @@ num_cpus = "1.8"
 enum-map = "=2.1"
 downcast-rs = "1.1.1"
 atomic-traits = "0.2.0"
-atomic = "0.4.6"
+atomic = "0.5.1"
 num-traits = "0.2"
 spin = "0.5.2"
 env_logger = "0.8.2"

src/policy/immix/block.rs

@@ -3,11 +3,11 @@ use super::defrag::Histogram;
 use super::line::Line;
 use super::ImmixSpace;
 use crate::util::constants::*;
+use crate::util::heap::blockpageresource::BlockPool;
 use crate::util::linear_scan::{Region, RegionIterator};
 use crate::util::metadata::side_metadata::{MetadataByteArrayRef, SideMetadataSpec};
 use crate::util::Address;
 use crate::vm::*;
-use spin::{Mutex, MutexGuard};
 use std::sync::atomic::Ordering;
 
 /// The block allocation state.
@@ -271,39 +271,51 @@ impl Block {
 }
 
 /// A non-block single-linked list to store blocks.
-#[derive(Default)]
-pub struct BlockList {
-    queue: Mutex<Vec<Block>>,
+pub struct ReusableBlockPool {
+    queue: BlockPool<Block>,
+    num_workers: usize,
 }
 
-impl BlockList {
+impl ReusableBlockPool {
+    /// Create empty block list
+    pub fn new(num_workers: usize) -> Self {
+        Self {
+            queue: BlockPool::new(num_workers),
+            num_workers,
+        }
+    }
+
     /// Get number of blocks in this list.
-    #[inline]
+    #[inline(always)]
     pub fn len(&self) -> usize {
-        self.queue.lock().len()
+        self.queue.len()
     }
 
     /// Add a block to the list.
-    #[inline]
+    #[inline(always)]
     pub fn push(&self, block: Block) {
-        self.queue.lock().push(block)
+        self.queue.push(block)
     }
 
     /// Pop a block out of the list.
-    #[inline]
+    #[inline(always)]
     pub fn pop(&self) -> Option<Block> {
-        self.queue.lock().pop()
+        self.queue.pop()
     }
 
     /// Clear the list.
-    #[inline]
-    pub fn reset(&self) {
-        *self.queue.lock() = Vec::new()
+    pub fn reset(&mut self) {
+        self.queue = BlockPool::new(self.num_workers);
     }
 
-    /// Get an array of all reusable blocks stored in this BlockList.
+    /// Iterate all the blocks in the queue. Call the visitor for each reported block.
     #[inline]
-    pub fn get_blocks(&self) -> MutexGuard<Vec<Block>> {
-        self.queue.lock()
+    pub fn iterate_blocks(&self, mut f: impl FnMut(Block)) {
+        self.queue.iterate_blocks(&mut f);
+    }
+
+    /// Flush the block queue
+    pub fn flush_all(&self) {
+        self.queue.flush_all();
     }
 }

src/policy/immix/chunk.rs

@@ -10,6 +10,8 @@ use crate::{
     MMTK,
 };
 use spin::Mutex;
+use std::sync::atomic::AtomicUsize;
+use std::sync::Arc;
 use std::{ops::Range, sync::atomic::Ordering};
 
 /// Data structure to reference a MMTk 4 MB chunk.
@@ -160,7 +162,19 @@ impl ChunkMap {
         space: &'static ImmixSpace<VM>,
     ) -> Vec<Box<dyn GCWork<VM>>> {
         space.defrag.mark_histograms.lock().clear();
-        self.generate_tasks(|chunk| Box::new(SweepChunk { space, chunk }))
+        let epilogue = Arc::new(FlushPageResource {
+            space,
+            counter: AtomicUsize::new(0),
+        });
+        let tasks = self.generate_tasks(|chunk| {
+            Box::new(SweepChunk {
+                space,
+                chunk,
+                epilogue: epilogue.clone(),
+            })
+        });
+        epilogue.counter.store(tasks.len(), Ordering::SeqCst);
+        tasks
     }
 }
 
@@ -174,6 +188,8 @@ impl Default for ChunkMap {
 struct SweepChunk<VM: VMBinding> {
     space: &'static ImmixSpace<VM>,
     chunk: Chunk,
+    /// A destructor invoked when all `SweepChunk` packets are finished.
+    epilogue: Arc<FlushPageResource<VM>>,
 }
 
 impl<VM: VMBinding> GCWork<VM> for SweepChunk<VM> {
@@ -184,5 +200,23 @@ impl<VM: VMBinding> GCWork<VM> for SweepChunk<VM> {
             self.chunk.sweep(self.space, &mut histogram);
         }
         self.space.defrag.add_completed_mark_histogram(histogram);
+        self.epilogue.finish_one_work_packet();
+    }
+}
+
+/// Count number of remaining work pacets, and flush page resource if all packets are finished.
+struct FlushPageResource<VM: VMBinding> {
+    space: &'static ImmixSpace<VM>,
+    counter: AtomicUsize,
+}
+
+impl<VM: VMBinding> FlushPageResource<VM> {
+    /// Called after a related work packet is finished.
+    fn finish_one_work_packet(&self) {
+        if 1 == self.counter.fetch_sub(1, Ordering::SeqCst) {
+            // We've finished releasing all the dead blocks to the BlockPageResource's thread-local queues.
+            // Now flush the BlockPageResource.
+            self.space.flush_page_resource()
+        }
     }
 }

src/policy/immix/defrag.rs

@@ -139,7 +139,7 @@ impl Defrag {
         spill_avail_histograms: &mut Histogram,
     ) -> usize {
         let mut total_available_lines = 0;
-        for block in space.reusable_blocks.get_blocks().iter() {
+        space.reusable_blocks.iterate_blocks(|block| {
             let bucket = block.get_holes();
             let unavailable_lines = match block.get_state() {
                 BlockState::Reusable { unavailable_lines } => unavailable_lines as usize,
@@ -148,7 +148,7 @@ impl Defrag {
             let available_lines = Block::LINES - unavailable_lines;
             spill_avail_histograms[bucket] += available_lines;
             total_available_lines += available_lines;
-        }
+        });
         total_available_lines
     }
 

src/policy/immix/immixspace.rs

@@ -11,6 +11,7 @@ use crate::policy::space::SpaceOptions;
 use crate::policy::space::{CommonSpace, Space};
 use crate::util::copy::*;
 use crate::util::heap::layout::heap_layout::{Mmapper, VMMap};
+use crate::util::heap::BlockPageResource;
 use crate::util::heap::HeapMeta;
 use crate::util::heap::PageResource;
 use crate::util::heap::VMRequest;
@@ -23,10 +24,7 @@ use crate::vm::*;
 use crate::{
     plan::ObjectQueue,
     scheduler::{GCWork, GCWorkScheduler, GCWorker, WorkBucketStage},
-    util::{
-        heap::FreeListPageResource,
-        opaque_pointer::{VMThread, VMWorkerThread},
-    },
+    util::opaque_pointer::{VMThread, VMWorkerThread},
     MMTK,
 };
 use atomic::Ordering;
@@ -37,15 +35,15 @@ pub(crate) const TRACE_KIND_DEFRAG: TraceKind = 1;
 
 pub struct ImmixSpace<VM: VMBinding> {
     common: CommonSpace<VM>,
-    pr: FreeListPageResource<VM>,
+    pr: BlockPageResource<VM, Block>,
     /// Allocation status for all chunks in immix space
     pub chunk_map: ChunkMap,
     /// Current line mark state
     pub line_mark_state: AtomicU8,
     /// Line mark state in previous GC
     line_unavail_state: AtomicU8,
     /// A list of all reusable blocks
-    pub reusable_blocks: BlockList,
+    pub reusable_blocks: ReusableBlockPool,
     /// Defrag utilities
     pub(super) defrag: Defrag,
     /// Object mark state
@@ -216,21 +214,38 @@ impl<VM: VMBinding> ImmixSpace<VM> {
         );
         ImmixSpace {
             pr: if common.vmrequest.is_discontiguous() {
-                FreeListPageResource::new_discontiguous(vm_map)
+                BlockPageResource::new_discontiguous(
+                    Block::LOG_PAGES,
+                    vm_map,
+                    scheduler.num_workers(),
+                )
             } else {
-                FreeListPageResource::new_contiguous(common.start, common.extent, vm_map)
+                BlockPageResource::new_contiguous(
+                    Block::LOG_PAGES,
+                    common.start,
+                    common.extent,
+                    vm_map,
+                    scheduler.num_workers(),
+                )
             },
             common,
             chunk_map: ChunkMap::new(),
             line_mark_state: AtomicU8::new(Line::RESET_MARK_STATE),
             line_unavail_state: AtomicU8::new(Line::RESET_MARK_STATE),
-            reusable_blocks: BlockList::default(),
+            reusable_blocks: ReusableBlockPool::new(scheduler.num_workers()),
             defrag: Defrag::default(),
             mark_state: Self::UNMARKED_STATE,
             scheduler,
         }
     }
 
+    /// Flush the thread-local queues in BlockPageResource
+    pub fn flush_page_resource(&self) {
+        self.reusable_blocks.flush_all();
+        #[cfg(target_pointer_width = "64")]
+        self.pr.flush_all()
+    }
+
     /// Get the number of defrag headroom pages.
     pub fn defrag_headroom_pages(&self) -> usize {
         self.defrag.defrag_headroom_pages(self)
@@ -339,7 +354,7 @@ impl<VM: VMBinding> ImmixSpace<VM> {
     /// Release a block.
     pub fn release_block(&self, block: Block) {
         block.deinit();
-        self.pr.release_pages(block.start());
+        self.pr.release_block(block);
     }
 
     /// Allocate a clean block.

src/scheduler/mod.rs

@@ -19,6 +19,7 @@ mod work_bucket;
 pub use work_bucket::WorkBucketStage;
 
 mod worker;
+pub(crate) use worker::current_worker_ordinal;
 pub use worker::GCWorker;
 
 mod controller;

src/scheduler/worker.rs

@@ -5,6 +5,7 @@ use crate::mmtk::MMTK;
 use crate::util::copy::GCWorkerCopyContext;
 use crate::util::opaque_pointer::*;
 use crate::vm::{Collection, GCThreadContext, VMBinding};
+use atomic::Atomic;
 use atomic_refcell::{AtomicRef, AtomicRefCell, AtomicRefMut};
 use crossbeam::deque::{self, Stealer};
 use crossbeam::queue::ArrayQueue;
@@ -15,6 +16,17 @@ use std::sync::{Arc, Mutex};
 /// Represents the ID of a GC worker thread.
 pub type ThreadId = usize;
 
+thread_local! {
+    /// Current worker's ordinal
+    static WORKER_ORDINAL: Atomic<Option<ThreadId>> = Atomic::new(None);
+}
+
+/// Get current worker ordinal. Return `None` if the current thread is not a worker.
+#[inline(always)]
+pub fn current_worker_ordinal() -> Option<ThreadId> {
+    WORKER_ORDINAL.with(|x| x.load(Ordering::Relaxed))
+}
+
 /// The part shared between a GCWorker and the scheduler.
 /// This structure is used for communication, e.g. adding new work packets.
 pub struct GCWorkerShared<VM: VMBinding> {
@@ -173,6 +185,7 @@ impl<VM: VMBinding> GCWorker<VM> {
     /// Entry of the worker thread. Resolve thread affinity, if it has been specified by the user.
     /// Each worker will keep polling and executing work packets in a loop.
     pub fn run(&mut self, tls: VMWorkerThread, mmtk: &'static MMTK<VM>) {
+        WORKER_ORDINAL.with(|x| x.store(Some(self.ordinal), Ordering::SeqCst));
         self.scheduler.resolve_affinity(self.ordinal);
         self.tls = tls;
         self.copy = crate::plan::create_gc_worker_context(tls, mmtk);