Skip to content

Commit 975cfa7

Browse files
ambvambv
committed
pythongh-91048: Also clear and set ts->asyncio_running_task with eager tasks
This was missing from pythongh-124640. It's already covered by the new test_asyncio/test_free_threading.py in combination with the runtime assertion in set_ts_asyncio_running_task.
1 parent 67d804b commit 975cfa7

1 file changed

Lines changed: 60 additions & 49 deletions

File tree

‎Modules/_asynciomodule.c‎

Lines changed: 60 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -2063,6 +2063,8 @@ static int task_call_step_soon(asyncio_state *state, TaskObj *, PyObject *);
20632063
static PyObject * task_wakeup(TaskObj *, PyObject *);
20642064
static PyObject * task_step(asyncio_state *, TaskObj *, PyObject *);
20652065
static int task_eager_start(asyncio_state *state, TaskObj *task);
2066+
static inline void clear_ts_asyncio_running_task(PyObject *loop);
2067+
static inline void set_ts_asyncio_running_task(PyObject *loop, PyObject *task);
20662068

20672069
/* ----- Task._step wrapper */
20682070

@@ -2236,47 +2238,7 @@ enter_task(asyncio_state *state, PyObject *loop, PyObject *task)
22362238

22372239
assert(task == item);
22382240
Py_CLEAR(item);
2239-
2240-
// This block is needed to enable `asyncio.capture_call_graph()` API.
2241-
// We want to be enable debuggers and profilers to be able to quickly
2242-
// introspect the asyncio running state from another process.
2243-
// When we do that, we need to essentially traverse the address space
2244-
// of a Python process and understand what every Python thread in it is
2245-
// currently doing, mainly:
2246-
//
2247-
// * current frame
2248-
// * current asyncio task
2249-
//
2250-
// A naive solution would be to require profilers and debuggers to
2251-
// find the current task in the "_asynciomodule" module state, but
2252-
// unfortunately that would require a lot of complicated remote
2253-
// memory reads and logic, as Python's dict is a notoriously complex
2254-
// and ever-changing data structure.
2255-
//
2256-
// So the easier solution is to put a strong reference to the currently
2257-
// running `asyncio.Task` on the interpreter thread state (we already
2258-
// have some asyncio state there.)
2259-
_PyThreadStateImpl *ts = (_PyThreadStateImpl *)_PyThreadState_GET();
2260-
if (ts->asyncio_running_loop == loop) {
2261-
// Protect from a situation when someone calls this method
2262-
// from another thread. This shouldn't ever happen though,
2263-
// as `enter_task` and `leave_task` can either be called by:
2264-
//
2265-
// - `asyncio.Task` itself, in `Task.__step()`. That method
2266-
// can only be called by the event loop itself.
2267-
//
2268-
// - third-party Task "from scratch" implementations, that
2269-
// our `capture_call_graph` API doesn't support anyway.
2270-
//
2271-
// That said, we still want to make sure we don't end up in
2272-
// a broken state, so we check that we're in the correct thread
2273-
// by comparing the *loop* argument to the event loop running
2274-
// in the current thread. If they match we know we're in the
2275-
// right thread, as asyncio event loops don't change threads.
2276-
assert(ts->asyncio_running_task == NULL);
2277-
ts->asyncio_running_task = Py_NewRef(task);
2278-
}
2279-
2241+
set_ts_asyncio_running_task(loop, task);
22802242
return 0;
22812243
}
22822244

@@ -2308,14 +2270,7 @@ leave_task(asyncio_state *state, PyObject *loop, PyObject *task)
23082270
// task was not found
23092271
return err_leave_task(Py_None, task);
23102272
}
2311-
2312-
// See the comment in `enter_task` for the explanation of why
2313-
// the following is needed.
2314-
_PyThreadStateImpl *ts = (_PyThreadStateImpl *)_PyThreadState_GET();
2315-
if (ts->asyncio_running_loop == NULL || ts->asyncio_running_loop == loop) {
2316-
Py_CLEAR(ts->asyncio_running_task);
2317-
}
2318-
2273+
clear_ts_asyncio_running_task(loop);
23192274
return res;
23202275
}
23212276

