Files
RustPython/extra_tests/snippets/stdlib_threading_itertools_tee.py
Jeong, YunWon 259ce8ce66 Fix the itertools.tee leak and race, the contextvars and generator races, _asyncio's InvalidStateError lookup, and _imp's frozen data (#8518)
* _imp: read get_frozen_object's data as a whole marshal value

The data argument went through deserialize_code(), which reads a code body
without the type byte in front of it, so nothing marshal.dumps() produces was
accepted. Read it with marshal.loads() and require a code object back.

Assisted-by: Claude

* _imp: bind find_frozen's arguments with FromArgs

The arity check was done by hand on a FuncArgs. Take the arguments through a
FromArgs struct instead, which makes withdata keyword-only, and fill in the
data it asks for: the frozen encoding is not marshal, so the code is
re-serialized into what get_frozen_object() reads back.

Assisted-by: Claude

* itertools: give tee's shared buffer a Python type

The buffer was a PyRc<PyItertoolsTeeData>, which is not a Python object, so the
collector could not walk into it and any cycle running through a tee was
uncollectable. Make it the _tee_dataobject type with a traverse, held by
PyRef, and split the rest along the same lines: tee() is a function, _tee is
the iterator type it builds, and _tee takes a single iterable rather than
returning a tuple from __new__. _tee is weak-referenceable, tee() rejects a
negative n with ValueError and reserves its tuple fallibly.

test_itertools.test_tee passes now; its expectedFailure marker is removed.

Assisted-by: Claude

* contextvars: hold the shared context state under locks

The variable map was a RefCell, the enter flag, the context index, the token's
used flag and the variable hash were Cells, and each carried an unsafe impl
Sync. A Context or ContextVar shared between threads overlapped their borrows
and panicked. The map and the per-variable cache are now PyMutex, the flags and
the index are atomics, entering a context is a compare_exchange, and the three
unsafe impl Sync are gone.

The cache also stopped being read through AtomicCell::as_ptr, which raced a
concurrent store on a value holding a PyObjectRef.

Values displaced from the map or the cache are dropped after the lock is
released: __del__ can come straight back into the same context, and the locks
are not reentrant.

Assisted-by: Claude

* generator: read the frame state under the running claim

send(), send_none(), throw() and close() read `closed` and `frame.lasti()`
before `running` was compare_exchanged, so the frame they went on to resume
could be one another thread had already advanced. A resume that decided from
`lasti() == 0` that the generator had not started pushes no value onto the
value stack, and the code after the yield pops one, which underflows the stack.

The compare_exchange now hands back a guard, taken before those reads and
released after maybe_close(), so the generator is retired while it is still
claimed.

Assisted-by: Claude

* itertools: claim a tee's position and its buffer as one step

`_tee::next` read `index` and moved it on afterwards, and
`_tee_dataobject::get_item` released `running` before the value it fetched
from the source was cached.

Two callers on one `_tee` then read the same index, hand out the same value
twice and advance past a value that was never cached, which leaves `index`
past `values.len()` and indexes the buffer out of bounds. Two callers at the
same index on separate tees each fetch a value from the source, and one of the
two is dropped without reaching a caller.

Both claims now cover the read and the update.

Assisted-by: Claude

* _asyncio: require InvalidStateError to be a type

new_invalid_state_error() called whatever `asyncio.exceptions.InvalidStateError`
names and unwrapped the downcast of the result, so a future asked for a result
it does not have panicked once that attribute was rebound to something that is
not an exception:

    asyncio.InvalidStateError = lambda *args: 42
    _asyncio.Future(loop=object()).result()

The type is looked up the way get_cancelled_error_type() looks its own up, and
raised with new_exception_msg; a lookup that does not produce an exception type
falls back to RuntimeError, as the other arms already did.

Assisted-by: Claude
2026-08-14 16:05:30 +09:00

58 lines
1.7 KiB
Python

"""Advance the iterators of one tee() from several threads at once.
Every tee iterator reads its position, asks the shared buffer for that item and
then moves the position on. Reading and moving it on has to be one step, and
the buffer has to stay claimed until the value it fetched from the source is
cached: otherwise two callers work on the same index, a fetched value is
dropped, and the buffer is left to be filled out of order.
A caller that loses the race gets a RuntimeError, never a value another caller
has already been handed.
"""
import itertools
import threading
ROUNDS = 200
WORKERS = 4
errors = []
def drain(iterator, out):
for _ in range(ROUNDS):
try:
out.append(next(iterator))
except StopIteration:
break
except RuntimeError:
# another thread is advancing this tee
pass
except Exception as exc: # noqa: BLE001
errors.append(exc)
break
for _ in range(10):
first, second = itertools.tee(iter(range(ROUNDS * WORKERS)))
taken = [[] for _ in range(WORKERS)]
threads = [
threading.Thread(target=drain, args=(first if i % 2 else second, taken[i]))
for i in range(WORKERS)
]
for t in threads:
t.start()
for t in threads:
t.join()
assert not errors, errors
for got in taken:
# one iterator hands out ascending values, each of them once
assert got == sorted(set(got)), got
for side in (taken[1], taken[3]), (taken[0], taken[2]):
# the two threads sharing an iterator split its values between them
shared = side[0] + side[1]
assert len(shared) == len(set(shared)), shared
print("ok")