FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

gh-154535: Fix crash when a contextvars.Context iterator is shared between threads by gianghungtien · Pull Request #154586 · python/cpython · GitHub

Repository navigation

Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension .c  (1) .h  (1) .py  (1) .rst  (1) All 4 file types selected
Viewed files
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Unified
Split
Hide whitespace
Diff view
Unified
Split
Hide whitespace
5 changes: 5 additions & 0 deletions Include/internal/pycore_hamt.h
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
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,11 @@ extern PyTypeObject _PyHamtItems_Type;
- i_nodes: an array of seven pointers to tree nodes
- i_level: the current node in i_nodes
- i_pos: an array of positions within nodes in i_nodes.

Because i_nodes holds borrowed references, advancing the state is only
safe while nothing else can observe it half-updated. When the state is
embedded in a PyHamtIterator (as opposed to living on the stack), it must
be advanced with the iterator locked -- see hamt_baseiter_tp_iternext().
*/
typedef struct {
PyHamtNode *i_nodes[_Py_HAMT_MAX_TREE_DEPTH];
Expand Down
39 changes: 39 additions & 0 deletions Lib/test/test_free_threading/test_iteration.py
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
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import contextvars
import sys
import threading
import unittest
Expand Down Expand Up @@ -114,6 +115,44 @@ def worker():
self.assert_iterator_results(results, list(seq))


class ContendedContextIterationTest(ContendedTupleIterationTest):
# A contextvars.Context is iterated by the HAMT iterator, which keeps its
# whole depth-first cursor -- including borrowed node pointers -- inside
# the iterator object, so sharing one between threads used to crash the
# interpreter (gh-154535).
def make_testdata(self, n):
variables = [contextvars.ContextVar(f'v{i}') for i in range(n)]
ctx = contextvars.Context()
def populate():
for i, var in enumerate(variables):
var.set(i)
ctx.run(populate)
return ctx

def test_restarted_shared_iterator(self):
"""Test a shared iterator that is restarted when it is exhausted"""
# A single pass over a shared iterator is over too quickly to reliably
# desync its cursor. Keep replacing the iterator once it is exhausted
# so that the threads stay on top of each other for a while.
seq = self.make_testdata(16)
cell = [iter(seq)]
results = []
start = threading.Barrier(NUMTHREADS)
def worker():
items = []
start.wait()
for _ in range(NUMITEMS):
try:
items.append(next(cell[0]))
except StopIteration:
cell[0] = iter(seq)
results.extend(items)
threads = self.run_threads(worker)
for t in threads:
t.join()
self.assert_iterator_results(results, seq)


class ContendedSeqIterExhaustionTest(unittest.TestCase):
"""Test draining a shared iter() fallback iterator (PySeqIter_Type).

Expand Down
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
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
Fix a crash on the :term:`free-threaded <free threading>` build when a single
iterator over a :class:`contextvars.Context` (as returned by :func:`iter`,
:meth:`~contextvars.Context.keys`, :meth:`~contextvars.Context.values` or
:meth:`~contextvars.Context.items`) is advanced from several threads at once.
The underlying HAMT iterator now advances its traversal cursor while holding
the iterator lock.
23 changes: 22 additions & 1 deletion Python/hamt.c
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
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#include "Python.h"
#include "pycore_bitutils.h" // _Py_popcount32()
#include "pycore_critical_section.h" // Py_BEGIN_CRITICAL_SECTION()
#include "pycore_hamt.h"
#include "pycore_initconfig.h" // _PyStatus_OK()
#include "pycore_long.h" // _PyLong_Format()
Expand Down Expand Up @@ -2480,7 +2481,27 @@ hamt_baseiter_tp_iternext(PyObject *op)
PyHamtIterator *it = (PyHamtIterator*)op;
PyObject *key;
PyObject *val;
hamt_iter_t res = hamt_iterator_next(&it->hi_iter, &key, &val);
hamt_iter_t res;

/* The depth-first cursor (i_level, i_pos and i_nodes) is stored in the
iterator and is updated by every step of the descent, so two threads
advancing the same iterator would interleave their updates and leave
i_level pointing at an i_nodes slot that was never filled in. Since
i_nodes holds borrowed pointers, that desync makes the next step
dereference a stale or NULL node. Serialize the whole descent on the
iterator; the HAMT itself is immutable, so nothing else needs locking.

key and val are borrowed references into the tree and stay valid after
the lock is released: they are pinned via hi_obj, which this function
never drops, and tp_clear (the only other code that drops it) cannot
run concurrently because an in-flight tp_iternext call implies the
caller owns a reference to the iterator, keeping it reachable.

Concurrent iteration can still skip or repeat items, which matches the
behaviour of the other free-threaded iterators. */
Py_BEGIN_CRITICAL_SECTION(op);
res = hamt_iterator_next(&it->hi_iter, &key, &val);
Py_END_CRITICAL_SECTION();

switch (res) {
case I_END:
Expand Down
Loading

Back | FazBrowse Home | New Git URL