gemini-code-assist[bot] commented on code in PR #593:
URL: https://github.com/apache/tvm-ffi/pull/593#discussion_r3265613443
##########
python/tvm_ffi/cython/object.pxi:
##########
@@ -88,6 +88,57 @@ class ObjectRValueRef:
self.obj = obj
+cdef inline void _detach_chandle_binding(CObject src) noexcept:
+ """Clear the prepended PyCustomizeAllocHeader binding when ``src`` is the
+ canonical wrapper for its chandle.
+
+ No-op when ``src.chandle == NULL`` or when the header points at someone
+ else (e.g. a no-cache wrapper that never installed itself).
+
+ Used by:
+ - ``__dealloc__`` — clear before ``DecRef`` so the C++ deleter's Weak
+ branch skips ``free_cb`` (state-(a)↔(b) invariant).
+ - ``__move_handle_from__`` — clear the source so the binding can be
+ transferred to the destination.
+ - ``TVMFFIPyArgSetterObjectRValueRef_`` — clear before the C++ side
+ nulls ``src.chandle``; otherwise the header keeps a stale pointer
+ at a still-live wrapper and the eventual deleter calls
+ ``PyObject_GC_Del`` on it (use-after-free).
+ """
+ cdef PyCustomizeAllocHeader* h
+ if src.chandle == NULL:
+ return
+ h = TVMFFIPyHeader(src.chandle)
+ if h.py_obj == <void*><PyObject*>src:
+ h.py_obj = NULL
+ h.cached_mem = NULL
+ h.free_cb = NULL
+
+
+cdef inline void _install_chandle_binding(CObject obj, void* chandle) except *:
+ """Set ``obj.chandle = chandle`` and bind the wrapper to the prepended
+ PyCustomizeAllocHeader (state (a) -> state (b)).
+
+ The caller must currently own a +1 strong ref on ``chandle``.
+
+ The header may already carry a different live wrapper for the same
+ chandle when, e.g., a Python field-getter handed us a fresh +1 but
+ another thread/path has just cached its own wrapper. That case is
+ handled by ``make_ret_object``'s fast path *before* this function is
+ called, so by the time we reach here we expect ``h.py_obj == NULL``.
+ Defensive: if not NULL, leave the header alone (don't clobber).
+ """
+ cdef PyCustomizeAllocHeader* h
+ obj.chandle = chandle
+ if chandle == NULL:
+ return
+ h = TVMFFIPyHeader(chandle)
+ if h.py_obj == NULL:
+ h.py_obj = <void*><PyObject*>obj
+ h.cached_mem = <void*><PyObject*>obj
+ h.free_cb = TVMFFIPyFreeCallback
Review Comment:
In `_install_chandle_binding`, if `h.py_obj` is `NULL` but `h.cached_mem` is
non-`NULL` (meaning the object is in state-c), the current code overwrites
`h.cached_mem` with the new wrapper's address. This causes a memory leak of the
previous "phantom" wrapper, as its reference count remains at 1 but it is no
longer reachable via the header for cleanup.
This scenario can occur during unpickling or manual handle movement where
`_install_chandle_binding` is called directly, bypassing the revive logic in
`make_ret_object`. To fix this, we should check for an existing cached wrapper
and release it if it's not the same as the one being installed.
```
if h.py_obj == NULL:
if h.cached_mem != NULL and h.cached_mem != <void*><PyObject*>obj:
Py_DecRef(<PyObject*>h.cached_mem)
h.py_obj = <void*><PyObject*>obj
h.cached_mem = <void*><PyObject*>obj
h.free_cb = TVMFFIPyFreeCallback
```
##########
python/tvm_ffi/cython/tvm_ffi_python_helpers.h:
##########
@@ -71,6 +105,35 @@
// prefixed with TVMFFIPy so they can be easily invoked through Cython.
///--------------------------------------------------------------------------------
+//------------------------------------------------------------------------------------
+// Python-object tying helpers.
+//
+// See include/tvm/ffi/py_object_tying.h for the design of the prepended
+// PyCustomizeAllocHeader. These helpers expose:
+// - TVMFFIPyHeader: pointer arithmetic to recover the header from a chandle.
+// - TVMFFIPyFreeCallback: default free_cb for installed wrappers, calls
+// PyObject_GC_Del. Stored on the header at install time so the C++ deleter
+// can free wrapper memory without depending on Python.h.
+//
+// PoC scope (state a <-> b only): wrapper destruction always clears the
+// header fields before DecRef'ing chandle, so the C++ deleter never observes
+// cached_mem != NULL. The free_cb path exists for future state-(c) work.
+//------------------------------------------------------------------------------------
+// Global alias so the Cython-generated C++ can reach the type without a
+// namespace qualifier (Cython does not emit `using namespace` directives).
+using PyCustomizeAllocHeader = ::tvm::ffi::PyCustomizeAllocHeader;
+
+TVM_FFI_INLINE PyCustomizeAllocHeader* TVMFFIPyHeader(void* chandle) {
+ return ::tvm::ffi::GetPyHeader(chandle);
+}
+
+TVM_FFI_INLINE void TVMFFIPyFreeCallback(void* mem) {
+ // Wrapper memory was allocated via the standard CPython object allocator.
+ // PyObject_GC_Del is the inverse of PyObject_GC_New / PyType_GenericAlloc.
+ // Caller must hold the GIL.
+ PyObject_GC_Del(mem);
+}
Review Comment:
The current implementation of `TVMFFIPyFreeCallback` is unsafe because it
invokes the Python C API (`PyObject_GC_Del`) without ensuring the Global
Interpreter Lock (GIL) is held. Since the C++ deleter can be triggered from any
thread (e.g., a background C++ thread dropping the last reference), this will
lead to crashes or undefined behavior.
Additionally, using `PyObject_GC_Del` directly on a resurrected wrapper
(state-c) is problematic because it bypasses the Python object's deallocator
(`tp_dealloc`). This will leak any Python-side attributes (like those stored in
`__dict__`) or other resources held by the wrapper.
It is recommended to acquire the GIL and use `Py_DecRef` instead, which will
correctly trigger the full deallocation sequence for the phantom reference.
```c
TVM_FFI_INLINE void TVMFFIPyFreeCallback(void* mem) {
PyGILState_STATE gstate = PyGILState_Ensure();
Py_DecRef(static_cast<PyObject*>(mem));
PyGILState_Release(gstate);
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
