Skip to content

Commit 5ab7663

Browse files
colesburycolesbury
committed
gh-xxxxxx: Generator send thread-safety
1 parent 4a8a36a commit 5ab7663

File tree

10 files changed

+111
-62
lines changed

10 files changed

+111
-62
lines changed

‎Include/internal/pycore_genobject.h‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,12 @@ extern "C" {
1010

1111
#include "pycore_frame.h"
1212

13+
#ifdef Py_GIL_DISABLED
14+
# define _PY_GEN_OUT_FRAME_STATE(prefix) int8_t *prefix##_out_frame_state
15+
#else
16+
# define _PY_GEN_OUT_FRAME_STATE(prefix)
17+
#endif
18+
1319
/* _PyGenObject_HEAD defines the initial segment of generator
1420
and coroutine objects. */
1521
#define _PyGenObject_HEAD(prefix) \
@@ -22,6 +28,7 @@ extern "C" {
2228
PyObject *prefix##_qualname; \
2329
_PyErr_StackItem prefix##_exc_state; \
2430
PyObject *prefix##_origin_or_finalizer; \
31+
_PY_GEN_OUT_FRAME_STATE(prefix); \
2532
char prefix##_hooks_inited; \
2633
char prefix##_closed; \
2734
char prefix##_running_async; \
@@ -52,6 +59,32 @@ PyGenObject *_PyGen_GetGeneratorFromFrame(_PyInterpreterFrame *frame)
5259
return (PyGenObject *)(((char *)frame) - offset_in_gen);
5360
}
5461

62+
static inline void
63+
_PyGen_SetFrameState(PyGenObject *gen, int8_t state)
64+
{
65+
#ifdef Py_GIL_DISABLED
66+
if (gen->gi_out_frame_state) {
67+
*gen->gi_out_frame_state = state;
68+
gen->gi_out_frame_state = NULL;
69+
}
70+
_Py_atomic_store_int8(&gen->gi_frame_state, state);
71+
#else
72+
gen->gi_frame_state = state;
73+
#endif
74+
}
75+
76+
static inline int
77+
_PyGen_TransitionFrameState(PyGenObject *gen, int8_t *from_state, int8_t state)
78+
{
79+
#ifdef Py_GIL_DISABLED
80+
return _Py_atomic_compare_exchange_int8(&gen->gi_frame_state, from_state, state);
81+
#else
82+
assert(gen->gi_frame_state == *from_state);
83+
gen->gi_frame_state = state;
84+
return 1;
85+
#endif
86+
}
87+
5588
PyAPI_FUNC(PyObject *)_PyGen_yf(PyGenObject *);
5689
extern void _PyGen_Finalize(PyObject *self);
5790

‎Include/internal/pycore_pyatomic_ft_wrappers.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@ extern "C" {
3535
_Py_atomic_load_uintptr_acquire(&value)
3636
#define FT_ATOMIC_LOAD_PTR_RELAXED(value) \
3737
_Py_atomic_load_ptr_relaxed(&value)
38+
#define FT_ATOMIC_LOAD_INT8_RELAXED(value) \
39+
_Py_atomic_load_int8_relaxed(&value)
3840
#define FT_ATOMIC_LOAD_UINT8(value) \
3941
_Py_atomic_load_uint8(&value)
4042
#define FT_ATOMIC_STORE_UINT8(value, new_value) \
@@ -55,6 +57,8 @@ extern "C" {
5557
_Py_atomic_store_uintptr_release(&value, new_value)
5658
#define FT_ATOMIC_STORE_SSIZE_RELAXED(value, new_value) \
5759
_Py_atomic_store_ssize_relaxed(&value, new_value)
60+
#define FT_ATOMIC_STORE_INT8(value, new_value) \
61+
_Py_atomic_store_int8(&value, new_value)
5862
#define FT_ATOMIC_STORE_UINT8_RELAXED(value, new_value) \
5963
_Py_atomic_store_uint8_relaxed(&value, new_value)
6064
#define FT_ATOMIC_STORE_UINT16_RELAXED(value, new_value) \
@@ -119,6 +123,7 @@ extern "C" {
119123
#define FT_ATOMIC_LOAD_PTR_ACQUIRE(value) value
120124
#define FT_ATOMIC_LOAD_UINTPTR_ACQUIRE(value) value
121125
#define FT_ATOMIC_LOAD_PTR_RELAXED(value) value
126+
#define FT_ATOMIC_LOAD_INT8_RELAXED(value) value
122127
#define FT_ATOMIC_LOAD_UINT8(value) value
123128
#define FT_ATOMIC_STORE_UINT8(value, new_value) value = new_value
124129
#define FT_ATOMIC_LOAD_UINT8_RELAXED(value) value
@@ -129,6 +134,7 @@ extern "C" {
129134
#define FT_ATOMIC_STORE_PTR_RELEASE(value, new_value) value = new_value
130135
#define FT_ATOMIC_STORE_UINTPTR_RELEASE(value, new_value) value = new_value
131136
#define FT_ATOMIC_STORE_SSIZE_RELAXED(value, new_value) value = new_value
137+
#define FT_ATOMIC_STORE_INT8(value, new_value) value = new_value
132138
#define FT_ATOMIC_STORE_UINT8_RELAXED(value, new_value) value = new_value
133139
#define FT_ATOMIC_STORE_UINT16_RELAXED(value, new_value) value = new_value
134140
#define FT_ATOMIC_STORE_UINT32_RELAXED(value, new_value) value = new_value

‎Lib/test/test_sys.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1662,7 +1662,10 @@ def bar(cls):
16621662
check(bar, size('PP'))
16631663
# generator
16641664
def get_gen(): yield 1
1665-
check(get_gen(), size('6P4c' + INTERPRETER_FRAME + 'P'))
1665+
if support.Py_GIL_DISABLED:
1666+
check(get_gen(), size('7P4c' + INTERPRETER_FRAME + 'P'))
1667+
else:
1668+
check(get_gen(), size('6P4c' + INTERPRETER_FRAME + 'P'))
16661669
# iterator
16671670
check(iter('abc'), size('lP'))
16681671
# callable-iterator

‎Objects/genobject.c‎

Lines changed: 39 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -190,9 +190,12 @@ gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
190190
{
191191
PyThreadState *tstate = _PyThreadState_GET();
192192
_PyInterpreterFrame *frame = &gen->gi_iframe;
193+
int8_t frame_state;
193194

194195
*presult = NULL;
195-
if (gen->gi_frame_state == FRAME_CREATED && arg && arg != Py_None) {
196+
retry:
197+
frame_state = FT_ATOMIC_LOAD_INT8_RELAXED(gen->gi_frame_state);
198+
if (frame_state == FRAME_CREATED && arg && arg != Py_None) {
196199
const char *msg = "can't send non-None value to a "
197200
"just-started generator";
198201
if (PyCoro_CheckExact(gen)) {
@@ -205,7 +208,7 @@ gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
205208
PyErr_SetString(PyExc_TypeError, msg);
206209
return PYGEN_ERROR;
207210
}
208-
if (gen->gi_frame_state == FRAME_EXECUTING) {
211+
if (frame_state == FRAME_EXECUTING) {
209212
const char *msg = "generator already executing";
210213
if (PyCoro_CheckExact(gen)) {
211214
msg = "coroutine already executing";
@@ -216,7 +219,7 @@ gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
216219
PyErr_SetString(PyExc_ValueError, msg);
217220
return PYGEN_ERROR;
218221
}
219-
if (FRAME_STATE_FINISHED(gen->gi_frame_state)) {
222+
if (FRAME_STATE_FINISHED(frame_state)) {
220223
if (PyCoro_CheckExact(gen) && !closing) {
221224
/* `gen` is an exhausted coroutine: raise an error,
222225
except when called from gen_close(), which should
@@ -234,8 +237,12 @@ gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
234237
return PYGEN_ERROR;
235238
}
236239

237-
assert((gen->gi_frame_state == FRAME_CREATED) ||
238-
FRAME_STATE_SUSPENDED(gen->gi_frame_state));
240+
assert((frame_state == FRAME_CREATED) ||
241+
FRAME_STATE_SUSPENDED(frame_state));
242+
243+
if (!_PyGen_TransitionFrameState(gen, &frame_state, FRAME_EXECUTING)) {
244+
goto retry;
245+
}
239246

240247
/* Push arg onto the frame's value stack */
241248
PyObject *arg_obj = arg ? arg : Py_None;
@@ -245,23 +252,36 @@ gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
245252
gen->gi_exc_state.previous_item = prev_exc_info;
246253
tstate->exc_info = &gen->gi_exc_state;
247254

255+
#ifdef Py_GIL_DISABLED
256+
/* Store the resulting frame_state in frame_state */
257+
assert(gen->gi_out_frame_state == NULL);
258+
gen->gi_out_frame_state = &frame_state;
259+
#endif
260+
248261
if (exc) {
249262
assert(_PyErr_Occurred(tstate));
250263
_PyErr_ChainStackItem();
251264
}
252265

253-
gen->gi_frame_state = FRAME_EXECUTING;
254266
EVAL_CALL_STAT_INC(EVAL_CALL_GENERATOR);
255267
PyObject *result = _PyEval_EvalFrame(tstate, frame, exc);
256-
assert(tstate->exc_info == prev_exc_info);
257-
assert(gen->gi_exc_state.previous_item == NULL);
258-
assert(gen->gi_frame_state != FRAME_EXECUTING);
268+
269+
#ifndef Py_GIL_DISABLED
270+
frame_state = gen->gi_frame_state;
271+
272+
// These asserts are not thread-safe in the free threaded build. Some
273+
// other thread can immediately start executing the generator here.
259274
assert(frame->previous == NULL);
275+
assert(gen->gi_exc_state.previous_item == NULL);
276+
#endif
277+
278+
assert(tstate->exc_info == prev_exc_info);
279+
assert(frame_state != FRAME_EXECUTING);
260280

261281
/* If the generator just returned (as opposed to yielding), signal
262282
* that the generator is exhausted. */
263283
if (result) {
264-
if (FRAME_STATE_SUSPENDED(gen->gi_frame_state)) {
284+
if (FRAME_STATE_SUSPENDED(frame_state)) {
265285
*presult = result;
266286
return PYGEN_NEXT;
267287
}
@@ -277,8 +297,10 @@ gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
277297
!PyErr_ExceptionMatches(PyExc_StopAsyncIteration));
278298
}
279299

300+
#ifndef Py_GIL_DISABLED
280301
assert(gen->gi_exc_state.exc_value == NULL);
281-
assert(gen->gi_frame_state == FRAME_CLEARED);
302+
#endif
303+
assert(frame_state == FRAME_CLEARED);
282304
*presult = result;
283305
return result ? PYGEN_RETURN : PYGEN_ERROR;
284306
}
@@ -918,6 +940,9 @@ make_gen(PyTypeObject *type, PyFunctionObject *func)
918940
gen->gi_weakreflist = NULL;
919941
gen->gi_exc_state.exc_value = NULL;
920942
gen->gi_exc_state.previous_item = NULL;
943+
#ifdef Py_GIL_DISABLED
944+
gen->gi_out_frame_state = NULL;
945+
#endif
921946
assert(func->func_name != NULL);
922947
gen->gi_name = Py_NewRef(func->func_name);
923948
assert(func->func_qualname != NULL);
@@ -1001,6 +1026,9 @@ gen_new_with_qualname(PyTypeObject *type, PyFrameObject *f,
10011026
gen->gi_weakreflist = NULL;
10021027
gen->gi_exc_state.exc_value = NULL;
10031028
gen->gi_exc_state.previous_item = NULL;
1029+
#ifdef Py_GIL_DISABLED
1030+
gen->gi_out_frame_state = NULL;
1031+
#endif
10041032
if (name != NULL)
10051033
gen->gi_name = Py_NewRef(name);
10061034
else

‎Python/bytecodes.c‎

Lines changed: 4 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1206,13 +1206,12 @@ dummy_func(
12061206
assert(frame->owner != FRAME_OWNED_BY_INTERPRETER);
12071207
if ((tstate->interp->eval_frame == NULL) &&
12081208
(Py_TYPE(receiver_o) == &PyGen_Type || Py_TYPE(receiver_o) == &PyCoro_Type) &&
1209-
((PyGenObject *)receiver_o)->gi_frame_state < FRAME_EXECUTING)
1209+
_PyGen_PrepareSend((PyGenObject *)receiver_o))
12101210
{
12111211
PyGenObject *gen = (PyGenObject *)receiver_o;
12121212
_PyInterpreterFrame *gen_frame = &gen->gi_iframe;
12131213
STACK_SHRINK(1);
12141214
_PyFrame_StackPush(gen_frame, v);
1215-
gen->gi_frame_state = FRAME_EXECUTING;
12161215
gen->gi_exc_state.previous_item = tstate->exc_info;
12171216
tstate->exc_info = &gen->gi_exc_state;
12181217
assert(INSTRUCTION_SIZE + oparg <= UINT16_MAX);
@@ -1253,12 +1252,11 @@ dummy_func(
12531252
op(_SEND_GEN_FRAME, (receiver, v -- receiver, gen_frame: _PyInterpreterFrame *)) {
12541253
PyGenObject *gen = (PyGenObject *)PyStackRef_AsPyObjectBorrow(receiver);
12551254
DEOPT_IF(Py_TYPE(gen) != &PyGen_Type && Py_TYPE(gen) != &PyCoro_Type);
1256-
DEOPT_IF(gen->gi_frame_state >= FRAME_EXECUTING);
1255+
DEOPT_IF(!_PyGen_PrepareSend(gen));
12571256
STAT_INC(SEND, hit);
12581257
gen_frame = &gen->gi_iframe;
12591258
_PyFrame_StackPush(gen_frame, v);
12601259
DEAD(v);
1261-
gen->gi_frame_state = FRAME_EXECUTING;
12621260
gen->gi_exc_state.previous_item = tstate->exc_info;
12631261
tstate->exc_info = &gen->gi_exc_state;
12641262
assert(INSTRUCTION_SIZE + oparg <= UINT16_MAX);
@@ -1281,7 +1279,6 @@ dummy_func(
12811279
PyGenObject *gen = _PyGen_GetGeneratorFromFrame(frame);
12821280
assert(FRAME_SUSPENDED_YIELD_FROM == FRAME_SUSPENDED + 1);
12831281
assert(oparg == 0 || oparg == 1);
1284-
gen->gi_frame_state = FRAME_SUSPENDED + oparg;
12851282
_PyStackRef temp = retval;
12861283
DEAD(retval);
12871284
SAVE_STACK();
@@ -1291,6 +1288,7 @@ dummy_func(
12911288
_PyInterpreterFrame *gen_frame = frame;
12921289
frame = tstate->current_frame = frame->previous;
12931290
gen_frame->previous = NULL;
1291+
_PyGen_SetFrameState(gen, FRAME_SUSPENDED + oparg);
12941292
/* We don't know which of these is relevant here, so keep them equal */
12951293
assert(INLINE_CACHE_ENTRIES_SEND == INLINE_CACHE_ENTRIES_FOR_ITER);
12961294
#if TIER_ONE
@@ -3321,18 +3319,10 @@ dummy_func(
33213319
op(_FOR_ITER_GEN_FRAME, (iter -- iter, gen_frame: _PyInterpreterFrame*)) {
33223320
PyGenObject *gen = (PyGenObject *)PyStackRef_AsPyObjectBorrow(iter);
33233321
DEOPT_IF(Py_TYPE(gen) != &PyGen_Type);
3324-
#ifdef Py_GIL_DISABLED
3325-
// Since generators can't be used by multiple threads anyway we
3326-
// don't need to deopt here, but this lets us work on making
3327-
// generators thread-safe without necessarily having to
3328-
// specialize them thread-safely as well.
3329-
DEOPT_IF(!_PyObject_IsUniquelyReferenced((PyObject *)gen));
3330-
#endif
3331-
DEOPT_IF(gen->gi_frame_state >= FRAME_EXECUTING);
3322+
DEOPT_IF(!_PyGen_PrepareSend(gen));
33323323
STAT_INC(FOR_ITER, hit);
33333324
gen_frame = &gen->gi_iframe;
33343325
_PyFrame_StackPush(gen_frame, PyStackRef_None);
3335-
gen->gi_frame_state = FRAME_EXECUTING;
33363326
gen->gi_exc_state.previous_item = tstate->exc_info;
33373327
tstate->exc_info = &gen->gi_exc_state;
33383328
gen_frame->previous = frame;

‎Python/ceval.c‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1735,13 +1735,15 @@ clear_gen_frame(PyThreadState *tstate, _PyInterpreterFrame * frame)
17351735
{
17361736
assert(frame->owner == FRAME_OWNED_BY_GENERATOR);
17371737
PyGenObject *gen = _PyGen_GetGeneratorFromFrame(frame);
1738-
gen->gi_frame_state = FRAME_CLEARED;
17391738
assert(tstate->exc_info == &gen->gi_exc_state);
17401739
tstate->exc_info = gen->gi_exc_state.previous_item;
1740+
PyObject *exc_value = gen->gi_exc_state.exc_value;
17411741
gen->gi_exc_state.previous_item = NULL;
1742+
gen->gi_exc_state.exc_value = NULL;
17421743
assert(frame->frame_obj == NULL || frame->frame_obj->f_frame == frame);
1744+
_PyGen_SetFrameState(gen, FRAME_CLEARED);
17431745
_PyFrame_ClearExceptCode(frame);
1744-
_PyErr_ClearExcState(&gen->gi_exc_state);
1746+
Py_XDECREF(exc_value);
17451747
frame->previous = NULL;
17461748
}
17471749

‎Python/ceval_macros.h‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -363,6 +363,17 @@ static inline void _Py_LeaveRecursiveCallPy(PyThreadState *tstate) {
363363
tstate->py_recursion_remaining++;
364364
}
365365

366+
static inline int
367+
_PyGen_PrepareSend(PyGenObject *gen)
368+
{
369+
assert(PyGen_CheckExact(gen) || PyCoro_CheckExact(gen));
370+
int8_t frame_state = FT_ATOMIC_LOAD_INT8_RELAXED(gen->gi_frame_state);
371+
if (frame_state >= FRAME_EXECUTING) {
372+
return 0;
373+
}
374+
return _PyGen_TransitionFrameState(gen, &frame_state, FRAME_EXECUTING);
375+
}
376+
366377
/* Implementation of "macros" that modify the instruction pointer,
367378
* stack pointer, or frame pointer.
368379
* These need to treated differently by tier 1 and 2.

‎Python/executor_cases.c.h‎

Lines changed: 3 additions & 15 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)