Fix callback lifetime management using PyCapsule-based cleanup

src/bun_compat.js

@@ -94,7 +94,7 @@ if (!("Deno" in globalThis) && "Bun" in globalThis) {
       }
 
       static value(ptr) {
-        return ptr;
+        return BigInt(ptr);
       }
     };
 

src/python.ts

@@ -3,13 +3,102 @@
 import { py } from "./ffi.ts";
 import { cstr, SliceItemRegExp } from "./util.ts";
 
-const refregistry = new FinalizationRegistry(py.Py_DecRef);
+/**
+ * Track which handles have been registered with refregistry to prevent duplicates.
+ * We track by handle value (bigint) not PyObject instance because multiple
+ * PyObject wrappers can exist for the same underlying Python object.
+ */
+const registeredHandles = new Set<bigint>();
+
+/**
+ * FinalizationRegistry for normal Python objects (NOT callbacks).
+ * When a PyObject wrapper is GC'd, this calls Py_DecRef to release the Python reference.
+ */
+const refregistry = new FinalizationRegistry((handle: Deno.PointerValue) => {
+  py.Py_DecRef(handle);
+});
+
+/**
+ * Global registry to keep Callbacks alive while Python holds references to them.
+ *
+ * Callbacks use a different lifetime management strategy than normal PyObjects:
+ * - They are NOT registered with refregistry
+ * - They are cleaned up via PyCapsule destructor when Python releases them
+ * - This registry prevents JS GC from collecting the Callback while Python needs it
+ */
+const callbackRegistry = new Map<number | bigint, Callback>();
 
-// FinalizationRegistry for auto-created callbacks
-// Closes the callback when the PyObject holding it is GC'd
+/**
+ * Keep track of handle buffers used by PyCapsule to prevent them from being GC'd.
+ * Each callback has a buffer that stores its handle value, which the capsule destructor
+ * uses to look up the callback when Python frees it.
+ */
+const handleBuffers = new Map<number | bigint, BigUint64Array>();
+
+/**
+ * Persistent global buffer for capsule name - MUST NOT be GC'd.
+ * PyCapsule API requires the name pointer to remain valid for the capsule's lifetime.
+ */
+const CAPSULE_NAME = new TextEncoder().encode("deno_python_callback\0");
+// @ts-ignore: globalThis augmentation
+globalThis.__deno_python_capsule_name = CAPSULE_NAME;
+
+/**
+ * PyCapsule destructor - called by Python when a PyCFunction's refcount reaches zero.
+ *
+ * This is the KEY to callback lifetime management:
+ * 1. When a callback is created, we store it in callbackRegistry with its handle
+ * 2. We create a PyCapsule containing the callback's handle
+ * 3. The PyCapsule is used as the m_self parameter of PyCFunction_NewEx
+ * 4. When Python frees the PyCFunction, it decrefs the capsule
+ * 5. When the capsule's refcount reaches 0, Python calls this destructor
+ * 6. We retrieve the handle, look up the callback, and destroy it
+ *
+ * This ensures callbacks are only freed when Python is truly done with them.
+ */
+const capsuleDestructor = new Deno.UnsafeCallback(
+  {
+    parameters: ["pointer"],
+    result: "void",
+  },
+  (capsulePtr: Deno.PointerValue) => {
+    if (capsulePtr === null) return;
+
+    try {
+      const dataPtr = py.PyCapsule_GetPointer(capsulePtr, CAPSULE_NAME);
+      if (dataPtr === null) return;
+
+      const view = new Deno.UnsafePointerView(dataPtr);
+      const handleValue = view.getBigUint64();
+
+      const callback = callbackRegistry.get(handleValue);
+      if (callback) {
+        callbackRegistry.delete(handleValue);
+        handleBuffers.delete(handleValue);
+        callback.destroy();
+      }
+    } catch (_e) {
+      // Silently ignore errors during cleanup
+    }
+  },
+);
+// @ts-ignore: globalThis augmentation
+globalThis.__deno_python_capsule_destructor = capsuleDestructor;
+
+/**
+ * FinalizationRegistry for auto-created callbacks (created from raw JS functions).
+ *
+ * Auto-created callbacks have two possible cleanup paths:
+ * 1. If passed to Python: cleaned up by capsule destructor (skip here)
+ * 2. If never passed to Python: cleaned up here when PyObject wrapper is GC'd
+ *
+ * We check callbackRegistry to determine which path applies.
+ */
 const callbackCleanupRegistry = new FinalizationRegistry(
-  (callback: Callback) => {
-    callback.destroy();
+  (data: { callback: Callback; handle: bigint }) => {
+    if (!callbackRegistry.has(data.handle)) {
+      data.callback.destroy();
+    }
   },
 );
 
@@ -206,8 +295,18 @@ export class Callback {
     );
   }
 
