From 6a3a1c68fb316e63fc52e6654b9dd81f037e92cd Mon Sep 17 00:00:00 2001 From: Serhiy Storchaka Date: Wed, 29 Jul 2026 16:29:29 +0300 Subject: [PATCH] gh-138765: Release items cached by itertools.tee once all iterators pass them A teedataobject no longer owns the items it caches. They are owned by its direct referrers: the tee objects positioned in the link, and the preceding link, which stands for the tee objects that have not reached it yet. Each referrer drops its reference as it steps over an item, so an item is released as soon as the last tee object passes it, instead of when the whole 57-item link is retired. Co-Authored-By: Claude Opus 5 (1M context) --- Doc/library/itertools.rst | 3 + Lib/test/test_itertools.py | 29 +++ ...-07-29-18-40-12.gh-issue-138765.Kf3Wq1.rst | 3 + Modules/itertoolsmodule.c | 170 ++++++++++++++++-- 4 files changed, 186 insertions(+), 19 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2026-07-29-18-40-12.gh-issue-138765.Kf3Wq1.rst diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst index e1730608887b3d..3b4e1b5899e555 100644 --- a/Doc/library/itertools.rst +++ b/Doc/library/itertools.rst @@ -757,6 +757,9 @@ loops that truncate the stream. most or all of the data before another iterator starts, it is faster to use :func:`list` instead of :func:`tee`. + .. versionchanged:: 3.16 + An item is released as soon as every returned iterator has passed it. + .. function:: zip_longest(*iterables, fillvalue=None) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index cf579d4da4e0df..d5c30c8da15b70 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -1391,6 +1391,35 @@ def test_tee_dealloc_segfault(self): ) script_helper.assert_python_ok("-c", script) + @support.cpython_only + def test_tee_releases_passed_items(self): + # gh-138765: an item is released once every branch has passed it, + # not when the internal buffer which caches it is recycled. + class Item: + pass + + refs = [] + def source(): + for i in range(200): + item = Item() + refs.append(weakref.ref(item)) + yield item + del item + + a, b = tee(source()) + for i in range(150): + next(a) + for i in range(150): + next(b) + # The generator is suspended at the yield of the last consumed item, + # so its local variable still keeps that one alive. + alive = [i for i, ref in enumerate(refs[:149]) if ref() is not None] + self.assertEqual(alive, []) + # Items which the lagging branch has not passed yet are kept. + for i in range(10): + next(a) + self.assertTrue(all(ref() is not None for ref in refs[150:160])) + # Issue 13454: Crash when deleting backward iterator from tee() def test_tee_del_backward(self): forward, backward = tee(repeat(None, 20000000)) diff --git a/Misc/NEWS.d/next/Library/2026-07-29-18-40-12.gh-issue-138765.Kf3Wq1.rst b/Misc/NEWS.d/next/Library/2026-07-29-18-40-12.gh-issue-138765.Kf3Wq1.rst new file mode 100644 index 00000000000000..626acc15c0bfab --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-07-29-18-40-12.gh-issue-138765.Kf3Wq1.rst @@ -0,0 +1,3 @@ +Items cached by :func:`itertools.tee` are now released as soon as every +iterator has passed them, +instead of being kept until the internal buffer holding them is recycled. diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index c0023c839ca7fe..24a19ed915d623 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -727,11 +727,24 @@ static PyType_Spec _grouper_spec = { */ #define LINKCELLS 57 +/* A teedataobject does not own the items it caches. They are owned by its + direct referrers: every tee object currently positioned in this link, and + the preceding link (which stands for all tee objects that have not reached + this link yet). A referrer owns one reference to every item it can still + reach -- items[index:numread] for a tee object at position index, and + items[0:numread] for the preceding link. A tee object drops its reference + as it steps over an item, so the item is released as soon as the last + referrer passes it, instead of waiting for the whole link to be retired. + nrefs is the number of direct referrers; it is the number of references + that a newly cached item starts with. +*/ + typedef struct { PyObject_HEAD PyObject *it; - int numread; /* 0 <= numread <= LINKCELLS */ - int running; + short numread; /* 0 <= numread <= LINKCELLS */ + short running; + Py_ssize_t nrefs; /* number of direct referrers */ PyObject *nextlink; PyObject *(values[LINKCELLS]); } teedataobject; @@ -759,12 +772,50 @@ teedataobject_newinternal(itertools_state *state, PyObject *it) tdo->running = 0; tdo->numread = 0; + tdo->nrefs = 1; /* the caller, which stores the only reference */ tdo->nextlink = NULL; tdo->it = Py_NewRef(it); PyObject_GC_Track(tdo); return (PyObject *)tdo; } +/* Register a new referrer of tdo, positioned at index: it takes a reference + to every item it can still reach. */ +static void +teedataobject_attach_lock_held(teedataobject *tdo, int index) +{ + tdo->nrefs++; + for (int i = index; i < tdo->numread; i++) + Py_INCREF(tdo->values[i]); +} + +static void +teedataobject_attach(teedataobject *tdo, int index) +{ + Py_BEGIN_CRITICAL_SECTION(tdo); + teedataobject_attach_lock_held(tdo, index); + Py_END_CRITICAL_SECTION(); +} + +/* Unregister a referrer of tdo positioned at index; the items it still owns + are released here. */ +static void +teedataobject_detach_lock_held(teedataobject *tdo, int index) +{ + assert(tdo->nrefs > 0); + tdo->nrefs--; + for (int i = index; i < tdo->numread; i++) + Py_DECREF(tdo->values[i]); +} + +static void +teedataobject_detach(teedataobject *tdo, int index) +{ + Py_BEGIN_CRITICAL_SECTION(tdo); + teedataobject_detach_lock_held(tdo, index); + Py_END_CRITICAL_SECTION(); +} + static PyObject * teedataobject_jumplink(itertools_state *state, teedataobject *tdo) { @@ -774,11 +825,17 @@ teedataobject_jumplink(itertools_state *state, teedataobject *tdo) tdo->nextlink = teedataobject_newinternal(state, tdo->it); link = Py_XNewRef(tdo->nextlink); Py_END_CRITICAL_SECTION(); + if (link != NULL) { + /* The tee object that jumps here becomes a new referrer. */ + teedataobject_attach(teedataobject_CAST(link), 0); + } return link; } +/* If steal is true, the caller takes over the reference owned by the tee + object it advances; the caller must advance past the item. */ static PyObject * -teedataobject_getitem_lock_held(teedataobject *tdo, int i) +teedataobject_getitem_lock_held(teedataobject *tdo, int i, int steal) { PyObject *value; @@ -798,18 +855,22 @@ teedataobject_getitem_lock_held(teedataobject *tdo, int i) tdo->running = 0; if (value == NULL) return NULL; + /* The new item is owned by every referrer of this link; PyIter_Next() + already provided one of those references. */ + assert(tdo->nrefs > 0); + _Py_RefcntAdd(value, tdo->nrefs - 1); tdo->numread++; tdo->values[i] = value; } - return Py_NewRef(value); + return steal ? value : Py_NewRef(value); } static PyObject * -teedataobject_getitem(teedataobject *tdo, int i) +teedataobject_getitem(teedataobject *tdo, int i, int steal) { PyObject *result; Py_BEGIN_CRITICAL_SECTION(tdo); - result = teedataobject_getitem_lock_held(tdo, i); + result = teedataobject_getitem_lock_held(tdo, i, steal); Py_END_CRITICAL_SECTION(); return result; } @@ -822,8 +883,13 @@ teedataobject_traverse(PyObject *op, visitproc visit, void * arg) Py_VISIT(Py_TYPE(tdo)); Py_VISIT(tdo->it); - for (i = 0; i < tdo->numread; i++) - Py_VISIT(tdo->values[i]); + /* The items of the next link are owned by this one, its own items are + not owned by it. */ + if (tdo->nextlink != NULL) { + teedataobject *next = teedataobject_CAST(tdo->nextlink); + for (i = 0; i < next->numread; i++) + Py_VISIT(next->values[i]); + } Py_VISIT(tdo->nextlink); return 0; } @@ -838,6 +904,8 @@ teedataobject_safe_decref(PyObject *obj) nextlink = tmp->nextlink; tmp->nextlink = NULL; Py_END_CRITICAL_SECTION(); + if (nextlink != NULL) + teedataobject_detach(teedataobject_CAST(nextlink), 0); Py_SETREF(obj, nextlink); } Py_XDECREF(obj); @@ -846,17 +914,19 @@ teedataobject_safe_decref(PyObject *obj) static int teedataobject_clear(PyObject *op) { - int i; PyObject *tmp; teedataobject *tdo = teedataobject_CAST(op); Py_BEGIN_CRITICAL_SECTION(op); Py_CLEAR(tdo->it); - for (i=0 ; inumread ; i++) - Py_CLEAR(tdo->values[i]); tmp = tdo->nextlink; tdo->nextlink = NULL; Py_END_CRITICAL_SECTION(); + /* The values are not owned by this link, they are released by their + owners. Only the reference to the items of the next link is dropped + here. */ + if (tmp != NULL) + teedataobject_detach(teedataobject_CAST(tmp), 0); teedataobject_safe_decref(tmp); return 0; } @@ -865,8 +935,16 @@ static void teedataobject_dealloc(PyObject *op) { PyTypeObject *tp = Py_TYPE(op); + teedataobject *tdo = teedataobject_CAST(op); PyObject_GC_UnTrack(op); (void)teedataobject_clear(op); + /* Referrers that never detached (a link created from Python code is + owned by the code that created it) still hold references. */ + while (tdo->nrefs > 0) { + tdo->nrefs--; + for (int i = 0; i < tdo->numread; i++) + Py_DECREF(tdo->values[i]); + } PyObject_GC_Del(op); Py_DECREF(tp); } @@ -904,13 +982,16 @@ itertools_teedataobject_impl(PyTypeObject *type, PyObject *it, Py_INCREF(tdo->values[i]); } /* len <= LINKCELLS < INT_MAX */ - tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int); + tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, short); if (len == LINKCELLS) { if (next != Py_None) { if (!Py_IS_TYPE(next, state->teedataobject_type)) goto err; assert(tdo->nextlink == NULL); + /* this link now owns the items of the next one; take the + references before teedataobject_traverse() reports them */ + teedataobject_attach(teedataobject_CAST(next), 0); tdo->nextlink = Py_NewRef(next); } } else { @@ -959,10 +1040,12 @@ tee_next(PyObject *op) if (link == NULL) { return NULL; } + teedataobject_detach(to->dataobj, to->index); Py_SETREF(to->dataobj, (teedataobject *)link); to->index = 0; } - value = teedataobject_getitem(to->dataobj, to->index); + /* Advancing past the item hands its reference to the caller. */ + value = teedataobject_getitem(to->dataobj, to->index, 1); if (value == NULL) { return NULL; } @@ -982,13 +1065,39 @@ tee_next(PyObject *op) Py_END_CRITICAL_SECTION(); if (index < LINKCELLS) { - value = teedataobject_getitem(dataobj, index); + /* If the item is already cached, hand over the reference owned by + this tee object and advance in the same atomic step: taking + both locks at once is what makes the hand over safe. */ + value = NULL; + Py_BEGIN_CRITICAL_SECTION2(op, (PyObject *)dataobj); + if (to->dataobj == dataobj && to->index == index + && index < dataobj->numread) + { + value = dataobj->values[index]; + to->index = index + 1; + } + Py_END_CRITICAL_SECTION2(); + if (value != NULL) { + Py_DECREF(dataobj); + return value; + } + /* The item is not there yet, it has to be read from the source + iterator, which cannot be done atomically with the advance. + The reference owned by this tee object cannot be handed to the + caller then, because the advance may lose the race below. */ + value = teedataobject_getitem(dataobj, index, 0); if (value != NULL) { + int advanced = 0; Py_BEGIN_CRITICAL_SECTION(op); if (to->dataobj == dataobj && to->index == index) { to->index = index + 1; + advanced = 1; } Py_END_CRITICAL_SECTION(); + if (advanced) { + /* stepping over the item releases it */ + Py_DECREF(value); + } } Py_DECREF(dataobj); return value; @@ -999,14 +1108,22 @@ tee_next(PyObject *op) Py_DECREF(dataobj); return NULL; } - Py_BEGIN_CRITICAL_SECTION(op); + /* Both locks are needed at once: this tee object stops referring to + dataobj exactly when it starts referring to the new link. */ + Py_BEGIN_CRITICAL_SECTION2(op, (PyObject *)dataobj); if (to->dataobj == dataobj) { + assert(to->index == LINKCELLS); + teedataobject_detach_lock_held(dataobj, to->index); Py_SETREF(to->dataobj, (teedataobject *)link); to->index = 0; link = NULL; } - Py_END_CRITICAL_SECTION(); - Py_XDECREF(link); + Py_END_CRITICAL_SECTION2(); + if (link != NULL) { + /* another thread jumped first: undo the attach */ + teedataobject_detach(teedataobject_CAST(link), 0); + Py_DECREF(link); + } Py_DECREF(dataobj); } #endif @@ -1017,6 +1134,11 @@ tee_traverse(PyObject *op, visitproc visit, void *arg) { teeobject *to = teeobject_CAST(op); Py_VISIT(Py_TYPE(to)); + /* the items which are not passed yet are owned by this tee object */ + if (to->dataobj != NULL) { + for (int i = to->index; i < to->dataobj->numread; i++) + Py_VISIT(to->dataobj->values[i]); + } Py_VISIT((PyObject *)to->dataobj); return 0; } @@ -1028,10 +1150,18 @@ tee_copy_impl(teeobject *to) if (newto == NULL) { return NULL; } + teedataobject *dataobj; + int index; Py_BEGIN_CRITICAL_SECTION(to); - newto->dataobj = (teedataobject *)Py_NewRef(to->dataobj); - newto->index = to->index; + dataobj = (teedataobject *)Py_NewRef(to->dataobj); + index = to->index; Py_END_CRITICAL_SECTION(); + /* The copy is a new referrer of the same link. Its references must be + taken before the copy is tracked, because tee_traverse() reports them + as soon as it is. */ + teedataobject_attach(dataobj, index); + newto->dataobj = dataobj; + newto->index = index; newto->weakreflist = NULL; newto->state = to->state; PyObject_GC_Track(newto); @@ -1103,6 +1233,8 @@ tee_clear(PyObject *op) teeobject *to = teeobject_CAST(op); if (to->weakreflist != NULL) PyObject_ClearWeakRefs(op); + if (to->dataobj != NULL) + teedataobject_detach(to->dataobj, to->index); Py_CLEAR(to->dataobj); return 0; }