Skip to content

gh-120496: Make enum_iter thread safe #120591

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
eendebakpt wants to merge 4 commits into from
Closed

gh-120496: Make enum_iter thread safe #120591

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3323cf0
Make enum_iter thread safe
eendebakpt Jun 16, 2024
98f72e6
pep8
eendebakpt Jun 16, 2024
ba17086
fix tests on wasi
eendebakpt Jun 16, 2024
07d4574
📜🤖 Added by blurb_it.
blurb-it[bot] Jun 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
Make enum_iter thread safe
  • Loading branch information
@eendebakpt
eendebakpt committed Jun 16, 2024
commit 3323cf09c5a0400688b731ddacd6476d9559ee5b
3 changes: 3 additions & 0 deletions Include/internal/pycore_critical_section.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,8 @@ extern "C" {
# define Py_END_CRITICAL_SECTION() \
_PyCriticalSection_End(&_cs); \
}
# define Py_EXIT_CRITICAL_SECTION() \
_PyCriticalSection_End(&_cs);

# define Py_BEGIN_CRITICAL_SECTION2(a, b) \
{ \
Expand Down Expand Up @@ -155,6 +157,7 @@ extern "C" {
# define Py_BEGIN_CRITICAL_SECTION_MUT(mut)
# define Py_BEGIN_CRITICAL_SECTION(op)
# define Py_END_CRITICAL_SECTION()
# define Py_EXIT_CRITICAL_SECTION()
# define Py_BEGIN_CRITICAL_SECTION2(a, b)
# define Py_END_CRITICAL_SECTION2()
# define Py_BEGIN_CRITICAL_SECTION_SEQUENCE_FAST(original)
Expand Down
23 changes: 23 additions & 0 deletions Lib/test/test_enumerate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import sys
import pickle
import gc
import threading

from test import support

Expand Down Expand Up @@ -292,5 +293,27 @@ def enum(self, iterable, start=sys.maxsize + 1):
(sys.maxsize+3,'c')]


class EnumerateThreading(unittest.TestCase):
@staticmethod
def work(index, enum, start):
while True:
try:
value = next(enum)
except StopIteration:
break
else:
if value[0] + start != value[1]:
raise ValueError(f'enumerate returned pair {value}')

def test_threading(self):
number_of_threads = 4
n = 100
start = sys.maxsize-10
enum = enumerate(range(start, start+n))
worker_threads = [threading.Thread(target=self.work, args=[ii, enum, start]) for ii in range(number_of_threads)]
_ = [t.start() for t in worker_threads]
_ = [t.join() for t in worker_threads]


if __name__ == "__main__":
unittest.main()
26 changes: 21 additions & 5 deletions Objects/enumobject.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

#include "Python.h"
#include "pycore_call.h" // _PyObject_CallNoArgs()
#include "pycore_critical_section.h" // Py_BEGIN_CRITICAL_SECTION()
#include "pycore_long.h" // _PyLong_GetOne()
#include "pycore_modsupport.h" // _PyArg_NoKwnames()
#include "pycore_object.h" // _PyObject_GC_TRACK()
Expand Down Expand Up @@ -228,23 +229,38 @@ enum_next(enumobject *en)
PyObject *it = en->en_sit;
PyObject *old_index;
PyObject *old_item;
int reuse_result = 0;

Py_BEGIN_CRITICAL_SECTION(en);

next_item = (*Py_TYPE(it)->tp_iternext)(it);
if (next_item == NULL)
if (next_item == NULL) {
Py_EXIT_CRITICAL_SECTION();
return NULL;

if (en->en_index == PY_SSIZE_T_MAX)
return enum_next_long(en, next_item);
}
if (en->en_index == PY_SSIZE_T_MAX) {
// we hold on to the lock for the entire call to enum_next_long
result = enum_next_long(en, next_item);
Py_EXIT_CRITICAL_SECTION();
return result;
}

next_index = PyLong_FromSsize_t(en->en_index);
if (next_index == NULL) {
Py_DECREF(next_item);
Py_EXIT_CRITICAL_SECTION();
return NULL;
}
en->en_index++;

if (Py_REFCNT(result) == 1) {
reuse_result = Py_REFCNT(result) == 1;
if (reuse_result) {
// increment within the critical section
Py_INCREF(result);
}
Py_END_CRITICAL_SECTION();

if (reuse_result) {

@eendebakpt eendebakpt Jun 16, 2024

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since we hold the only references to result this can be outside the critical section.

old_index = PyTuple_GET_ITEM(result, 0);
old_item = PyTuple_GET_ITEM(result, 1);
PyTuple_SET_ITEM(result, 0, next_index);
Expand Down