diff --git a/Misc/NEWS.d/next/Library/2018-09-06-15-40-37.bpo-34410.Mk8NtD.rst b/Misc/NEWS.d/next/Library/2018-09-06-15-40-37.bpo-34410.Mk8NtD.rst new file mode 100644 index 00000000000000..c1851258ebac8f --- /dev/null +++ b/Misc/NEWS.d/next/Library/2018-09-06-15-40-37.bpo-34410.Mk8NtD.rst @@ -0,0 +1 @@ +``itertools.tee()`` fix it can segfault interpreter when wrapped iterator releases GIL. \ No newline at end of file diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 8a36755bfa7265..59d378cf762e6e 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -391,6 +391,14 @@ static PyTypeObject _grouper_type = { */ #define LINKCELLS 57 +struct lead_rlock { + PyThread_type_lock lock; /* lead iterator lock */ + unsigned long owner; + unsigned long count; + unsigned long sharenum; + int needed; +}; + typedef struct { PyObject_HEAD PyObject *it; @@ -404,10 +412,71 @@ typedef struct { teedataobject *dataobj; int index; /* 0 <= index <= LINKCELLS */ PyObject *weakreflist; + struct lead_rlock *rlock; } teeobject; static PyTypeObject teedataobject_type; +static PyLockStatus +lead_rlock_smart_acquire(struct lead_rlock *rlock) +{ + unsigned long tid; + PyLockStatus r = PY_LOCK_ACQUIRED; + + tid = PyThread_get_thread_ident(); + if (rlock->owner == tid) + { + unsigned long count = rlock->count + 1; + if (count <= rlock->count) { + PyErr_SetString(PyExc_OverflowError, + "Internal lock count overflowed"); + return PY_LOCK_FAILURE; + } + rlock->count = count; + return PY_LOCK_ACQUIRED; + } + + if(rlock->owner != 0 && rlock->owner != tid) { + if(rlock->count == 0) { + /* It's safe to release. */ + PyThread_release_lock(rlock->lock); + rlock->owner = 0; + }else { + rlock->needed = 1; + } + } + + if (!PyThread_acquire_lock_timed(rlock->lock, 0, NOWAIT_LOCK)) { + Py_BEGIN_ALLOW_THREADS + r = PyThread_acquire_lock_timed(rlock->lock, -1, WAIT_LOCK); + Py_END_ALLOW_THREADS + } + + if (r == PY_LOCK_ACQUIRED) { + assert(rlock->count == 0); + rlock->owner = tid; + rlock->count = 1; + return r; + } + return PY_LOCK_FAILURE; +} + +static void +lead_rlock_smart_release(struct lead_rlock *rlock) +{ + if (rlock->count == 0) { + PyErr_SetString(PyExc_RuntimeError, + "cannot release un-acquired lock"); + return; + } + if (--rlock->count == 0) { + if(rlock->needed) { + PyThread_release_lock(rlock->lock); + rlock->owner = 0; + } + } +} + static PyObject * teedataobject_newinternal(PyObject *it) { @@ -634,7 +703,15 @@ tee_next(teeobject *to) Py_SETREF(to->dataobj, (teedataobject *)link); to->index = 0; } - value = teedataobject_getitem(to->dataobj, to->index); + + if (to->index < to->dataobj->numread) { + value = to->dataobj->values[to->index]; + Py_INCREF(value); + } else { + lead_rlock_smart_acquire(to->rlock); + value = teedataobject_getitem(to->dataobj, to->index); + lead_rlock_smart_release(to->rlock); + } if (value == NULL) return NULL; to->index++; @@ -658,6 +735,8 @@ tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored)) return NULL; Py_INCREF(to->dataobj); newto->dataobj = to->dataobj; + newto->rlock = to->rlock; + newto->rlock->sharenum++; newto->index = to->index; newto->weakreflist = NULL; PyObject_GC_Track(newto); @@ -691,6 +770,14 @@ tee_fromiterable(PyObject *iterable) } to->index = 0; + to->rlock = (struct lead_rlock *)PyMem_RawMalloc(sizeof(struct lead_rlock)); + if (to->rlock == NULL) + return NULL; + to->rlock->lock = PyThread_allocate_lock(); + to->rlock->owner = 0; + to->rlock->count = 0; + to->rlock->sharenum = 1; + to->rlock->needed = 0; to->weakreflist = NULL; PyObject_GC_Track(to); done: @@ -720,6 +807,11 @@ tee_clear(teeobject *to) static void tee_dealloc(teeobject *to) { + if(--to->rlock->sharenum == 0) { + if (to->rlock->count > 0) + PyThread_release_lock(to->rlock->lock); + PyMem_RawFree((void *) to->rlock); + } PyObject_GC_UnTrack(to); tee_clear(to); PyObject_GC_Del(to);