Moving GC Extensions to Julia's FFI

base/Base_compiler.jl

@@ -204,6 +204,7 @@ const _DOCS_ALIASING_WARNING = """
 """
 
 include("ctypes.jl")
+include("gc-explicit-pinning.jl")
 include("gcutils.jl")
 include("generator.jl")
 include("runtime_internals.jl")

base/gc-explicit-pinning.jl

@@ -0,0 +1,162 @@
+# This file is a part of Julia. License is MIT: https://julialang.org/license
+
+export increment_pin_count!, decrement_pin_count!,
+       increment_tpin_count!, decrement_tpin_count!,
+       get_pin_count, get_tpin_count
+
+"""
+    increment_pin_count!(obj)
+
+Increment the pin count of `obj` to preserve it beyond a lexical scope.
+
+This ensures that `obj` is not moved or collected by the garbage collector.
+It is crucial for safely passing references to foreign code.
+
+Each call increments the pin count by one. The object remains pinned and
+alive until the count is decremented to zero via `decrement_pin_count!`.
+
+# Examples
+
+```julia
+x = SomeObject()
+increment_pin_count!(x)   # x is now pinned (count = 1)
+
+increment_pin_count!(x)   # pin count is now 2
+"""
+function increment_pin_count!(obj)
+    ccall(:jl_increment_pin_count, Cvoid, (Any,), obj)
+end
+
+"""
+    decrement_pin_count!(obj)
+
+Decrement the pin count of `obj`.
+
+When the count drops to zero, the object is no longer pinned and may be
+moved (or collected by the garbage collector, if no other references
+exist).
+
+This is necessary to release objects that were previously preserved for
+foreign code.
+
+# Examples
+
+```julia
+x = SomeObject()
+increment_pin_count!(x)   # x is now pinned (count = 1)
+
+increment_pin_count!(x)   # pin count is now 2
+
+decrement_pin_count!(x)   # reduces pin count to 1
+
+decrement_pin_count!(x)   # count is 0; x may now be collected
+"""
+function decrement_pin_count!(obj)
+    ccall(:jl_decrement_pin_count, Cvoid, (Any,), obj)
+end
+
+"""
+    increment_tpin_count!(obj)
+
+Increment the transitive pin count of `obj` to preserve it beyond a
+lexical scope.
+
+This ensures that `obj` and any other objects reachable from it are not
+moved or collected by the garbage collector. This is crucial for safely
+passing references to foreign code.
+
+Each call increments the transitive pin count by one. The object remains
+transitively pinned and alive until the count is decremented to zero via
+`decrement_tpin_count!`.
+
+# Examples
+
+```julia
+x = SomeObject()
+increment_tpin_count!(x)   # x is now transitively pinned (count = 1)
+
+increment_tpin_count!(x)   # transitive pin count is now 2
+"""
+function increment_tpin_count!(obj)
+    ccall(:jl_increment_tpin_count, Cvoid, (Any,), obj)
+end
+
+"""
+    decrement_tpin_count!(obj)
+
+Decrement the transitive pin count of `obj`.
+
+When the count drops to zero, `obj` and any objects reachable from it are
+no longer pinned and may be moved (if no other pins exist) or collected by
+the garbage collector (if no other references exist).
+
+This is necessary to release object graphs that were previously preserved
+for foreign code.
+
+# Examples
+
+```julia
+x = SomeObject()
+increment_tpin_count!(x)     # pins x and reachable objects (count = 1)
+
+decrement_tpin_count!(x)     # reduces transitive pin count to 0
+                             # objects may now be collected
+"""
+function decrement_tpin_count!(obj)
+    ccall(:jl_decrement_tpin_count, Cvoid, (Any,), obj)
+end
+
+"""
+    get_pin_count(obj)
+
+Return the current pin count of `obj`.
+
+This indicates how many times `obj` has been explicitly pinned via
+`increment_pin_count!`. A nonzero count means the object is currently
+pinned and will not be moved or collected by the garbage collector (GC).
+
+# Examples
+
+```julia
+x = SomeObject()
+increment_pin_count!(x)
+get_pin_count(x)   # returns 1
+
+increment_pin_count!(x)
+get_pin_count(x)   # returns 2
+
+decrement_pin_count!(x)
+get_pin_count(x)   # returns 1
+"""
+function get_pin_count(obj)
+    c = ccall(:jl_get_pin_count, Csize_t, (Any,), obj)
+    return Int64(c)
+end
+
+"""
+    get_tpin_count(obj)
+
+Return the current transitive pin count of `obj`.
+
+This indicates how many times `obj` has been explicitly transitively pinned
+via `increment_tpin_count!`. A nonzero count means `obj` and all objects
+reachable from it are currently pinned and will not be moved or collected
+by the garbage collector (GC).
+
+# Examples
+
+```julia
+x = SomeObject()
+increment_tpin_count!(x)
+get_tpin_count(x)   # returns 1
+
+increment_tpin_count!(x)
+get_tpin_count(x)   # returns 2
+
+decrement_tpin_count!(x)
+get_tpin_count(x)   # returns 1
+"""
+function get_tpin_count(obj)
+    c = ccall(:jl_get_tpin_count, Csize_t, (Any,), obj)
+    return Int64(c)
+end

