Skip to main content

runmat_gc/
lib.rs

1use parking_lot::{Mutex, RwLock};
2// Use a local trait-alias shim to avoid compile-time dependency ordering issues.
3// Downstream crates in the workspace provide runmat_builtins; during GC unit tests, we provide a minimal Value.
4use runmat_builtins::Value;
5use runmat_time::Instant;
6use std::collections::HashSet;
7use std::sync::atomic::{AtomicBool, Ordering};
8use std::sync::Arc;
9use thiserror::Error;
10
11pub mod allocator;
12pub mod barriers;
13pub mod collector;
14pub mod config;
15pub mod gc_ptr;
16pub use runmat_gc_api::GcPtr;
17pub mod generations;
18pub mod roots;
19pub mod stats;
20
21// Finalizer support
22use std::collections::HashMap;
23
24/// A finalizer that runs when a GC-managed Value is collected.
25///
26/// Finalizers must be fast, non-panicking, and avoid allocating.
27pub trait GcFinalizer: Send + Sync {
28    fn finalize(&self);
29}
30
31static FINALIZERS: once_cell::sync::Lazy<
32    parking_lot::Mutex<HashMap<usize, std::sync::Arc<dyn GcFinalizer>>>,
33> = once_cell::sync::Lazy::new(|| parking_lot::Mutex::new(HashMap::new()));
34
35/// Register a finalizer for the provided GC pointer.
36pub fn gc_register_finalizer(ptr: GcPtr<Value>, f: std::sync::Arc<dyn GcFinalizer>) -> Result<()> {
37    let addr = unsafe { ptr.as_raw() } as usize;
38    if addr == 0 {
39        return Ok(());
40    }
41    FINALIZERS.lock().insert(addr, f);
42    Ok(())
43}
44
45/// Remove any registered finalizer for the provided GC pointer.
46pub fn gc_unregister_finalizer(ptr: GcPtr<Value>) -> Result<()> {
47    let addr = unsafe { ptr.as_raw() } as usize;
48    if addr == 0 {
49        return Ok(());
50    }
51    FINALIZERS.lock().remove(&addr);
52    Ok(())
53}
54
55/// Internal: run and remove finalizer for an address if present.
56pub(crate) fn gc_run_finalizer_for_addr(addr: usize) {
57    if let Some(f) = FINALIZERS.lock().remove(&addr) {
58        // Run finalizer; avoid panicking across GC boundaries
59        f.finalize();
60    }
61}
62
63#[cfg(feature = "pointer-compression")]
64pub mod compression;
65
66pub use allocator::{AllocatorStats, GenerationalAllocator, SizeClass};
67use barriers::WriteBarrierManager;
68pub use barriers::{CardTable, WriteBarrier};
69pub use collector::MarkSweepCollector;
70pub use config::{GcConfig, GcConfigBuilder};
71// Re-export unified handle from API crate
72pub use generations::{Generation, GenerationStats, GenerationalHeap, GenerationalHeapStats};
73pub use roots::{GcRoot, GlobalRoot, RootId, RootScanner, StackRoot, VariableArrayRoot};
74pub use stats::{CollectionEvent, CollectionType, GcStats};
75
76#[cfg(feature = "pointer-compression")]
77pub use compression::*;
78
79/// Errors that can occur during garbage collection
80#[derive(Error, Debug)]
81pub enum GcError {
82    #[error("Out of memory: {0}")]
83    OutOfMemory(String),
84
85    #[error("Invalid GC pointer: {0}")]
86    InvalidPointer(String),
87
88    #[error("Collection failed: {0}")]
89    CollectionFailed(String),
90
91    #[error("Root registration failed: {0}")]
92    RootRegistrationFailed(String),
93
94    #[error("Configuration error: {0}")]
95    ConfigError(String),
96
97    #[error("Thread synchronization error: {0}")]
98    SyncError(String),
99}
100
101pub type Result<T> = std::result::Result<T, GcError>;
102
103// Legacy handle/object table removed in favor of allocator-backed pointers and address-keyed maps
104
105/// Garbage collector with safe handle-based design
106pub struct GarbageCollector {
107    /// Configuration
108    config: Arc<RwLock<GcConfig>>,
109
110    /// Generational allocator (owns heap memory)
111    allocator: Mutex<GenerationalAllocator>,
112
113    /// Mark-and-sweep collector
114    collector: Mutex<MarkSweepCollector>,
115
116    /// Explicit roots stored as raw addresses (address-keyed)
117    root_ptrs: Arc<Mutex<HashSet<usize>>>,
118
119    /// Collection state
120    collection_in_progress: AtomicBool,
121
122    /// Statistics
123    stats: Arc<GcStats>,
124}
125
126impl GarbageCollector {
127    pub fn new() -> Result<Self> {
128        Self::with_config(GcConfig::default())
129    }
130
131    pub fn with_config(config: GcConfig) -> Result<Self> {
132        let allocator = GenerationalAllocator::new(&config);
133        let collector = MarkSweepCollector::new(&config);
134
135        Ok(Self {
136            config: Arc::new(RwLock::new(config)),
137            allocator: Mutex::new(allocator),
138            collector: Mutex::new(collector),
139            root_ptrs: Arc::new(Mutex::new(HashSet::new())),
140            collection_in_progress: AtomicBool::new(false),
141            stats: Arc::new(GcStats::new()),
142        })
143    }
144
145    /// Allocate a new Value object
146    pub fn allocate(&self, value: Value) -> Result<GcPtr<Value>> {
147        // Check if collection is needed
148        if self.should_collect() {
149            let _ = self.collect_minor();
150        }
151
152        // Capture GPU handle if present to attach a finalizer after allocation
153        let gpu_handle: Option<runmat_accelerate_api::GpuTensorHandle> =
154            if let Value::GpuTensor(h) = &value {
155                Some(h.clone())
156            } else {
157                None
158            };
159
160        let mut allocator = self.allocator.lock();
161        let ptr = allocator.allocate(value, &self.stats)?;
162        let usage = allocator.young_generation_usage();
163        let alloc_count = allocator.young_allocations_count();
164        let cfg = self.config.read().clone();
165        drop(allocator);
166
167        // Register finalizer for GPU tensors so buffers are freed on collection
168        if let Some(handle) = gpu_handle {
169            struct GpuTensorFinalizer {
170                handle: runmat_accelerate_api::GpuTensorHandle,
171            }
172            impl GcFinalizer for GpuTensorFinalizer {
173                fn finalize(&self) {
174                    if let Some(p) = runmat_accelerate_api::provider() {
175                        let _ = p.free(&self.handle);
176                        runmat_accelerate_api::clear_handle_logical(&self.handle);
177                    }
178                }
179            }
180            let fin = std::sync::Arc::new(GpuTensorFinalizer { handle });
181            let _ = gc_register_finalizer(ptr.clone(), fin);
182        }
183
184        // Heuristic triggers:
185        // - Utilization threshold
186        // - Aggressive mode: periodic minor GC by allocation count to satisfy stress configs
187        if usage > cfg.minor_gc_threshold
188            || (cfg.minor_gc_threshold <= 0.35 && alloc_count > 0 && alloc_count.is_multiple_of(32))
189        {
190            let _ = self.collect_minor();
191        }
192        Ok(ptr)
193    }
194
195    // get_value by handle removed in favor of direct GcPtr usage
196
197    /// Check if collection should be triggered
198    fn should_collect(&self) -> bool {
199        if self.collection_in_progress.load(Ordering::Acquire) {
200            return false;
201        }
202
203        let threshold = self.config.read().minor_gc_threshold;
204        let allocator = self.allocator.lock();
205        let utilization = allocator.young_generation_usage();
206        utilization > threshold
207    }
208
209    /// Perform minor collection (young generation)
210    pub fn collect_minor(&self) -> Result<usize> {
211        if self
212            .collection_in_progress
213            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
214            .is_err()
215        {
216            return Ok(0);
217        }
218
219        let start_time = Instant::now();
220
221        // Build combined roots: explicit + external (root scanner + barriers)
222        let mut combined_roots: Vec<GcPtr<Value>> = Vec::new();
223        {
224            let roots = self.root_ptrs.lock();
225            combined_roots.extend(
226                roots
227                    .iter()
228                    .map(|&addr| unsafe { GcPtr::from_raw(addr as *const Value) }),
229            );
230        }
231        // External roots
232        if let Ok(mut ext) = ROOT_SCANNER.scan_roots() {
233            combined_roots.append(&mut ext);
234        }
235        combined_roots.extend(gc_barrier_minor_roots());
236
237        // Collect young generation via unified collector/allocator
238        let mut allocator = self.allocator.lock();
239        let mut collector = self.collector.lock();
240        let collected_count =
241            collector.collect_young_generation(&mut allocator, &combined_roots, &self.stats)?;
242
243        // Clear dirty cards; keep remembered set for next minor cycle
244        WRITE_BARRIERS.clear_after_minor_gc();
245
246        let duration = start_time.elapsed();
247        self.stats
248            .record_minor_collection(collected_count, duration);
249
250        self.collection_in_progress.store(false, Ordering::Release);
251
252        log::info!("Minor GC: collected {collected_count} objects in {duration:?}");
253
254        Ok(collected_count)
255    }
256
257    // /// Get object by handle
258    // Address-keyed design: object table lookup removed
259
260    // /// Mark objects transitively (follow references)
261    // Transitive marking handled by collector over allocator-backed objects
262
263    // /// Find the GcObject that contains a specific Value
264    // Value-equality scans removed in favor of address-keyed marking
265
266    /// Perform major collection (all generations)
267    pub fn collect_major(&self) -> Result<usize> {
268        if self
269            .collection_in_progress
270            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
271            .is_err()
272        {
273            return Ok(0);
274        }
275
276        let start_time = Instant::now();
277
278        // Build combined roots: explicit + external
279        let mut combined_roots: Vec<GcPtr<Value>> = Vec::new();
280        {
281            let roots = self.root_ptrs.lock();
282            combined_roots.extend(
283                roots
284                    .iter()
285                    .map(|&addr| unsafe { GcPtr::from_raw(addr as *const Value) }),
286            );
287        }
288        if let Ok(mut ext) = ROOT_SCANNER.scan_roots() {
289            combined_roots.append(&mut ext);
290        }
291        combined_roots.extend(gc_barrier_minor_roots());
292
293        let mut allocator = self.allocator.lock();
294        let mut collector = self.collector.lock();
295        let collected_count =
296            collector.collect_all_generations(&mut allocator, &combined_roots, &self.stats)?;
297
298        // Clear all barriers after major GC
299        WRITE_BARRIERS.clear_after_major_gc();
300        allocator.clear_promotion_state();
301
302        let duration = start_time.elapsed();
303        self.stats
304            .record_major_collection(collected_count, duration);
305
306        self.collection_in_progress.store(false, Ordering::Release);
307
308        log::info!("Major GC: collected {collected_count} objects in {duration:?}");
309
310        Ok(collected_count)
311    }
312
313    /// Add a root to protect an object from collection
314    pub fn add_root(&self, root: GcPtr<Value>) -> Result<()> {
315        let value_ptr = unsafe { root.as_raw() } as usize;
316        if value_ptr == 0 {
317            return Ok(());
318        }
319        self.root_ptrs.lock().insert(value_ptr);
320        Ok(())
321    }
322
323    /// Remove a root
324    pub fn remove_root(&self, root: GcPtr<Value>) -> Result<()> {
325        let value_ptr = unsafe { root.as_raw() } as usize;
326        if value_ptr == 0 {
327            return Ok(());
328        }
329        self.root_ptrs.lock().remove(&value_ptr);
330        Ok(())
331    }
332
333    // /// Find the handle for an object that contains the given value pointer
334    // Address-keyed roots: no handle lookup needed
335    /// Get GC statistics
336    pub fn stats(&self) -> GcStats {
337        self.stats.as_ref().clone()
338    }
339
340    /// Configure the GC
341    pub fn configure(&self, config: GcConfig) -> Result<()> {
342        *self.config.write() = config.clone();
343        {
344            // Rebuild allocator to support changes like num_generations and sizes
345            let mut allocator = self.allocator.lock();
346            *allocator = GenerationalAllocator::new(&config);
347        }
348        {
349            let mut collector = self.collector.lock();
350            collector.reconfigure(&config)?;
351        }
352        Ok(())
353    }
354
355    /// Get the current GC configuration
356    pub fn get_config(&self) -> GcConfig {
357        self.config.read().clone()
358    }
359}
360
361// Safety: All shared data is protected by proper synchronization
362unsafe impl Send for GarbageCollector {}
363unsafe impl Sync for GarbageCollector {}
364
365/// Global garbage collector instance
366static GC: once_cell::sync::Lazy<Arc<GarbageCollector>> =
367    once_cell::sync::Lazy::new(|| Arc::new(GarbageCollector::new().expect("Failed to create GC")));
368static ROOT_SCANNER: once_cell::sync::Lazy<Arc<RootScanner>> =
369    once_cell::sync::Lazy::new(|| Arc::new(RootScanner::new()));
370
371/// Global write barrier manager
372static WRITE_BARRIERS: once_cell::sync::Lazy<Arc<WriteBarrierManager>> =
373    once_cell::sync::Lazy::new(|| Arc::new(WriteBarrierManager::new(true, false)));
374
375/// Helper function to dereference a GcPtr safely (now just uses normal dereferencing)
376pub fn gc_deref(ptr: GcPtr<Value>) -> Value {
377    (*ptr).clone()
378}
379
380/// Global GC functions for easy access
381pub fn gc_allocate(value: Value) -> Result<GcPtr<Value>> {
382    GC.allocate(value)
383}
384
385pub fn gc_collect_minor() -> Result<usize> {
386    // Stage external roots inside GC and perform unified minor collection
387    GC.collect_minor()
388}
389
390pub fn gc_collect_major() -> Result<usize> {
391    GC.collect_major()
392}
393
394pub fn gc_add_root(root: GcPtr<Value>) -> Result<()> {
395    GC.add_root(root)
396}
397
398pub fn gc_remove_root(root: GcPtr<Value>) -> Result<()> {
399    GC.remove_root(root)
400}
401
402pub fn gc_stats() -> GcStats {
403    GC.stats()
404}
405
406pub fn gc_configure(config: GcConfig) -> Result<()> {
407    GC.configure(config)
408}
409
410pub fn gc_get_config() -> GcConfig {
411    GC.get_config()
412}
413
414/// Simplified root registration for backwards compatibility
415pub fn gc_register_root(root: Box<dyn GcRoot>) -> Result<RootId> {
416    ROOT_SCANNER.register_root(root)
417}
418
419pub fn gc_unregister_root(root_id: RootId) -> Result<()> {
420    ROOT_SCANNER.unregister_root(root_id)
421}
422
423/// Record a write for GC barriers (approximate old->young tracking)
424pub fn gc_record_write(old: &Value, new: &Value) {
425    // Generation-aware barrier: record only if old is logically old and new is young
426    let old_ptr = old as *const Value as *const u8;
427    let young_ptr = new as *const Value as *const u8;
428    // Query allocator logical generations
429    let alloc = GC.allocator.lock();
430    let old_gen = alloc.logical_generation(old_ptr).unwrap_or(0);
431    let new_gen = alloc.logical_generation(young_ptr).unwrap_or(0);
432    drop(alloc);
433    if old_gen > new_gen {
434        WRITE_BARRIERS.record_reference(old_ptr, young_ptr);
435    }
436}
437
438/// Get barrier-derived roots for minor GC
439pub fn gc_barrier_minor_roots() -> Vec<GcPtr<Value>> {
440    WRITE_BARRIERS
441        .get_minor_gc_roots()
442        .into_iter()
443        .map(|p| unsafe { GcPtr::from_raw(p as *const Value) })
444        .collect()
445}
446
447/// Force a garbage collection for testing/debugging
448#[cfg(any(test, feature = "debug-gc"))]
449pub fn gc_force_collect() -> Result<usize> {
450    gc_collect_major()
451}
452
453/// Reset GC for testing - always available
454pub fn gc_reset_for_test() -> Result<()> {
455    // Reset statistics
456    GC.stats.reset();
457
458    // Reset allocator/collector and roots
459    {
460        let config = GcConfig::default();
461        *GC.config.write() = config.clone();
462        let mut alloc = GC.allocator.lock();
463        *alloc = GenerationalAllocator::new(&config);
464        let mut coll = GC.collector.lock();
465        *coll = MarkSweepCollector::new(&config);
466    }
467
468    {
469        let mut roots = GC.root_ptrs.lock();
470        roots.clear();
471    }
472
473    // Ensure collection is not in progress
474    GC.collection_in_progress.store(false, Ordering::Relaxed);
475
476    // Configure with default settings
477    gc_configure(GcConfig::default())?;
478
479    Ok(())
480}
481
482/// Create an isolated test context with clean GC state
483pub fn gc_test_context<F, R>(test_fn: F) -> R
484where
485    F: FnOnce() -> R,
486{
487    // Use a global test mutex to ensure tests run sequentially when needed
488    static TEST_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());
489
490    // Handle poisoned mutex gracefully
491    let _guard = match TEST_MUTEX.lock() {
492        Ok(guard) => guard,
493        Err(poisoned) => {
494            // Clear the poison and continue
495            poisoned.into_inner()
496        }
497    };
498
499    // Reset GC to clean state
500    gc_reset_for_test().expect("GC reset should succeed");
501
502    // Run the test
503    let result = test_fn();
504
505    // Clean up after test (optional, but good practice)
506    let _ = gc_reset_for_test();
507
508    result
509}
510
511#[cfg(test)]
512mod tests {
513    use super::*;
514
515    #[test]
516    fn test_basic_allocation() {
517        gc_test_context(|| {
518            let value = Value::Num(42.0);
519            let ptr = gc_allocate(value).expect("allocation failed");
520            assert_eq!(*ptr, Value::Num(42.0));
521        });
522    }
523
524    #[test]
525    fn test_collection() {
526        gc_test_context(|| {
527            let mut _ptrs = Vec::new();
528            for i in 0..100 {
529                let ptr = gc_allocate(Value::Num(i as f64)).expect("allocation failed");
530                _ptrs.push(ptr);
531            }
532
533            let _collected = gc_collect_minor().expect("collection failed");
534            drop(_ptrs);
535        });
536    }
537
538    #[test]
539    fn test_thread_safety() {
540        use std::thread;
541
542        gc_test_context(|| {
543            let handles: Vec<_> = (0..4)
544                .map(|i| {
545                    thread::spawn(move || {
546                        let mut ptrs = Vec::new();
547                        for j in 0..100 {
548                            let value = Value::Num((i * 100 + j) as f64);
549                            let ptr = gc_allocate(value).expect("allocation failed");
550                            ptrs.push(ptr);
551                        }
552                        ptrs
553                    })
554                })
555                .collect();
556
557            for handle in handles {
558                let _ptrs = handle.join().expect("thread failed");
559            }
560
561            let _ = gc_collect_major();
562        });
563    }
564
565    #[test]
566    fn test_root_protection() {
567        gc_test_context(|| {
568            let protected = gc_allocate(Value::Num(42.0)).expect("allocation failed");
569            gc_add_root(protected.clone()).expect("root registration failed");
570
571            for i in 0..60 {
572                let _ = gc_allocate(Value::String(format!("garbage_{i}")));
573            }
574
575            let _ = gc_collect_minor().expect("collection failed");
576            assert_eq!(*protected, Value::Num(42.0));
577
578            gc_remove_root(protected).expect("root removal failed");
579        });
580    }
581
582    #[test]
583    fn test_gc_allocation_and_roots() {
584        gc_test_context(|| {
585            let v = gc_allocate(Value::Num(7.0)).expect("allocation failed");
586            assert_eq!(*v, Value::Num(7.0));
587
588            gc_add_root(v.clone()).expect("root add failed");
589            let _ = gc_collect_minor().expect("collection failed");
590            assert_eq!(*v, Value::Num(7.0));
591
592            gc_remove_root(v).expect("root remove failed");
593        });
594    }
595}