+  /**
+   * Destroys the callback by closing the UnsafeCallback resource.
+   * This is idempotent and safe to call multiple times.
+   */
   destroy() {
-    this.unsafe.close();
+    if (this.unsafe && this.unsafe.pointer) {
+      try {
+        this.unsafe.close();
+      } catch (_e) {
+        // Ignore double-close errors (BadResource)
+      }
+    }
   }
 }
 
@@ -235,10 +334,10 @@ export class Callback {
  */
 export class PyObject {
   /**
-   * A Python callabale object as Uint8Array
-   * This is used with `PyCFunction_NewEx` in order to extend its liftime and not allow v8 to release it before its actually used
+   * Reference to the Callback object if this PyObject wraps a callback.
+   * Stored so we can register it with callbackRegistry when .owned is called.
    */
-  #pyMethodDef?: Uint8Array;
+  #callback?: Callback;
   constructor(public handle: Deno.PointerValue) {}
 
   /**
@@ -253,10 +352,33 @@ export class PyObject {
 
   /**
    * Increases ref count of the object and returns it.
+   * This is idempotent - calling it multiple times on the same instance
+   * will only increment the refcount and register once.
+   *
+   * IMPORTANT: Callbacks are handled differently:
+   * - They do NOT get registered with refregistry
+   * - They do NOT call Py_IncRef (the initial ref from PyCFunction_NewEx is sufficient)
+   * - They are cleaned up via the capsule destructor when Python releases them
+   *
+   * This prevents the refcount from being artificially inflated, allowing Python
+   * to properly free callbacks when they're no longer needed.
    */
   get owned(): PyObject {
-    py.Py_IncRef(this.handle);
-    refregistry.register(this, this.handle);
+    const handleValue = Deno.UnsafePointer.value(this.handle);
+
+    if (!registeredHandles.has(handleValue)) {
+      // Check if this is a callback (has #callback reference)
+      if (this.#callback) {
+        // Callback - register with callbackRegistry now that it's being passed to Python
+        callbackRegistry.set(handleValue, this.#callback);
+      } else {
+        // Normal PyObject - use refregistry
+        py.Py_IncRef(this.handle);
+        refregistry.register(this, this.handle);
+      }
+      registeredHandles.add(handleValue);
+    }
+
     return this;
   }
 
@@ -545,18 +667,42 @@ export class PyObject {
             ),
             LE,
           );
+
+          // Create a temporary buffer to store callback ID (will be filled after fn is created)
+          const handleBuffer = new BigUint64Array([0n]);
+
+          // Create a capsule with destructor to track when Python frees the callback
+          const capsule = py.PyCapsule_New(
+            handleBuffer,
+            CAPSULE_NAME,
+            capsuleDestructor.pointer,
+          );
+
+          // Use the capsule as m_self parameter so it gets freed with the function
+          // This is KEY - when PyCFunction is freed, Python will decref m_self (capsule),
+          // triggering our destructor
           const fn = py.PyCFunction_NewEx(
             v._pyMethodDef,
-            PyObject.from(null).handle,
+            capsule,
             null,
           );
 
-          // NOTE: we need to extend `pyMethodDef` lifetime
-          // Otherwise V8 can release it before the callback is called
-          // Is this still needed (after the change of pinning fields to the callabck) ? might be
+          // Now fill in the handle buffer with the actual function pointer
+          const handleValue = Deno.UnsafePointer.value(fn);
+          handleBuffer[0] = handleValue;
+
+          // Store handle buffer for capsule destructor lookup
+          // The callback will be added to callbackRegistry when .owned is called
+          handleBuffers.set(handleValue, handleBuffer);
+
+          // Decref capsule so only PyCFunction holds it
+          // When PyCFunction is freed, capsule refcount goes to 0, triggering our destructor
+          py.Py_DecRef(capsule);
+
           const pyObject = new PyObject(fn);
-          pyObject.#pyMethodDef = methodDef;
-          // Note: explicit Callback objects are user-managed, not auto-cleaned
+          pyObject.#callback = v; // Store callback reference for registration in .owned
+
+          // Do NOT register with refregistry - callbacks use capsule-based cleanup
           return pyObject;
         } else if (v instanceof PyObject) {
           return v;
@@ -614,8 +760,13 @@ export class PyObject {
         if (typeof v === "function") {
           const callback = new Callback(v);
           const pyObject = PyObject.from(callback);
-          // Register cleanup to close callback when PyObject is GC'd
-          callbackCleanupRegistry.register(pyObject, callback);
+
+          // Register with cleanup registry for auto-created callbacks
+          const handleValue = Deno.UnsafePointer.value(pyObject.handle);
+          callbackCleanupRegistry.register(pyObject, {
+            callback,
+            handle: handleValue,
+          });
           return pyObject;
         }
       }

src/symbols.ts

@@ -259,4 +259,19 @@ export const SYMBOLS = {
     parameters: ["pointer"],
     result: "pointer",
   },
+
+  PyCapsule_New: {
+    parameters: ["buffer", "buffer", "pointer"], // pointer, name, destructor
+    result: "pointer",
+  },
+
+  PyCapsule_GetPointer: {
+    parameters: ["pointer", "buffer"], // capsule, name
+    result: "pointer",
+  },
+
+  PyCapsule_SetPointer: {
+    parameters: ["pointer", "pointer"], // capsule, pointer
+    result: "i32",
+  },
 } as const;

test/test_with_gc.ts

@@ -12,9 +12,6 @@ def call_the_callback(cb):
   );
 
   assertEquals(pyModule.call_the_callback(() => 4).valueOf(), 5);