src/Makefile

@@ -48,7 +48,7 @@ FLAGS += -I$(LOCALBASE)/include
 endif
 
 # GC source code. It depends on which GC implementation to use.
-GC_SRCS := gc-common gc-stacks gc-alloc-profiler gc-heap-snapshot gc-pinning-log
+GC_SRCS := gc-common gc-stacks gc-alloc-profiler gc-heap-snapshot gc-pinning-log gc-explicit-pinning
 ifeq (${USE_THIRD_PARTY_GC},mmtk)
 GC_SRCS += gc-mmtk
 else

src/gc-explicit-pinning.cpp

@@ -0,0 +1,102 @@
+// This file is a part of Julia. License is MIT: https://julialang.org/license
+
+#include <map>
+#include <mutex>
+
+#include "julia.h"
+
+// Let's just use global maps in this first implementation
+// They could cause contention in multi-threaded code, so we might need to optimize them later
+
+// Pinning
+std::map<void *, size_t> pin_count_map;
+std::mutex pin_count_map_lock;
+// Transitive Pinning
+std::map<void *, size_t> tpin_count_map;
+std::mutex tpin_count_map_lock;
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Pinning
+JL_DLLEXPORT void jl_increment_pin_count(void *obj) {
+    pin_count_map_lock.lock();
+    if (pin_count_map.find(obj) == pin_count_map.end()) {
+        pin_count_map[obj] = 0;
+    }
+    pin_count_map[obj]++;
+    pin_count_map_lock.unlock();
+}
+JL_DLLEXPORT void jl_decrement_pin_count(void *obj) {
+    pin_count_map_lock.lock();
+    auto it = pin_count_map.find(obj);
+    if (it != pin_count_map.end()) {
+        if (it->second == 1) {
+            pin_count_map.erase(it);
+        } else {
+            it->second--;
+        }
+    }
+    pin_count_map_lock.unlock();
+}
+
+// Transitive Pinning
+JL_DLLEXPORT void jl_increment_tpin_count(void *obj) {
+    tpin_count_map_lock.lock();
+    if (tpin_count_map.find(obj) == tpin_count_map.end()) {
+        tpin_count_map[obj] = 0;
+    }
+    tpin_count_map[obj]++;
+    tpin_count_map_lock.unlock();
+}
+JL_DLLEXPORT void jl_decrement_tpin_count(void *obj) {
+    tpin_count_map_lock.lock();
+    auto it = tpin_count_map.find(obj);
+    if (it != tpin_count_map.end()) {
+        if (it->second == 1) {
+            tpin_count_map.erase(it);
+        } else {
+            it->second--;
+        }
+    }
+    tpin_count_map_lock.unlock();
+}
+
+// Retrieve Pinning and Transitive Pinning counts for a given object
+JL_DLLEXPORT size_t jl_get_pin_count(void *obj) {
+    pin_count_map_lock.lock();
+    auto it = pin_count_map.find(obj);
+    size_t count = (it != pin_count_map.end()) ? it->second : 0;
+    pin_count_map_lock.unlock();
+    return count;
+}
+JL_DLLEXPORT size_t jl_get_tpin_count(void *obj) {
+    tpin_count_map_lock.lock();
+    auto it = tpin_count_map.find(obj);
+    size_t count = (it != tpin_count_map.end()) ? it->second : 0;
+    tpin_count_map_lock.unlock();
+    return count;
+}
+
+// Returns all pinned and transitively pinned objects
+// Argument should have been initialized by the caller
+// TODO: add a few assertions to check this?
+JL_DLLEXPORT void jl_dump_all_pinned_objects(arraylist_t *objects) {
+    pin_count_map_lock.lock();
+    for (const auto &pair : pin_count_map) {
+        arraylist_push(objects, pair.first);
+    }
+    pin_count_map_lock.unlock();
+}
+JL_DLLEXPORT void jl_dump_all_tpinned_objects(arraylist_t *objects) {
+    tpin_count_map_lock.lock();
+    for (const auto &pair : tpin_count_map) {
+        arraylist_push(objects, pair.first);
+    }
+    tpin_count_map_lock.unlock();
+}
+
+#ifdef __cplusplus
+}
+#endif

