fix: correct rare-case bugs from adversarial audit across deadlocks, caching, serialization, memory safety, and Python frontend - #290
Open
XuehaiPan wants to merge 5 commits into
Conversation
XuehaiPan
added
bug
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #290 +/- ##
=========================================
Coverage 100.00% 100.00%
=========================================
Files 15 15
Lines 1585 1617 +32
Branches 212 218 +6
=========================================
+ Hits 1585 1617 +32 Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. |
XuehaiPan
force-pushed
the
fix/correctness-audit
branch
10 times, most recently
from
July 30, 2026 07:36
55bb6ae to
948330e
Compare
No Python may run while `sm_mutex` is held. A concurrent flatten holds the GIL and waits on the read lock, so anything under the write lock that executes bytecode can hand the GIL to that thread and never get it back, wedging the whole interpreter. Several sites broke that rule. - `Register` and `Unregister` ran the namedtuple / PyStructSequence classification and `PyErr_WarnEx` under the write lock. Both run Python and can release the GIL. Classify and warn before taking the lock; emitting the override warning before any insert also keeps registration atomic under warnings-as-errors. - Their error paths still formatted messages with `PyRepr`, which runs the type's `__repr__` as bytecode. `RegisterImpl` and `UnregisterImpl` now only report a `RegistryStatus`, and the caller formats and raises after unlocking. An ordinary `enum.Enum` is enough to hit this: its metaclass `__repr__` is plain Python. - `Unregister` and `Clear` dropped the registry's own references while still locked. Dropping the last reference to a flatten function runs `__del__` and weakref callbacks, which can re-enter the registry. The registrations are now detached under the lock and destroyed after it. - `Init` recorded the interpreter as alive before `atexit.register` succeeded. That call runs Python under the lock, and if it raises the id is left behind with no callback able to remove it. Register the cleanup first, mark the interpreter only on success, and roll back on throw. - `Lookup`, `Register`, `Unregister` and `GetRegistrySize` called `GetSingleton()` under the lock. Under `per_interpreter_gil` that releases the GIL once any subinterpreter has imported the module. Acquire the singletons first. - `FlattenInto` read-locked `sm_dict_order_mutex` and then called `IsDictInsertionOrdered`, which locks it again. `std::shared_mutex` is not recursive, so the nested shared acquisition is undefined behaviour and deadlocks under writer preference. Add `GetDictInsertionOrderedFlags`, which returns both flags from one lock, and make `IsDictInsertionOrdered` delegate to it. - `GetRegistrySize` took `Size()` twice, dropping the lock between the two reads, so a concurrent registration could slip in and trip the consistency check. Split out `SizeImpl` and read both registries under a single lock. Re-entrancy is the other half. `PyTreeIter::Next` releases the GIL, takes the non-recursive `m_mutex`, and only then runs the leaf predicate and any custom `flatten_func`. A callback that advances the same iterator blocks on a mutex its own thread holds, with the GIL released, so nothing can interrupt it. Track the thread inside `Next` and reject re-entry with a `RuntimeError`, mirroring CPython's "generator already executing"; `ToString` and `HashValue` already guard themselves this way. `RegisterImpl` and `UnregisterImpl` no longer need to be templated on `NoneIsLeaf` now that the callers hold both singletons, so they become plain member functions.
…rpreter-safe The namedtuple classification, PyStructSequence classification and struct sequence field-name caches each hand-rolled the same cache, weakref and locking dance. Extract one `WeakKeyCache` class and route all three through it, then fix the subinterpreter and lifetime hazards in one place. Entries are keyed by `(interpreter_id, object address)` rather than the address alone. The address can be shared across interpreters, since a type like `int` is immortal and lives at the same address in each, while a computed reference value belongs to the interpreter that produced it. An address-only key would hand that value to another interpreter, which could use it after the owner freed it on finalization. A per-entry weakref evicts an entry when its key is collected, so a later key that reuses the freed address cannot read a stale value. A per-interpreter `atexit` callback clears that interpreter's entries, covering what the weakref cannot: an immortal key is never collected, and interpreter ids restart from 0 after a `Py_Finalize` and `Py_Initialize` cycle, so a fresh interpreter must not inherit a finalized one's entries. Two orderings matter, and both were wrong. The lookup releases the read lock and re-acquires the GIL in that order before touching the borrowed value, so the GIL is never re-acquired while the lock is held, which would invert the order against the weakref eviction callback. The insert registers the interpreter's cleanup callback and takes ownership of the value *before* publishing the entry, marking the interpreter as registered only once that succeeds and rolling back if it does not. Publishing first left a window where a raising `atexit.register` abandoned an entry whose value was unowned and whose key had no eviction callback, so the next lookup read it after it was freed. `GetCurrentPyInterpreterID` sits on that lookup path, twice per leaf through the `Is*Instance` caches, so it tests the documented failure returns rather than probing the global error state: neither `PyInterpreterState_Get()` nor `PyInterpreterState_GetID()` can set an exception where it is called. `hashing.h` gains `std::equal_to` and `std::not_equal_to` for the two pair key types in use. They are full specializations, not a partial one over `std::pair<T, U>`: a partial specialization of a class template already instantiated elsewhere in the program is ill-formed, and GCC rejects it across every translation unit that has already used the pair while clang accepts it silently. The duplication is therefore deliberate. `std::hash<std::pair<T, U>>` can stay generic because no primary `std::hash<std::pair>` exists to pre-instantiate. The header also gained the `<string>` it had always used through a transitive pybind11 include. Value-returning helpers in `pytypes.h` are marked `[[nodiscard]]`, and the unused `stdutils.h` include is dropped.
XuehaiPan
force-pushed
the
fix/correctness-audit
branch
from
July 30, 2026 11:46
948330e to
984d7c2
Compare
…led state Both areas run through `serialization.cpp` and the struct sequence helpers. `tp_members` lists only the named fields, so indexing it by position mislabelled every slot after the first unnamed one: `os.stat_result` slots 7, 8 and 9 were reported as `st_atime`, `st_mtime` and `st_ctime`. Map by the byte offset each member carries instead, and fill the remaining slots with the unnamed marker. The pure Python `structseq_fields` cannot read those offsets, so it recovers each named field's position from a probe instance holding distinct sentinels and matches them by identity; unnamed slots may sit anywhere, not only at the tail. The marker is exported as `PyStructSequence_UnnamedField` and is now the single test for "this slot is unnamed": the treespec repr renders such a slot as `<unnamed@N>`, and `StructSequenceEntry` compares against the constant rather than guessing from `str.isidentifier()`, so its repr reports the index instead of printing the raw marker as if it were a field name, and `codify` falls back to index access. `PyStructSequenceUnnamedField()` works around the symbol being unexported before 3.11.0a2, where referencing it broke the import. `FromPicklable` trusted its input, so a crafted pickle could build a spec that later read out of bounds or aborted in repr. Validate the kind before the narrowing cast, the arity against the children the traversal provides, node data on childless kinds, dict key count and distinctness, defaultdict shape, deque `maxlen`, namedtuple and struct sequence arity, and the whole traversal rather than only its last node. A `Custom` node must also name a type whose registration is a custom one: the built-in registrations share the same map but carry empty flatten and unflatten callables, which `MakeNode` would have called. Aliasing had to be fixed in both directions. `ToPicklable` copies the node's mutable containers so the state cannot alias the spec, and `FromPicklable` copies the containers it is given so the restored spec cannot alias the caller's state, which mutating that state afterwards would otherwise reorder. Both go through `ListCopy` and `DictCopy`, which use the C API rather than a Python-level `.copy()` and so cost no attribute lookup, bound-method allocation or vectorcall per node. `__reduce__` makes protocols 0 and 1 reconstruct through `cls.__new__(cls)` instead of aborting. `is_namedtuple_instance()` and `is_structseq_instance()` accept `object`, matching what they already did at runtime.
Namespace merging. `compose`, `broadcast_to_common_suffix`, `transform` and `treespec_from_collection` let an empty namespace adopt the other side's namespace. If a custom node resolved globally but resolves to a *different* registration under the adopted namespace, the result claimed the namespace while carrying the wrong registration. Add `FindReregisteredCustomType` and reject those merges; a type registered only globally still resolves by fallback and is allowed. The dict constructors also dropped the namespace, losing insertion-ordered key order, and `MakeFromCollection` ignored the return value of `PyErr_WarnEx`, so an escalated warning surfaced as a confusing `SystemError`. A childless root goes the other way: a leaf or `None` resolves no custom type and no key order, so it has no namespace dependency, and keeping the caller's made otherwise identical treespecs compare unequal. GC. `PyTreeIter` never visited or cleared its leaf predicate, so a cycle through that callback leaked. `PyTreeSpec` reports a registration's members once, and only when its own nodes hold every reference to it: the registration holds one reference to each member however many nodes point at it, so reporting per node decrements the same object once per node and underflows the collector's shadow refcount, which is fatal on debug builds. Counting the holders per treespec, rather than testing one node for sole ownership, lets a spec with repeated custom nodes report correctly. A cycle spanning several treespecs still survives, since none of them owns every reference; a strict `xfail` pins that. Neither slot may throw either: both are invoked from `extern "C"` code, where an escaping exception calls `std::terminate`, and `tp_clear` empties the traversal itself, so the sanity check would abort on a cleared but still-live spec. Walker safety. `PathsImpl`, `AccessorsImpl` and `BroadcastToCommonSuffixImpl` recurse once per tree level, so a deeply nested spec overflowed the native stack and crashed instead of raising. Thread a depth through them and raise `RecursionError` at `MAX_RECURSION_DEPTH`, lowered further for debug builds on Windows, where the frames are large and the throw still has to unwind them. A walk also hands control to arbitrary user code before its own state is consistent. `FlattenUpTo` pre-sized a `py::list`, which is garbage-collected with NULL slots from the moment it is created, so a custom `flatten_func` or a dict key `__hash__` could reach it through `gc.get_objects()` and read an unwritten slot; collect the leaves into a vector and build the list once every element is known. `ListGetItemAs` read the length once and then ran that same user code, so a list shrunk mid-loop was indexed past its buffer on versions without `PyList_GetItemRef`; bounds-check on every version instead. Also fixes `IsPrefix` snapshotting a dict node's subtree from the pristine traversal rather than the working copy (a processed ancestor may have relocated it), `broadcast_to_common_suffix` sorting the argument spec's live key list in place while building an error message, custom node entries being dropped when broadcasting, and `Transform` reading pending counts after `pop_back`. The `traverse` and `walk` stubs gain overloads for the accumulating form, and the pickle and iterator constructors document that `cls.__new__(cls)` leaves the C++ value unbuilt.
…otal ordering An integer accessor entry indexes the children a registered flatten function emitted, not every declared field. `children_fields` derived them from a field predicate instead, so a class holding a metadata or non-`init` field resolved an integer entry to the wrong attribute, and a class registered through the generic `register_pytree_node()` raised `IndexError`. Read the children from the record `register_node()` leaves on the class, taken from its own `__dict__` so a subclass does not inherit a stale one, and fall back to the declared field order when there is none. `register_node` set its `_FIELDS` guard before `register_pytree_node()` ran, so a failed registration left the class marked as registered and unusable afterwards. Set the guard only on success, and say in the error which API to use to register the class in another namespace. `InitVar` pseudo-fields are excluded from `dataclasses.fields()`, so they are neither children nor metadata and cannot round trip: reject them and point at the generic API. Both modules now document that the generated unflatten rebuilds via `cls(**fields)`, which re-runs `__init__`, so the round trip is exact only when it returns each field unchanged. Accessor identity was wrong in three ways. Equality and hashing compared the bytecode of `__call__` and `codify` rather than the methods, so two entry classes that happened to share an implementation compared equal. `PyTreeAccessor` overrode `__eq__` but inherited `tuple.__ne__`, leaving both `==` and `!=` false against a plain tuple of equal entries. `GetAttrEntry.codify` emitted invalid syntax for a field name that is not an identifier, now a `getattr()` call. `tree_broadcast_common` fills an internal tree with a private sentinel and then re-applied the caller's `is_leaf` to it, so a predicate that inspects leaf values could crash on the sentinel or collapse the filled subtree and under-replicate. Do not apply the predicate to the internal tree. The single-input `tree_broadcast_map` variants flattened their tree twice, which broke a one-shot flatten function and contradicted their documented equivalence with `tree_map`; they now delegate to it. The transpose-map family and `tree_partition` no longer require a leaf in the outer structure when an explicit `inner_treespec` leaves nothing to infer. `prefix_errors` raised `AssertionError` for a custom node whose per-instance entries differ while its metadata does not, a pair `broadcast_prefix` accepts, and raised `TypeError` while formatting a key mismatch between keys of different types. Both total-order sorts fell back incorrectly. `total_order_sorted` handed `key` to `sorted()`, so a `TypeError` from the callback was mistaken for a comparison failure and swallowed, and the callback ran twice per element; apply it up front instead. `TotalOrderSort` sorted the same list on both attempts, committing a partial order when the second comparison also raised; each attempt now sorts a copy and only a fully sorted one is kept. `treespec_entry()` and `treespec_child()` accept any `SupportsInt` or `SupportsIndex` object, which their annotations rejected. `treespec_is_prefix` and `treespec_is_suffix` gain the preorder semantics they actually implement: reflexive and transitive, but neither antisymmetric nor total. `tree_broadcast_common` no longer claims its results share one structure, and `broadcast_common` documents that it returns two lists of leaves rather than two pytrees.
XuehaiPan
force-pushed
the
fix/correctness-audit
branch
from
July 30, 2026 13:02
984d7c2 to
3c20a98
Compare
XuehaiPan
changed the title
There was a problem hiding this comment.
Pull request overview
This PR applies a broad adversarial correctness audit across optree’s C++ core and Python frontend, focusing on rare-path failures involving lock ordering, interpreter-scoped caching, treespec serialization safety, GC integration, and broadcasting/ordering edge cases.
Changes:
- Fixes deadlocks and re-entrancy hazards in the registry and
PyTreeIter, including GIL/lock-order inversions and re-entrantnext()hangs. - Hardens treespec behavior: safer pickling/unpickling validation, namespace merge safety checks, recursion-depth guards, and safer traversal/flattening behavior under adversarial user code.
- Updates Python utilities and APIs (
total_order_sorted, broadcasting/map helpers, accessor codification/equality) and adds extensive regression tests + changelog entries.
Reviewed changes
Copilot reviewed 34 out of 34 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_utils.py | Adds regression coverage for total_order_sorted key-callback error propagation and call counts. |
| tests/test_typing.py | Adds structseq field/accessor tests (incl. CPython unnamed slots) and subprocess tests for interpreter-safe type caches. |
| tests/test_registry.py | Adds subprocess deadlock/regression tests for registry locking, error paths, and namespace lookup precedence. |
| tests/test_prefix_errors.py | Adds regressions for heterogeneous dict keys in error formatting and custom-node entry variability handling. |
| tests/test_ops.py | Adds regressions for dict ordering with mid-sort failures, iterator re-entrancy, flatten safety, list shrinking under is_leaf, zero-leaf partitioning, and broadcasting sentinel behavior. |
| tests/test_dataclasses.py | Adds characterization + regressions for dataclass reconstruction pitfalls, InitVar rejection, guard leakage, and entry indexing. |
| tests/test_accessors.py | Adds regressions for accessor __ne__, entry hashing/equality robustness, codify escaping, and unnamed structseq entry repr. |
| tests/integrations/test_attrs.py | Adds characterization + regressions for attrs reconstruction/converters, guard leakage, and entry indexing. |
| tests/helpers.py | Adds skip marker for deferred type refs on free-threaded builds < 3.14. |
| tests/concurrent/test_subinterpreters.py | Adds subinterpreter churn tests for type-cache correctness and registry init failure rollback/deadlock avoidance. |
| src/treespec/treespec.cpp | Adds namespace rebind guards, recursion-depth protection in broadcasting/walking, and avoids in-place key sorting corruption. |
| src/treespec/traversal.cpp | Adds re-entrant iterator next() detection and strengthens iterator safety. |
| src/treespec/serialization.cpp | Improves treespec repr and significantly hardens pickled-state copying + validation. |
| src/treespec/richcomparison.cpp | Fixes prefix comparison to snapshot subtrees before overlapping permutations. |
| src/treespec/gc.cpp | Adjusts GC traversal/clear semantics for registrations and iterator leaf predicate visibility. |
| src/treespec/flatten.cpp | Fixes dict-order flag reads, avoids exposing partially-filled GC-tracked leaf lists, and hardens flattening behavior. |
| src/treespec/constructors.cpp | Uses consistent dict-order flag snapshotting, improves warning/error handling, and adds namespace rebind guard. |
| src/registry.cpp | Refactors (un)registration to avoid Python execution under locks, adds rollback for init failures, and fixes deadlock-prone ordering. |
| src/optree.cpp | Exposes unnamed-field marker, improves pickle protocol support via explicit __reduce__, and documents pybind11 construction limitations. |
| optree/utils.py | Reworks total_order_sorted to avoid swallowed callback errors, avoid partial in-place sort artifacts, and ensure stable fallbacks. |
| optree/typing.py | Makes type checks more robust, exposes unnamed-field marker, and fixes structseq field mapping (incl. unnamed-slot behavior). |
| optree/ops.py | Fixes zero-leaf transpose/partition behavior, rewrites broadcast helpers to avoid reapplying is_leaf to internal sentinels, and improves docs/types. |
| optree/integrations/attrs.py | Fixes integer entry field selection, improves docs/errors, and prevents guard leakage on failed registrations. |
| optree/dataclasses.py | Adds InitVar rejection, improves docs/errors, and prevents guard leakage on failed registrations. |
| optree/accessors.py | Fixes entry equality/hashing robustness, corrects codify for non-identifiers, adds unnamed-slot handling, and fixes accessor __ne__. |
| optree/_C.pyi | Updates typing: overloads for walk/traverse, broader accepted index types, and exposes unnamed-field marker. |
| include/optree/treespec.h | Adds recursion-depth special-casing for Windows debug/free-threaded, introduces iterator re-entrancy state, and adds helper declarations. |
| include/optree/registry.h | Adds dict-order flag snapshot API, registry status enum, and consistent sizing under a single lock. |
| include/optree/pytypes.h | Introduces WeakKeyCache, safer list access pre-3.13, C-API copy helpers, structseq field mapping by offset, and safer total-order sorting. |
| include/optree/pymacros.h | Adds safe access to PyStructSequence_UnnamedField across Python versions and hardens interpreter-id retrieval. |
| include/optree/hashing.h | Adds hashing/equality for (interpreter_id, handle) cache keys. |
| docs/source/spelling_wordlist.txt | Adds new technical terms used in docs. |
| CHANGELOG.md | Adds detailed “Fixed” entries for the audit findings and their resolutions. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+134
to
+137
| if (self.m_leaf_predicate) [[likely]] { | ||
| // Drop the owned leaf predicate reference to break cycles that pass through it. | ||
| Py_CLEAR(self.m_leaf_predicate->ptr()); | ||
| } |
Comment on lines
+452
to
+456
| // The keys must be hashable and distinct; a duplicate or unhashable key would | ||
| // collapse or fail when the dict is rebuilt, desyncing the keys from the children. | ||
| if (DistinctCount(node.node_data) != node.arity) [[unlikely]] { | ||
| throw malformed("the keys are not distinct"); | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
A systematic adversarial correctness audit of optree, landed as five self-contained commits, one per area. Every finding is a latent bug on a non-common path rather than a reported regression, and most are only observable under free-threading, subinterpreters, debug assertions, or on Windows. The series is kept bisectable rather than squashed, since a failure in this matrix is far easier to attribute per commit.
fix(registry,traversal): correct lock ordering and re-entrancyrefactor(pytypes): extract WeakKeyCache and make the type caches interpreter-safefix(treespec): map struct sequence fields by offset and validate pickled statefix(treespec): correct namespace merging, GC reporting and walker safetyfix(ops,accessors,utils): correct field selection, broadcasting and total orderingEach commit message describes its own fixes, and the
CHANGELOG.mdentries are distributed to match.Motivation and Context
These are latent correctness bugs rather than reported regressions. Reaching any of them takes a deliberate, non-common path: registering a node concurrently with flattening on another thread, an
is_leafpredicate that mutates the container being flattened, a__hash__that walksgc.get_objects(), a hand-built pickle payload, or a class registered generically rather than through the integration helper. Most are invisible on a release build of a single interpreter, which is why they had not been caught.The set grew over three adversarial review rounds, each trying to break the code rather than confirm it. Beyond the original audit it now also covers two segfaults (a pickled
CUSTOMnode naming a built-in type;flatten_up_to()exposing a garbage-collected list with uninitialized slots), an abort from a C++ exception crossing atp_traverseslot, three further places where Python ran while the registry lock was held, an out-of-bounds list read on Python 3.9-3.12, and an uninterruptible hang when a leaf predicate advances its own iterator.Not linked to an existing issue; the changes came out of an adversarial audit rather than a bug report.
Types of changes
What types of changes does your code introduce? Put an
xin all the boxes that apply:Implemented Tasks
sm_mutex, reject re-entrantPyTreeIteriteration, and read the registry size under a single lockWeakKeyCache, key the type caches by(interpreter_id, address), and take ownership before publishing an entryis_leaf, and keep insertion order when keys cannot be comparedChecklist
Go over all the following points, and put an
xin all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!
make format. (required)make lint. (required)make testpass. (required)