Skip to content

Commit 317a7e2

Browse files
colesburycolesbury
committed
pythongh-120974: Make asyncio swap_current_task safe in free-threaded build
Use a critical section around the modifications to `current_tasks`.
1 parent 7c29218 commit 317a7e2

3 files changed

Lines changed: 44 additions & 31 deletions

File tree

‎Include/internal/pycore_dict.h‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,8 +108,12 @@ PyAPI_FUNC(PyObject *)_PyDict_LoadGlobal(PyDictObject *, PyDictObject *, PyObjec
108108
/* Consumes references to key and value */
109109
PyAPI_FUNC(int) _PyDict_SetItem_Take2(PyDictObject *op, PyObject *key, PyObject *value);
110110
extern int _PyDict_SetItem_LockHeld(PyDictObject *dict, PyObject *name, PyObject *value);
111+
// Export for '_asyncio' shared extension
112+
PyAPI_FUNC(int) _PyDict_SetItem_KnownHash_LockHeld(PyDictObject *mp, PyObject *key,
113+
PyObject *value, Py_hash_t hash);
114+
// Export for '_asyncio' shared extension
115+
PyAPI_FUNC(int) _PyDict_GetItemRef_KnownHash(PyDictObject *op, PyObject *key, Py_hash_t hash, PyObject **result);
111116
extern int _PyDict_GetItemRef_Unicode_LockHeld(PyDictObject *op, PyObject *key, PyObject **result);
112-
extern int _PyDict_GetItemRef_KnownHash(PyDictObject *op, PyObject *key, Py_hash_t hash, PyObject **result);
113117
extern int _PyObjectDict_SetItem(PyTypeObject *tp, PyObject *obj, PyObject **dictptr, PyObject *name, PyObject *value);
114118

115119
extern int _PyDict_Pop_KnownHash(

‎Modules/_asynciomodule.c‎

Lines changed: 24 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2026,6 +2026,25 @@ leave_task(asyncio_state *state, PyObject *loop, PyObject *task)
20262026
return res;
20272027
}
20282028

2029+
static PyObject *
2030+
swap_current_task_lock_held(PyDictObject *current_tasks, PyObject *loop,
2031+
Py_hash_t hash, PyObject *task)
2032+
{
2033+
PyObject *prev_task;
2034+
if (_PyDict_GetItemRef_KnownHash(current_tasks, loop, hash, &prev_task) < 0) {
2035+
return NULL;
2036+
}
2037+
if (prev_task == NULL) {
2038+
prev_task = Py_None;
2039+
}
2040+
Py_INCREF(prev_task);
2041+
if (_PyDict_SetItem_KnownHash_LockHeld(current_tasks, loop, task, hash) < 0) {
2042+
Py_DECREF(prev_task);
2043+
return NULL;
2044+
}
2045+
return prev_task;
2046+
}
2047+
20292048
static PyObject *
20302049
swap_current_task(asyncio_state *state, PyObject *loop, PyObject *task)
20312050
{
@@ -2041,24 +2060,15 @@ swap_current_task(asyncio_state *state, PyObject *loop, PyObject *task)
20412060
return prev_task;
20422061
}
20432062

2044-
Py_hash_t hash;
2045-
hash = PyObject_Hash(loop);
2063+
Py_hash_t hash = PyObject_Hash(loop);
20462064
if (hash == -1) {
20472065
return NULL;
20482066
}
2049-
prev_task = _PyDict_GetItem_KnownHash(state->current_tasks, loop, hash);
2050-
if (prev_task == NULL) {
2051-
if (PyErr_Occurred()) {
2052-
return NULL;
2053-
}
2054-
prev_task = Py_None;
2055-
}
2056-
Py_INCREF(prev_task);
2057-
if (_PyDict_SetItem_KnownHash(state->current_tasks, loop, task, hash) == -1) {
2058-
Py_DECREF(prev_task);
2059-
return NULL;
2060-
}
20612067

2068+
PyDictObject *current_tasks = (PyDictObject *)state->current_tasks;
2069+
Py_BEGIN_CRITICAL_SECTION(current_tasks);
2070+
prev_task = swap_current_task_lock_held(current_tasks, loop, hash, task);
2071+
Py_END_CRITICAL_SECTION();
20622072
return prev_task;
20632073
}
20642074

‎Objects/dictobject.c‎

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2460,33 +2460,32 @@ setitem_lock_held(PyDictObject *mp, PyObject *key, PyObject *value)
24602460

24612461

24622462
int
2463-
_PyDict_SetItem_KnownHash(PyObject *op, PyObject *key, PyObject *value,
2464-
Py_hash_t hash)
2463+
_PyDict_SetItem_KnownHash_LockHeld(PyDictObject *mp, PyObject *key, PyObject *value,
2464+
Py_hash_t hash)
24652465
{
2466-
PyDictObject *mp;
2466+
PyInterpreterState *interp = _PyInterpreterState_GET();
2467+
if (mp->ma_keys == Py_EMPTY_KEYS) {
2468+
return insert_to_emptydict(interp, mp, Py_NewRef(key), hash, Py_NewRef(value));
2469+
}
2470+
/* insertdict() handles any resizing that might be necessary */
2471+
return insertdict(interp, mp, Py_NewRef(key), hash, Py_NewRef(value));
2472+
}
24672473

2474+
int
2475+
_PyDict_SetItem_KnownHash(PyObject *op, PyObject *key, PyObject *value,
2476+
Py_hash_t hash)
2477+
{
24682478
if (!PyDict_Check(op)) {
24692479
PyErr_BadInternalCall();
24702480
return -1;
24712481
}
24722482
assert(key);
24732483
assert(value);
24742484
assert(hash != -1);
2475-
mp = (PyDictObject *)op;
24762485

24772486
int res;
2478-
PyInterpreterState *interp = _PyInterpreterState_GET();
2479-
2480-
Py_BEGIN_CRITICAL_SECTION(mp);
2481-
2482-
if (mp->ma_keys == Py_EMPTY_KEYS) {
2483-
res = insert_to_emptydict(interp, mp, Py_NewRef(key), hash, Py_NewRef(value));
2484-
}
2485-
else {
2486-
/* insertdict() handles any resizing that might be necessary */
2487-
res = insertdict(interp, mp, Py_NewRef(key), hash, Py_NewRef(value));
2488-
}
2489-
2487+
Py_BEGIN_CRITICAL_SECTION(op);
2488+
res = _PyDict_SetItem_KnownHash_LockHeld((PyDictObject *)op, key, value, hash);
24902489
Py_END_CRITICAL_SECTION();
24912490
return res;
24922491
}

0 commit comments

Comments
 (0)