fix: fix issues for rare cases across deadlocking, caching, serialization, GC, and Python frontend - #290
Open
XuehaiPan wants to merge 4 commits into
Open
fix: fix issues for rare cases across deadlocking, caching, serialization, GC, and Python frontend#290XuehaiPan wants to merge 4 commits into
XuehaiPan wants to merge 4 commits into
Conversation
XuehaiPan
added
bug
XuehaiPan
force-pushed
the
fix/correctness-audit
branch
2 times, most recently
from
July 27, 2026 17:43
01e1568 to
df646e0
Compare
Several registry entry points released the GIL while holding `sm_mutex`, which inverts the GIL to `sm_mutex` order against a concurrent flatten that holds the GIL and waits on the read lock, wedging the interpreter. - `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. - `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. `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 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. The lookup also 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. `hashing.h` gains generic `std::equal_to` and `std::not_equal_to` for `std::pair<T, U>`, mirroring the `std::hash` specialization already there, so the pair key works without a third copy of the element-wise comparison.
…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. `structseq_fields` is fixed the same way, the marker is exported as `PyStructSequence_UnnamedField`, the repr renders an unnamed slot as `<unnamed@N>`, and `StructSequenceEntry.codify` emits index access for a slot whose name is not an identifier. `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. `ToPicklable` copies the node's mutable containers so the state cannot alias the spec, and `__reduce__` makes protocols 0 and 1 reconstruct through `cls.__new__(cls)` instead of aborting.
…ursion 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`. GC. `PyTreeIter` never visited or cleared its leaf predicate, so a cycle through that callback leaked. `PyTreeSpec` reports the registration's members only when its node is their sole owner: the registration holds one reference to each member however many nodes point at it, so reporting while shared decrements the same object once per node and underflows the collector's shadow refcount, which is fatal on debug builds. Recursion. `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`, which is lowered further for debug builds on Windows, where the frames are large and the throw still has to unwind them. 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`.
XuehaiPan
force-pushed
the
fix/correctness-audit
branch
from
July 28, 2026 09:02
0a46312 to
2905055
Compare
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
1 participant
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 correctness audit of optree, landed as five self-contained commits. They are being pushed one at
a time so that CI validates each in isolation: much of what they fix is only observable under
free-threading, subinterpreters, debug assertions, or on Windows, so a green run per commit is worth
more than one green run at the end. Each commit compiles standalone with
OPTREE_CXX_WERROR=ON, sothe series stays bisectable. See the individual commit messages for details.
fix(registry): correct GIL and registry lock orderingrefactor(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 recursionfix(dataclasses,attrs,accessors,ops): correct field selection and broadcastingMotivation and Context
These are latent correctness bugs rather than reported regressions: thread-safety and memory-safety
issues that surface only under free-threading, subinterpreters, debug builds or crafted input, plus
a few semantic bugs in namespace handling and field selection. Most are invisible on a release build
of a single interpreter, which is why they had not been caught.
Not linked to an existing issue; the changes came out of a systematic 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
WeakKeyCacheand key the type caches by(interpreter_id, address)with per-interpreter cleanup<unnamed@N>, and validate pickled treespec statePyTreeIterGC hooks, and bound walker recursionis_leafto internal sentinelsChecklist
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)Local verification of the full series
pytest tests/: 93860 passed, 5 skipped with all five commits applied.OPTREE_CXX_WERROR=ON.ruff,pylint(10.00/10),mypy(no issues in 19 source files),clang-formatandaddlicenseclean.make doctestcannot complete locally becausetorchandjaxare not installed; the samedoctests fail identically on
main(for exampleoptree/functools.py, untouched here), so CIcovers them.
Windows stack limits) is not observable on a macOS release build, which is why each commit gets
its own CI run.