@@ -2342,6 +2297,7 @@ swap_current_task(asyncio_state *state, PyObject *loop, PyObject *task)
23422297
{
23432298
PyObject *prev_task;
23442299

2300+
clear_ts_asyncio_running_task(loop);
23452301
if (task == Py_None) {
23462302
if (PyDict_Pop(state->current_tasks, loop, &prev_task) < 0) {
23472303
return NULL;
@@ -2361,9 +2317,64 @@ swap_current_task(asyncio_state *state, PyObject *loop, PyObject *task)
23612317
Py_BEGIN_CRITICAL_SECTION(current_tasks);
23622318
prev_task = swap_current_task_lock_held(current_tasks, loop, hash, task);
23632319
Py_END_CRITICAL_SECTION();
2320+
set_ts_asyncio_running_task(loop, task);
23642321
return prev_task;
23652322
}
23662323

2324+
static inline void
2325+
set_ts_asyncio_running_task(PyObject *loop, PyObject *task)
2326+
{
2327+
// This is needed to enable `asyncio.capture_call_graph()` API.
2328+
// We want to be enable debuggers and profilers to be able to quickly
2329+
// introspect the asyncio running state from another process.
2330+
// When we do that, we need to essentially traverse the address space
2331+
// of a Python process and understand what every Python thread in it is
2332+
// currently doing, mainly:
2333+
//
2334+
// * current frame
2335+
// * current asyncio task
2336+
//
2337+
// A naive solution would be to require profilers and debuggers to
2338+
// find the current task in the "_asynciomodule" module state, but
2339+
// unfortunately that would require a lot of complicated remote
2340+
// memory reads and logic, as Python's dict is a notoriously complex
2341+
// and ever-changing data structure.
2342+
//
2343+
// So the easier solution is to put a strong reference to the currently
2344+
// running `asyncio.Task` on the interpreter thread state (we already
2345+
// have some asyncio state there.)
2346+
_PyThreadStateImpl *ts = (_PyThreadStateImpl *)_PyThreadState_GET();
2347+
if (ts->asyncio_running_loop == loop) {
2348+
// Protect from a situation when someone calls this method
2349+
// from another thread. This shouldn't ever happen though,
2350+
// as `enter_task` and `leave_task` can either be called by:
2351+
//
2352+
// - `asyncio.Task` itself, in `Task.__step()`. That method
2353+
// can only be called by the event loop itself.
2354+
//
2355+
// - third-party Task "from scratch" implementations, that
2356+
// our `capture_call_graph` API doesn't support anyway.
2357+
//
2358+
// That said, we still want to make sure we don't end up in
2359+
// a broken state, so we check that we're in the correct thread
2360+
// by comparing the *loop* argument to the event loop running
2361+
// in the current thread. If they match we know we're in the
2362+
// right thread, as asyncio event loops don't change threads.
2363+
assert(ts->asyncio_running_task == NULL);
2364+
ts->asyncio_running_task = Py_NewRef(task);
2365+
}
2366+
}
2367+
2368+
static inline void
2369+
clear_ts_asyncio_running_task(PyObject *loop)
2370+
{
2371+
// See comment in set_ts_asyncio_running_task() for details.
2372+
_PyThreadStateImpl *ts = (_PyThreadStateImpl *)_PyThreadState_GET();
2373+
if (ts->asyncio_running_loop == NULL || ts->asyncio_running_loop == loop) {
2374+
Py_CLEAR(ts->asyncio_running_task);
2375+
}
2376+
}
2377+
23672378
/* ----- Task */
23682379

23692380
/*[clinic input]

0 commit comments

Comments
 (0)