junrushao commented on code in PR #593: URL: https://github.com/apache/tvm-ffi/pull/593#discussion_r3576609743
########## python/tvm_ffi/cython/tvm_ffi_python_object.h: ########## @@ -0,0 +1,1012 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +/* + * \file tvm_ffi_python_object.h + * \brief PyObject-tying state machine: binds one Python wrapper to one C++ FFI object + * ("chandle") for the object's lifetime so identity is stable (``a.x is a.x``, + * stable ``id()`` across drop+refetch, ``f(x) is x`` for FFI returns). + * + * Split out of tvm_ffi_python_helpers.h. The design overview is the banner comment below. + */ +#ifndef TVM_FFI_PYTHON_OBJECT_H_ +#define TVM_FFI_PYTHON_OBJECT_H_ + +#include <Python.h> +#include <tvm/ffi/c_api.h> +#include <tvm/ffi/memory.h> + +// Define here to avoid dependencies on non-c headers for now +#ifndef TVM_FFI_INLINE +#if defined(_MSC_VER) +#define TVM_FFI_INLINE [[msvc::forceinline]] inline +#else +#define TVM_FFI_INLINE [[gnu::always_inline]] inline +#endif +#endif + +// Managed-dict (`__slots__ = ("__dict__",)` without an explicit dictoffset) +// is a CPython 3.11+ feature. On 3.9/3.10 such types instead use a regular +// ``tp_dictoffset != 0``, which the inactive-eligibility check catches anyway, +// so defining the flag as 0 here yields the correct (no-op) behavior. +#ifndef Py_TPFLAGS_MANAGED_DICT +#define Py_TPFLAGS_MANAGED_DICT 0 +#endif + +#include <atomic> +#include <cassert> +#include <cstring> +#include <utility> + +// ``_Interlocked*`` intrinsics for the MSVC arm of the spin-lock leaves below. <intrin.h>, not +// <windows.h>, to keep min/max etc. macros out of the Cython TU. +#if defined(_MSC_VER) && defined(Py_GIL_DISABLED) +#include <intrin.h> +#endif + +//================================================================================ +// PyObject-tying state machine. +// +// Ties one Python wrapper to one C++ chandle so that +// - ``a.x is a.x`` while the wrapper is live; +// - ``id(a.x)`` is stable across drop+refetch (when other C++ holders keep +// the chandle alive); +// - ``f(x) is x`` whenever an FFI function returns a chandle that already +// has a canonical wrapper. +// +// Layout +// ------ +// Every Object allocated through the registered Python allocator +// (`TVMFFIPyAllocate`) is preceded by a 16-byte ``PyCustomAllocHeader``: +// +// malloc start +// +-------------------+--------------------------+--------+ +// | tagged_pyobj | TVMFFIObjectAllocHeader | T | +// | (offset 0..8) | delete_space (8..16) | | +// +-------------------+--------------------------+--------+ +// ^ ptr = malloc + 16 +// +// ``tagged_pyobj`` is a tagged pointer to the canonical Python wrapper. The +// wrapper is >= 16-aligned, so the low 4 bits are free; two encode the state +// (see below) without growing the header past its fixed 16 bytes. +// +// States +// ------ +// Bit 0 (Inactive) and bit 1 (InTransit) tag ``tagged_pyobj`` into four states: +// Detached: ``tagged_pyobj == NULL`` -- no wrapper bound to this chandle. +// Active: ``ptr, bits == 00`` -- the live canonical wrapper. +// Inactive: ``ptr | Inactive`` -- dead, untracked allocation cached for +// address-stable revival (settled). +// InTransit: ``ptr | Inactive | InTransit`` -- a transition on this binding is in +// flight (Inactive stays set, so ``TVMFFIPyTagIsInactive`` matches too). +// +// Invariants +// ---------- +// I1. When a PyObject goes out of scope (no Python var refers to it), its +// +1 on chandle is always released (in ``__dealloc__`` -> +// ``TVMFFIPyTpDealloc``). +// I2. When a chandle is destroyed, its cached allocation (if any) is +// reclaimed. +// I3'. ``wrapper.chandle`` is only ever a real C++ object pointer or NULL, +// never a sentinel. A non-NULL chandle owns +1, except inside the +// wrapper's own dealloc window (where it is kept only as a header locator). +// I4. Every ``PyObject*`` the Cython side passes to a helper here is a live +// wrapper (tag bits 0); only this header sets or clears the tag bits. +// I5. InTransit is the dealloc handshake's baton and nothing else: it overlays a +// non-live binding (Inactive(W) or Detached(NULL)) while the two teardown sides +// settle, never the live Active wrapper. Any reader that sees it -- a peer +// settler, or make_ret's classify -- waits the transition out. +// +// The dealloc handshake +// --------------------- +// One allocation can be torn down from two directions, and the handshake stops them +// from racing into a double free or a leak: +// * from Python -- the wrapper's refcount hits 0, so ``tp_dealloc`` -> ``tp_free`` run; +// * from C++ -- the chandle's weak count hits 0, so its Weak deleter fires +// ``TVMFFIPyDeleteSpace``. +// ``tp_dealloc`` cannot read the chandle refcount to decide which side will be last (an +// FFI ``DecRef`` may race it from another thread), so instead of deciding up front it +// pre-tags ``Inactive | InTransit`` and ``DecRef``s unconditionally; the InTransit bit is a +// baton the FIRST settler clears (deferring), so the SECOND finds it clear and does the free. +// +// Flow 1 -- wrapper dies, chandle outlives it (cache the allocation): +// tp_dealloc : Active -> Inactive +// InTransit 0 -> 1 +// DecRef (chandle still has refs, so no deleter fires) +// tp_free : InTransit 1 -> 0, keep ``self`` cached Inactive. +// ... +// (later in C++, the chandle dies, delete_space fires) +// delete_space : InTransit == 0, so reclaim the cached wrapper and free the block. +// +// Flow 2 -- wrapper held the last ref (free the allocation now): +// tp_dealloc : Active -> Inactive +// InTransit 0 -> 1 +// DecRef (last ref dropped, so reentrantly fires delete_space) +// delete_space : InTransit 1 -> 0, defer the free back to tp_free. +// tp_free : InTransit == 0, so free the C++ block here. +// +// Where transitions happen +// ------------------------ +// ``TVMFFIPyMakeRetObject`` (this header), behind ``make_ret_object`` +// (object.pxi) -- owns the whole return-object transition in one frame: +// Detached/Active/Inactive -> Active : fresh / cached / revived-in-place. +// +// ``TVMFFIPyTpDealloc`` (CObject.__dealloc__) -- runs when the wrapper's +// refcount hits 0, before the free: +// Active -> Inactive : eligible; tag Inactive | InTransit, DecRef (the +// handshake; ``tp_free`` / ``TVMFFIPyDeleteSpace`` +// settle it). +// Active -> Detached : type not eligible; detach first, then DecRef. +// +// ``TVMFFIPyArgSetterObjectRValueRef_`` (function.pxi), +// ``__move_handle_from__`` (object.pxi): +// Active -> Detached : detach the binding before a move nulls the +// source chandle. +// +// ``TVMFFIPyDeleteSpace`` (Weak deleter) -- the chandle's weak count hit 0: +// Inactive|InTransit : in-flight dealloc; defer both frees to ``tp_free``. +// Inactive (settled) : reclaim the cached wrapper and free the C++ block. +// +// Slot install +// ------------ +// Two slot families, each unmissable over a different scope: +// * ``tp_dealloc`` (correctness, I1): unmissable over the WHOLE CObject hierarchy -- defined +// once on CObject and inherited by every subtype, so nothing to install per type. +// * ``tp_alloc`` / ``tp_free`` (cache-&-revive optimization): unmissable over the REGISTERED +// types -- ``_update_registry`` (object.pxi) is the sole choke point all registration funnels +// through, and it installs there. An unregistered subtype just fails eligibility and +// genuine-frees (loses stable-id-across-drop, not correctness). +// ``tp_dealloc`` works with either pairing: with the custom ``tp_alloc`` / ``tp_free`` it caches +// (eligible), and with the generic ones it detaches + genuine-frees -- the eligibility gate keys +// on the same ``tp_free``, so the two can never disagree. +// +// Shutdown guard +// -------------- +// ``TVMFFIPyMarkPythonFinalizing`` is wired to atexit from Cython module +// init. After it fires, inactive cached allocations on still-live chandles are +// intentionally leaked (process exiting; OS reclaims) rather than reaching +// for ``PyGILState_Ensure`` on a teardown interpreter. +// +// Free-threaded builds (``Py_GIL_DISABLED``) +// ------------------------------------------ +// Without the GIL the bare ``tagged_pyobj`` reads/writes above race -- the Active-hit +// read is a use-after-free (``make_ret`` reads the wrapper, a concurrent dealloc frees +// it before the IncRef). The tie stays enabled; three FT-only mechanisms close the gap, +// all behind ``#ifdef Py_GIL_DISABLED`` so the GIL build is byte-for-byte unchanged: +// * The word is its own spin-lock (a Locked tag bit, CAS-acquired via the portable +// pointer-atomic leaves -- ``__atomic_*`` on GCC/Clang, ``_Interlocked*`` on MSVC), +// so every transition serializes its word edits. Details in the word-access leaves. +// * The Active hit uses ``PyUnstable_TryIncRef`` (inc-if-nonzero), not ``Py_INCREF``, +// so it fails on a wrapper a concurrent dealloc is collecting -- closing the UAF. +//================================================================================ + +/*! + * \brief Python-side derived header. ``base.delete_space`` sits at + * ``ptr - sizeof(TVMFFIObjectAllocHeader)`` so the generic C++ + * deleter (which knows nothing about Python) can find it. + */ +struct PyCustomAllocHeader { + PyObject* tagged_pyobj; Review Comment: ```cpp /*! * \brief Tagged-pointer encoding for a chandle's canonical Python wrapper. * * ``PyCustomAllocHeader::tagged_pyobj`` stores both the canonical wrapper * address and its lifecycle state. Python wrapper allocations are at least * 16-byte aligned, leaving their low four address bits available as tags: * * \code * Bit 0: Inactive * The wrapper has reached refcount zero. Its untracked allocation is retained * so it can later be revived at the same address. * * Bit 1: InTransit * Python and native teardown are settling ownership of the allocation. * * Bit 2: Locked * Free-threaded builds only. Acts as a per-header spin lock protecting state * transitions. It is orthogonal to the lifecycle state. * \endcode * * Ignoring the temporary Locked bit, the stable states are: * * \code * nullptr Detached * wrapper Active * wrapper | Inactive Inactive * wrapper | Inactive | InTransit InTransit * \endcode * * Detached means no canonical wrapper is associated with the chandle. Active * identifies the live canonical wrapper. Inactive identifies dead, untracked * wrapper storage that can be revived in place. InTransit is an intermediate * teardown state used by ``tp_free`` and ``delete_space`` to avoid double-free * and leak races: the first side to settle clears InTransit and defers cleanup; * the second side observes the settled state and performs the required cleanup. * * The helpers below operate on this encoding: * * - ``TVMFFIPyTagIsInactive`` tests bit 0 and therefore matches both Inactive * and InTransit states. * - ``TVMFFIPyTagInTransit`` tests whether the teardown handshake is active. * - ``TVMFFIPyRemoveTag`` removes all tag bits and recovers the wrapper address. * - ``TVMFFIPyTagClearInTransit`` clears only bit 1, normally transitioning * Inactive|InTransit to settled Inactive. * * These helpers only inspect or transform the encoded word. They do not change * reference counts, allocate memory, or free either the wrapper or chandle. */ ``` -- 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]
