Close callbacks when PyObject is GCed by JS

Author: sigmaSd

src/python.ts

@@ -6,14 +6,10 @@ import { cstr, SliceItemRegExp } from "./util.ts";
 const refregistry = new FinalizationRegistry(py.Py_DecRef);
 
 // FinalizationRegistry for auto-created callbacks
-const callbackRegistry = new FinalizationRegistry(
-  (
-    unsafe: Deno.UnsafeCallback<{
-      parameters: ["pointer", "pointer", "pointer"];
-      result: "pointer";
-    }>,
-  ) => {
-    unsafe.close();
+// Closes the callback when the PyObject holding it is GC'd (meaning Python released it)
+const callbackCleanupRegistry = new FinalizationRegistry(
+  (callback: Callback) => {
+    callback.destroy();
   },
 );
 
@@ -213,21 +209,6 @@ export class Callback {
   destroy() {
     this.unsafe.close();
   }
-
-  /**
-   * Marks this Callback instance as "owned" by registering it with the
-   * FinalizationRegistry for automatic cleanup when the Callback is garbage collected.
-   *
-   * This allows the underlying Deno.UnsafeCallback to be closed automatically
-   * when the Callback object is no longer referenced, helping to prevent resource leaks.
-   * This is useful if you don't want to call destroy manually.
-   *
-   * @returns {Callback} The current Callback instance.
-   */
-  get owned(): Callback {
-    callbackRegistry.register(this, this.unsafe);
-    return this;
-  }
 }
 
 /**
@@ -575,6 +556,7 @@ export class PyObject {
           // Is this still needed (after the change of pinning fields to the callabck) ? might be
           const pyObject = new PyObject(fn);
           pyObject.#pyMethodDef = methodDef;
+          // Note: explicit Callback objects are user-managed, not auto-cleaned
           return pyObject;
         } else if (v instanceof PyObject) {
           return v;
@@ -630,7 +612,11 @@ export class PyObject {
         }
 
         if (typeof v === "function") {
-          return PyObject.from(new Callback(v).owned);
+          const callback = new Callback(v);
+          const pyObject = PyObject.from(callback);
+          // Register cleanup to close callback when PyObject is GC'd
+          callbackCleanupRegistry.register(pyObject, callback);
+          return pyObject;
         }
       }
 

test/test_with_gc.ts

@@ -1,6 +1,22 @@
-import python, { Callback } from "../mod.ts";
+import python, { Callback, PyObject } from "../mod.ts";
 import { assertEquals } from "./asserts.ts";
 
+Deno.test("js fns are automaticlly converted to callbacks", () => {
+  const pyModule = python.runModule(
+    `
+def call_the_callback(cb):
+  result = cb()
+  return result + 1
+  `,
+    "test_module",
+  );
+
+  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", () => {
   const pyModule = python.runModule(
     `
@@ -46,3 +62,64 @@ def call_stored_callback():
   assertEquals(callCount, 3);
   callbackObj.destroy();
 });
+
+Deno.test("callbacks are not gc'd while still needed by python (function version)", () => {
+  const pyModule = python.runModule(
+    `
+stored_callback = None
+
+def store_and_call_callback(cb):
+  global stored_callback
+  stored_callback = cb
+  return stored_callback()
+
+def call_stored_callback():
+  global stored_callback
+  if stored_callback is None:
+    return -1
+  return stored_callback()
+  `,
+    "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();
+  }
+
+  // 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);
+
+  // Call it again to be sure
+  assertEquals(pyModule.call_stored_callback().valueOf(), 30);
+  assertEquals(callCount, 3);
+});
+
+Deno.test("auto-created callbacks are cleaned up after gc", () => {
+  // Create callback and explicitly null it out to help GC
+  // @ts-ignore PyObject can be created from fns its just the types are not exposed
+  // deno-lint-ignore no-explicit-any
+  let _f: any = PyObject.from(() => 5);
+
+  // 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();
+  }
+});