src/gc-mmtk.c

@@ -776,9 +776,12 @@ void trace_partial_globally_rooted(RootsWorkClosure* closure, RootsWorkBuffer* b
 JL_DLLEXPORT void jl_gc_scan_vm_specific_roots(RootsWorkClosure* closure)
 {
 
-    // Create a new buf
+    // Create a new buf for roots and nodes that have been pinned through the FFI functions
     RootsWorkBuffer buf = (closure->report_nodes_func)((void**)0, 0, 0, closure->data, true);
     size_t len = 0;
+    // Create a new buf for nodes that have been transitively pinned through the FFI functions
+    RootsWorkBuffer tpin_buf = (closure->report_tpinned_nodes_func)((void**)0, 0, 0, closure->data, true);
+    size_t tpin_len = 0;
 
     // globally rooted
     trace_full_globally_rooted(closure, &buf, &len);
@@ -825,6 +828,24 @@ JL_DLLEXPORT void jl_gc_scan_vm_specific_roots(RootsWorkClosure* closure)
         }
     }
 
+    // Trace pinned and transitively pinned objects
+    arraylist_t all_pinned_objects;
+    arraylist_t all_tpinned_objects;
+    arraylist_new(&all_pinned_objects, 0);
+    arraylist_new(&all_tpinned_objects, 0);
+    jl_dump_all_pinned_objects(&all_pinned_objects);
+    jl_dump_all_tpinned_objects(&all_tpinned_objects);
+    for (size_t i = 0; i < all_pinned_objects.len; i++) {
+        void *obj = all_pinned_objects.items[i];
+        add_node_to_roots_buffer(closure, &buf, &len, obj);
+    }
+    for (size_t i = 0; i < all_tpinned_objects.len; i++) {
+        void *obj = all_tpinned_objects.items[i];
+        add_node_to_roots_buffer(closure, &tpin_buf, &tpin_len, obj);
+    }
+    arraylist_free(&all_pinned_objects);
+    arraylist_free(&all_tpinned_objects);
+
     // // add module
     // add_node_to_roots_buffer(closure, &buf, &len, jl_main_module);
 
@@ -850,6 +871,8 @@ JL_DLLEXPORT void jl_gc_scan_vm_specific_roots(RootsWorkClosure* closure)
 
     // Push the result of the work.
     (closure->report_nodes_func)(buf.ptr, len, buf.cap, closure->data, false);
+    // Push the result of the transitively pinned work.
+    (closure->report_tpinned_nodes_func)(tpin_buf.ptr, tpin_len, tpin_buf.cap, closure->data, false);
 }
 
 JL_DLLEXPORT void jl_gc_scan_julia_exc_obj(void* obj_raw, void* closure, ProcessSlotFn process_slot) {

src/julia.h

@@ -1395,6 +1395,8 @@ typedef bool (*check_alive_fn_type)(void *);
 JL_DLLEXPORT void jl_set_check_alive_type(check_alive_fn_type fn);
 JL_DLLEXPORT void jl_log_pinning_event(void *pinned_object, const char *filename, int lineno);
 JL_DLLEXPORT void jl_print_pinning_log(void);
+JL_DLLEXPORT void jl_dump_all_pinned_objects(arraylist_t *objects);
+JL_DLLEXPORT void jl_dump_all_tpinned_objects(arraylist_t *objects);
 
 #define ENABLE_PINNING_LOGGING
 #ifdef ENABLE_PINNING_LOGGING