gh-119333: Add C api to have contextvar enter/exit callbacks (#119335)

Co-authored-by: Erlend E. Aasland <erlend.aasland@protonmail.com>
This commit is contained in:
Jason Fried 2024-09-23 20:40:17 -07:00 committed by GitHub
parent ad7c778546
commit d87482bc4e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 402 additions and 0 deletions

View file

@ -99,6 +99,80 @@ PyContext_CopyCurrent(void)
return (PyObject *)context_new_from_vars(ctx->ctx_vars);
}
static const char *
context_event_name(PyContextEvent event) {
switch (event) {
case Py_CONTEXT_EVENT_ENTER:
return "Py_CONTEXT_EVENT_ENTER";
case Py_CONTEXT_EVENT_EXIT:
return "Py_CONTEXT_EVENT_EXIT";
default:
return "?";
}
Py_UNREACHABLE();
}
static void notify_context_watchers(PyContextEvent event, PyContext *ctx)
{
assert(Py_REFCNT(ctx) > 0);
PyInterpreterState *interp = _PyInterpreterState_GET();
assert(interp->_initialized);
uint8_t bits = interp->active_context_watchers;
int i = 0;
while (bits) {
assert(i < CONTEXT_MAX_WATCHERS);
if (bits & 1) {
PyContext_WatchCallback cb = interp->context_watchers[i];
assert(cb != NULL);
if (cb(event, ctx) < 0) {
PyErr_FormatUnraisable(
"Exception ignored in %s watcher callback for %R",
context_event_name(event), ctx);
}
}
i++;
bits >>= 1;
}
}
int
PyContext_AddWatcher(PyContext_WatchCallback callback)
{
PyInterpreterState *interp = _PyInterpreterState_GET();
assert(interp->_initialized);
for (int i = 0; i < CONTEXT_MAX_WATCHERS; i++) {
if (!interp->context_watchers[i]) {
interp->context_watchers[i] = callback;
interp->active_context_watchers |= (1 << i);
return i;
}
}
PyErr_SetString(PyExc_RuntimeError, "no more context watcher IDs available");
return -1;
}
int
PyContext_ClearWatcher(int watcher_id)
{
PyInterpreterState *interp = _PyInterpreterState_GET();
assert(interp->_initialized);
if (watcher_id < 0 || watcher_id >= CONTEXT_MAX_WATCHERS) {
PyErr_Format(PyExc_ValueError, "Invalid context watcher ID %d", watcher_id);
return -1;
}
if (!interp->context_watchers[watcher_id]) {
PyErr_Format(PyExc_ValueError, "No context watcher set for ID %d", watcher_id);
return -1;
}
interp->context_watchers[watcher_id] = NULL;
interp->active_context_watchers &= ~(1 << watcher_id);
return 0;
}
static int
_PyContext_Enter(PyThreadState *ts, PyObject *octx)
@ -118,6 +192,7 @@ _PyContext_Enter(PyThreadState *ts, PyObject *octx)
ts->context = Py_NewRef(ctx);
ts->context_ver++;
notify_context_watchers(Py_CONTEXT_EVENT_ENTER, ctx);
return 0;
}
@ -151,6 +226,7 @@ _PyContext_Exit(PyThreadState *ts, PyObject *octx)
return -1;
}
notify_context_watchers(Py_CONTEXT_EVENT_EXIT, ctx);
Py_SETREF(ts->context, (PyObject *)ctx->ctx_prev);
ts->context_ver++;