Generalise weak reference processing for other languages (mmtk#700)

This commit adds a language-neutral API to let the VM binding handle weak
references and finalisers.

Added VMRefClosure and VMRefForwarding work bucket stages to replace the
Java-specific stages.

Added a "sentinel" mechanism to execute a work packet when a bucket is drained,
and prevent the GC from entering the next stage, making it possible to expand
the transitive closure multiple times. It replaces GCWorkScheduler::closure_end.

Extended the Collection::process_weak_refs method to allow the binding to trace
objects during weak reference processing.

Renamed Collection::vm_release to Collection::post_forwarding, and it is now
executed in a dedicated work packet.

Author: wks

src/memory_manager.rs

@@ -900,9 +900,3 @@ pub fn add_work_packets<VM: VMBinding>(
 ) {
     mmtk.scheduler.work_buckets[bucket].bulk_add(packets)
 }
-
-/// Add a callback to be notified after the transitive closure is finished.
-/// The callback should return true if it add more work packets to the closure bucket.
-pub fn on_closure_end<VM: VMBinding>(mmtk: &'static MMTK<VM>, f: Box<dyn Send + Fn() -> bool>) {
-    mmtk.scheduler.on_closure_end(f)
-}

src/plan/markcompact/global.rs

@@ -121,10 +121,6 @@ impl<VM: VMBinding> Plan for MarkCompact<VM> {
             scheduler.work_buckets[WorkBucketStage::PhantomRefClosure]
                 .add(PhantomRefProcessing::<MarkingProcessEdges<VM>>::new());
 
-            // VM-specific weak ref processing
-            scheduler.work_buckets[WorkBucketStage::WeakRefClosure]
-                .add(VMProcessWeakRefs::<MarkingProcessEdges<VM>>::new());
-
             use crate::util::reference_processor::RefForwarding;
             scheduler.work_buckets[WorkBucketStage::RefForwarding]
                 .add(RefForwarding::<ForwardingProcessEdges<VM>>::new());
@@ -147,6 +143,17 @@ impl<VM: VMBinding> Plan for MarkCompact<VM> {
                 .add(ForwardFinalization::<ForwardingProcessEdges<VM>>::new());
         }
 
+        // VM-specific weak ref processing
+        scheduler.work_buckets[WorkBucketStage::VMRefClosure]
+            .set_sentinel(Box::new(VMProcessWeakRefs::<MarkingProcessEdges<VM>>::new()));
+
+        // VM-specific weak ref forwarding
+        scheduler.work_buckets[WorkBucketStage::VMRefForwarding]
+            .add(VMForwardWeakRefs::<ForwardingProcessEdges<VM>>::new());
+
+        // VM-specific work after forwarding, possible to implement ref enququing.
+        scheduler.work_buckets[WorkBucketStage::Release].add(VMPostForwarding::<VM>::default());
+
         // Analysis GC work
         #[cfg(feature = "analysis")]
         {

src/scheduler/gc_work.rs

@@ -108,9 +108,9 @@ impl<C: GCWorkContext> Release<C> {
 impl<C: GCWorkContext + 'static> GCWork<C::VM> for Release<C> {
     fn do_work(&mut self, worker: &mut GCWorker<C::VM>, mmtk: &'static MMTK<C::VM>) {
         trace!("Release Global");
+
         self.plan.base().gc_trigger.policy.on_gc_release(mmtk);
 
-        <C::VM as VMBinding>::VMCollection::vm_release();
         // We assume this is the only running work packet that accesses plan at the point of execution
         #[allow(clippy::cast_ref_to_mut)]
         let plan_mut: &mut C::PlanType = unsafe { &mut *(self.plan as *const _ as *mut _) };
@@ -252,24 +252,190 @@ impl<VM: VMBinding> GCWork<VM> for EndOfGC {
 
 impl<VM: VMBinding> CoordinatorWork<VM> for EndOfGC {}
 
-/// Delegate to the VM binding for reference processing.
+/// This implements `ObjectTracer` by forwarding the `trace_object` calls to the wrapped
+/// `ProcessEdgesWork` instance.
+struct ProcessEdgesWorkTracer<E: ProcessEdgesWork> {
+    process_edges_work: E,
+    stage: WorkBucketStage,
+}
+
+impl<E: ProcessEdgesWork> ObjectTracer for ProcessEdgesWorkTracer<E> {
+    /// Forward the `trace_object` call to the underlying `ProcessEdgesWork`,
+    /// and flush as soon as the underlying buffer of `process_edges_work` is full.
+    ///
+    /// This function is inlined because `trace_object` is probably the hottest function in MMTk.
+    /// If this function is called in small closures, please profile the program and make sure the
+    /// closure is inlined, too.
+    #[inline(always)]
+    fn trace_object(&mut self, object: ObjectReference) -> ObjectReference {
+        let result = self.process_edges_work.trace_object(object);
+        self.flush_if_full();
+        result
+    }
+}
+
+impl<E: ProcessEdgesWork> ProcessEdgesWorkTracer<E> {
+    #[inline(always)]
+    fn flush_if_full(&mut self) {
+        if self.process_edges_work.nodes.is_full() {
+            self.flush();
+        }
+    }
+
+    pub fn flush_if_not_empty(&mut self) {
+        if !self.process_edges_work.nodes.is_empty() {
+            self.flush();
+        }
+    }
+
+    #[cold]
+    fn flush(&mut self) {
+        let next_nodes = self.process_edges_work.pop_nodes();
+        assert!(!next_nodes.is_empty());
+        let work_packet = self.process_edges_work.create_scan_work(next_nodes, false);
+        let worker = self.process_edges_work.worker();
+        worker.scheduler().work_buckets[self.stage].add(work_packet);
+    }
+}
+
+/// This type implements `ObjectTracerContext` by creating a temporary `ProcessEdgesWork` during
+/// the call to `with_tracer`, making use of its `trace_object` method.  It then creates work
+/// packets using the methods of the `ProcessEdgesWork` and add the work packet into the given
+/// `stage`.
+struct ProcessEdgesWorkTracerContext<E: ProcessEdgesWork> {
+    stage: WorkBucketStage,
+    phantom_data: PhantomData<E>,
+}
+
+impl<E: ProcessEdgesWork> Clone for ProcessEdgesWorkTracerContext<E> {
+    fn clone(&self) -> Self {
+        Self { ..*self }
+    }
+}
+
+impl<E: ProcessEdgesWork> ObjectTracerContext<E::VM> for ProcessEdgesWorkTracerContext<E> {
+    type TracerType = ProcessEdgesWorkTracer<E>;
+
+    fn with_tracer<R, F>(&self, worker: &mut GCWorker<E::VM>, func: F) -> R
+    where
+        F: FnOnce(&mut Self::TracerType) -> R,
+    {
+        let mmtk = worker.mmtk;
+
+        // Prepare the underlying ProcessEdgesWork
+        let mut process_edges_work = E::new(vec![], false, mmtk);
+        // FIXME: This line allows us to omit the borrowing lifetime of worker.
+        // We should refactor ProcessEdgesWork so that it uses `worker` locally, not as a member.
+        process_edges_work.set_worker(worker);
+
+        // Cretae the tracer.
+        let mut tracer = ProcessEdgesWorkTracer {
+            process_edges_work,
+            stage: self.stage,
+        };
+
+        // The caller can use the tracer here.
+        let result = func(&mut tracer);
+
+        // Flush the queued nodes.
+        tracer.flush_if_not_empty();
+
+        result
+    }
+}
+
+/// Delegate to the VM binding for weak reference processing.
 ///
 /// Some VMs (e.g. v8) do not have a Java-like global weak reference storage, and the
 /// processing of those weakrefs may be more complex. For such case, we delegate to the
 /// VM binding to process weak references.
-#[derive(Default)]
-pub struct VMProcessWeakRefs<E: ProcessEdgesWork>(PhantomData<E>);
+///
+/// NOTE: This will replace `{Soft,Weak,Phantom}RefProcessing` and `Finalization` in the future.
+pub struct VMProcessWeakRefs<E: ProcessEdgesWork> {
+    phantom_data: PhantomData<E>,
+}
 
 impl<E: ProcessEdgesWork> VMProcessWeakRefs<E> {
     pub fn new() -> Self {
-        Self(PhantomData)
+        Self {
+            phantom_data: PhantomData,
+        }
     }
 }
 
 impl<E: ProcessEdgesWork> GCWork<E::VM> for VMProcessWeakRefs<E> {
     fn do_work(&mut self, worker: &mut GCWorker<E::VM>, _mmtk: &'static MMTK<E::VM>) {
-        trace!("ProcessWeakRefs");
-        <E::VM as VMBinding>::VMCollection::process_weak_refs(worker); // TODO: Pass a factory/callback to decide what work packet to create.
+        trace!("VMProcessWeakRefs");
+
+        let stage = WorkBucketStage::VMRefClosure;
+
+        let need_to_repeat = {
+            let tracer_factory = ProcessEdgesWorkTracerContext::<E> {
+                stage,
+                phantom_data: PhantomData,
+            };
+            <E::VM as VMBinding>::VMScanning::process_weak_refs(worker, tracer_factory)
+        };
+
+        if need_to_repeat {
+            // Schedule Self as the new sentinel so we'll call `process_weak_refs` again after the
+            // current transitive closure.
+            let new_self = Box::new(Self::new());
+
+            worker.scheduler().work_buckets[stage].set_sentinel(new_self);
+        }
+    }
+}
+
+/// Delegate to the VM binding for forwarding weak references.
+///
+/// Some VMs (e.g. v8) do not have a Java-like global weak reference storage, and the
+/// processing of those weakrefs may be more complex. For such case, we delegate to the
+/// VM binding to process weak references.
+///
+/// NOTE: This will replace `RefForwarding` and `ForwardFinalization` in the future.
+pub struct VMForwardWeakRefs<E: ProcessEdgesWork> {
+    phantom_data: PhantomData<E>,
+}
+
+impl<E: ProcessEdgesWork> VMForwardWeakRefs<E> {
+    pub fn new() -> Self {
+        Self {
+            phantom_data: PhantomData,
+        }
+    }
+}
+
+impl<E: ProcessEdgesWork> GCWork<E::VM> for VMForwardWeakRefs<E> {
+    fn do_work(&mut self, worker: &mut GCWorker<E::VM>, _mmtk: &'static MMTK<E::VM>) {
+        trace!("VMForwardWeakRefs");
+
+        let stage = WorkBucketStage::VMRefForwarding;
+
+        let tracer_factory = ProcessEdgesWorkTracerContext::<E> {
+            stage,
+            phantom_data: PhantomData,
+        };
+        <E::VM as VMBinding>::VMScanning::forward_weak_refs(worker, tracer_factory)
+    }
+}
+
+/// This work packet calls `Collection::post_forwarding`.
+///
+/// NOTE: This will replace `RefEnqueue` in the future.
+///
+/// NOTE: Although this work packet runs in parallel with the `Release` work packet, it does not
+/// access the `Plan` instance.
+#[derive(Default)]
+pub struct VMPostForwarding<VM: VMBinding> {
+    phantom_data: PhantomData<VM>,
+}
+
+impl<VM: VMBinding> GCWork<VM> for VMPostForwarding<VM> {
+    fn do_work(&mut self, worker: &mut GCWorker<VM>, _mmtk: &'static MMTK<VM>) {
+        trace!("VMPostForwarding start");
+        <VM as VMBinding>::VMCollection::post_forwarding(worker.tls);
+        trace!("VMPostForwarding end");
     }
 }
 
@@ -678,38 +844,22 @@ pub trait ScanObjectsWork<VM: VMBinding>: GCWork<VM> + Sized {
 
         // If any object does not support edge-enqueuing, we process them now.
         if !scan_later.is_empty() {
-            // We create an instance of E to use its `trace_object` method and its object queue.
-            let mut process_edges_work = Self::E::new(vec![], false, mmtk);
-            let mut closure = |object| process_edges_work.trace_object(object);
-
-            // Scan objects and trace their edges at the same time.
-            for object in scan_later.iter().copied() {
-                <VM as VMBinding>::VMScanning::scan_object_and_trace_edges(
-                    tls,
-                    object,
-                    &mut closure,
-                );
-                self.post_scan_object(object);
-            }
-
-            // Create work packets to scan adjacent objects.  We skip ProcessEdgesWork and create
-            // object-scanning packets directly, because the edges are already traced.
-            if !process_edges_work.nodes.is_empty() {
-                let next_nodes = process_edges_work.nodes.take();
-                let make_packet = |nodes| {
-                    let work_packet = self.make_another(nodes);
-                    memory_manager::add_work_packet(mmtk, WorkBucketStage::Closure, work_packet);
-                };
-
-                // Divide the resulting nodes into appropriately sized packets.
-                if next_nodes.len() <= Self::E::CAPACITY {
-                    make_packet(next_nodes);
-                } else {
-                    for chunk in next_nodes.chunks(Self::E::CAPACITY) {
-                        make_packet(chunk.into());
-                    }
+            let object_tracer_context = ProcessEdgesWorkTracerContext::<Self::E> {
+                stage: WorkBucketStage::Closure,
+                phantom_data: PhantomData,
+            };
+
+            object_tracer_context.with_tracer(worker, |object_tracer| {
+                // Scan objects and trace their edges at the same time.
+                for object in scan_later.iter().copied() {
+                    <VM as VMBinding>::VMScanning::scan_object_and_trace_edges(
+                        tls,
+                        object,
+                        object_tracer,
+                    );
+                    self.post_scan_object(object);
                 }
-            }
+            });
         }
     }
 }