-
-  // @ts-ignore:requires: --v8-flags=--expose-gc
-  gc(); // if this is commented out, the test will fail beacuse the callback was not freed
 });
 
 Deno.test("callbacks are not gc'd while still needed by python", () => {
@@ -47,10 +44,8 @@ def call_stored_callback():
   assertEquals(pyModule.store_and_call_callback(callbackObj).valueOf(), 10);
   assertEquals(callCount, 1);
 
-  for (let i = 0; i < 10; i++) {
-    // @ts-ignore:requires: --v8-flags=--expose-gc
-    gc();
-  }
+  // @ts-ignore:requires: --v8-flags=--expose-gc
+  gc();
 
   // If the callback was incorrectly GC'd, this should segfault
   // But it should work because Python holds a reference
@@ -78,33 +73,35 @@ def call_stored_callback():
   if stored_callback is None:
     return -1
   return stored_callback()
+
+def clear_callback():
+  global stored_callback
+  stored_callback = None
   `,
     "test_gc_module",
   );
 
-  let callCount = 0;
-  const callback = () => {
-    callCount++;
-    return callCount * 10;
-  };
-
   // Store the callback in Python and call it
-  assertEquals(pyModule.store_and_call_callback(callback).valueOf(), 10);
-  assertEquals(callCount, 1);
-
-  for (let i = 0; i < 10; i++) {
-    // @ts-ignore:requires: --v8-flags=--expose-gc
-    gc();
+  {
+    assertEquals(pyModule.store_and_call_callback(() => 4).valueOf(), 4);
   }
 
+  // @ts-ignore:requires: --v8-flags=--expose-gc
+  gc();
+
   // If the callback was incorrectly GC'd, this should segfault
   // But it should work because Python holds a reference
-  assertEquals(pyModule.call_stored_callback().valueOf(), 20);
-  assertEquals(callCount, 2);
+  assertEquals(pyModule.call_stored_callback().valueOf(), 4);
 
   // Call it again to be sure
-  assertEquals(pyModule.call_stored_callback().valueOf(), 30);
-  assertEquals(callCount, 3);
+  assertEquals(pyModule.call_stored_callback().valueOf(), 4);
+
+  // Clean up Python's reference to allow callback to be freed
+  pyModule.clear_callback();
+
+  // Force GC to trigger capsule destructor
+  // @ts-ignore:requires: --v8-flags=--expose-gc
+  gc();
 });
 
 Deno.test("auto-created callbacks are cleaned up after gc", () => {
@@ -116,10 +113,6 @@ Deno.test("auto-created callbacks are cleaned up after gc", () => {
   // Explicitly null the reference
   _f = null;
 
-  // Now f is null, trigger GC to clean it up
-  // Run many GC cycles with delays to ensure finalizers execute
-  for (let i = 0; i < 10; i++) {
-    // @ts-ignore:requires: --v8-flags=--expose-gc
-    gc();
-  }
+  // @ts-ignore:requires: --v8-flags=--expose-gc
+  gc();
 });