Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
57 changes: 57 additions & 0 deletions Lib/test/test_free_threading/test_context.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
import contextvars
import unittest
from threading import Event, Thread

from test.support import threading_helper


@threading_helper.requires_working_threading()
class TestContext(unittest.TestCase):
def test_racing_read_write(self):
# gh-154535: reading a Context object from one thread while another
# thread sets variables in it used to crash. The readers looked at
# Context.ctx_vars without owning a reference to it, so the writer
# could deallocate the mapping while a reader was walking it.
ctx = contextvars.Context()
cvars = [contextvars.ContextVar(f"cvar{i}") for i in range(64)]
done = Event()
errors = []

def writer():
def body():
i = 0
while not done.is_set():
cvars[i % len(cvars)].set(i)
i += 1
try:
ctx.run(body)
except BaseException as e:
errors.append(e)

def reader():
try:
for _ in range(200):
ctx.copy()
len(ctx)
list(ctx)
list(ctx.items())
list(ctx.keys())
list(ctx.values())
cvars[0] in ctx
ctx.get(cvars[0])
ctx == ctx
except BaseException as e:
errors.append(e)
finally:
done.set()

threads = [Thread(target=writer)]
threads += [Thread(target=reader) for _ in range(4)]
with threading_helper.start_threads(threads, done.set):
pass

self.assertEqual(errors, [], msg=f"unexpected errors: {errors}")


if __name__ == "__main__":
unittest.main()
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Avoid a data-race in free-threaded builds when reading and writing context
variables from different threads.
113 changes: 92 additions & 21 deletions Python/context.c
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
#include "Python.h"
#include "pycore_call.h" // _PyObject_VectorcallTstate()
#include "pycore_context.h"
#include "pycore_critical_section.h" // Py_BEGIN_CRITICAL_SECTION()
#include "pycore_freelist.h" // _Py_FREELIST_FREE(), _Py_FREELIST_POP()
#include "pycore_gc.h" // _PyObject_GC_MAY_BE_TRACKED()
#include "pycore_hamt.h"
#include "pycore_initconfig.h" // _PyStatus_OK()
#include "pycore_object.h"
#include "pycore_pyatomic_ft_wrappers.h" // FT_ATOMIC_LOAD_INT_RELAXED()
#include "pycore_pyerrors.h"
#include "pycore_pystate.h" // _PyThreadState_GET()

Expand Down Expand Up @@ -64,6 +66,41 @@ contextvar_set(PyContextVar *var, PyObject *val);
static int
contextvar_del(PyContextVar *var);

static inline PyHamtObject *
context_get_vars(PyContext *ctx)
{
PyHamtObject *vars;
Py_BEGIN_CRITICAL_SECTION(ctx);
vars = ctx->ctx_vars;
assert(vars != NULL);
Py_INCREF(vars);
Py_END_CRITICAL_SECTION();
return vars;
}

static inline PyHamtObject *
context_get_current_vars(PyContext *ctx)
{
// ctx_vars written only by the owning thread, and read by other threads
// only under the context's lock, a plain (non-atomic) load is okay
PyHamtObject *vars = ctx->ctx_vars;
assert(vars != NULL);
return vars;
}

// Note: steals a reference to new_vars and must only be called by the thread
// that has `ctx` as its current context.
static inline void
context_set_vars(PyContext *ctx, PyHamtObject *new_vars)
{
PyHamtObject *old_vars;
Py_BEGIN_CRITICAL_SECTION(ctx);
old_vars = ctx->ctx_vars;
ctx->ctx_vars = new_vars;
Py_END_CRITICAL_SECTION();
Py_XDECREF(old_vars);
}


PyObject *
_PyContext_NewHamtForTests(void)
Expand All @@ -84,7 +121,10 @@ PyContext_Copy(PyObject * octx)
{
ENSURE_Context(octx, NULL)
PyContext *ctx = (PyContext *)octx;
return (PyObject *)context_new_from_vars(ctx->ctx_vars);
PyHamtObject *vars = context_get_vars(ctx);
PyObject *res = (PyObject *)context_new_from_vars(vars);
Py_DECREF(vars);
return res;
}


Expand All @@ -96,7 +136,7 @@ PyContext_CopyCurrent(void)
return NULL;
}

return (PyObject *)context_new_from_vars(ctx->ctx_vars);
return (PyObject *)context_new_from_vars(context_get_current_vars(ctx));
}

static const char *
Expand Down Expand Up @@ -298,7 +338,7 @@ PyContextVar_Get(PyObject *ovar, PyObject *def, PyObject **val)
#endif

assert(PyContext_CheckExact(ts->context));
PyHamtObject *vars = ((PyContext *)ts->context)->ctx_vars;
PyHamtObject *vars = context_get_current_vars((PyContext *)ts->context);

PyObject *found = NULL;
int res = _PyHamt_Find(vars, (PyObject*)var, &found);
Expand Down Expand Up @@ -354,7 +394,8 @@ PyContextVar_Set(PyObject *ovar, PyObject *val)
}

PyObject *old_val = NULL;
int found = _PyHamt_Find(ctx->ctx_vars, (PyObject *)var, &old_val);
int found = _PyHamt_Find(context_get_current_vars(ctx), (PyObject *)var,
&old_val);
if (found < 0) {
return NULL;
}
Expand Down Expand Up @@ -552,7 +593,10 @@ static PyObject *
context_tp_iter(PyObject *op)
{
PyContext *self = _PyContext_CAST(op);
return _PyHamt_NewIterKeys(self->ctx_vars);
PyHamtObject *vars = context_get_vars(self);
PyObject *res = _PyHamt_NewIterKeys(vars);
Py_DECREF(vars);
return res;
}

static PyObject *
Expand All @@ -564,8 +608,11 @@ context_tp_richcompare(PyObject *v, PyObject *w, int op)
Py_RETURN_NOTIMPLEMENTED;
}

int res = _PyHamt_Eq(
((PyContext *)v)->ctx_vars, ((PyContext *)w)->ctx_vars);
PyHamtObject *v_vars = context_get_vars((PyContext *)v);
PyHamtObject *w_vars = context_get_vars((PyContext *)w);
int res = _PyHamt_Eq(v_vars, w_vars);
Py_DECREF(v_vars);
Py_DECREF(w_vars);
if (res < 0) {
return NULL;
}
Expand All @@ -586,7 +633,10 @@ static Py_ssize_t
context_tp_len(PyObject *op)
{
PyContext *self = _PyContext_CAST(op);
return _PyHamt_Len(self->ctx_vars);
PyHamtObject *vars = context_get_vars(self);
Py_ssize_t res = _PyHamt_Len(vars);
Py_DECREF(vars);
return res;
}

static PyObject *
Expand All @@ -597,15 +647,18 @@ context_tp_subscript(PyObject *op, PyObject *key)
}
PyObject *val = NULL;
PyContext *self = _PyContext_CAST(op);
int found = _PyHamt_Find(self->ctx_vars, key, &val);
PyHamtObject *vars = context_get_vars(self);
int found = _PyHamt_Find(vars, key, &val);
Py_XINCREF(val);
Py_DECREF(vars);
if (found < 0) {
return NULL;
}
if (found == 0) {
PyErr_SetObject(PyExc_KeyError, key);
return NULL;
}
return Py_NewRef(val);
return val;
}

static int
Expand All @@ -616,7 +669,10 @@ context_tp_contains(PyObject *op, PyObject *key)
}
PyObject *val = NULL;
PyContext *self = _PyContext_CAST(op);
return _PyHamt_Find(self->ctx_vars, key, &val);
PyHamtObject *vars = context_get_vars(self);
int res = _PyHamt_Find(vars, key, &val);
Py_DECREF(vars);
return res;
}


Expand All @@ -643,14 +699,17 @@ _contextvars_Context_get_impl(PyContext *self, PyObject *key,
}

PyObject *val = NULL;
int found = _PyHamt_Find(self->ctx_vars, key, &val);
PyHamtObject *vars = context_get_vars(self);
int found = _PyHamt_Find(vars, key, &val);
Py_XINCREF(val);
Py_DECREF(vars);
if (found < 0) {
return NULL;
}
if (found == 0) {
return Py_NewRef(default_value);
}
return Py_NewRef(val);
return val;
}


Expand All @@ -666,7 +725,10 @@ static PyObject *
_contextvars_Context_items_impl(PyContext *self)
/*[clinic end generated code: output=fa1655c8a08502af input=00db64ae379f9f42]*/
{
return _PyHamt_NewIterItems(self->ctx_vars);
PyHamtObject *vars = context_get_vars(self);
PyObject *res = _PyHamt_NewIterItems(vars);
Py_DECREF(vars);
return res;
}


Expand All @@ -680,7 +742,10 @@ static PyObject *
_contextvars_Context_keys_impl(PyContext *self)
/*[clinic end generated code: output=177227c6b63ec0e2 input=114b53aebca3449c]*/
{
return _PyHamt_NewIterKeys(self->ctx_vars);
PyHamtObject *vars = context_get_vars(self);
PyObject *res = _PyHamt_NewIterKeys(vars);
Py_DECREF(vars);
return res;
}


Expand All @@ -694,7 +759,10 @@ static PyObject *
_contextvars_Context_values_impl(PyContext *self)
/*[clinic end generated code: output=d286dabfc8db6dde input=ce8075d04a6ea526]*/
{
return _PyHamt_NewIterValues(self->ctx_vars);
PyHamtObject *vars = context_get_vars(self);
PyObject *res = _PyHamt_NewIterValues(vars);
Py_DECREF(vars);
return res;
}


Expand All @@ -708,7 +776,10 @@ static PyObject *
_contextvars_Context_copy_impl(PyContext *self)
/*[clinic end generated code: output=30ba8896c4707a15 input=ebafdbdd9c72d592]*/
{
return (PyObject *)context_new_from_vars(self->ctx_vars);
PyHamtObject *vars = context_get_vars(self);
PyObject *res = (PyObject *)context_new_from_vars(vars);
Py_DECREF(vars);
return res;
}


Expand Down Expand Up @@ -796,12 +867,12 @@ contextvar_set(PyContextVar *var, PyObject *val)
}

PyHamtObject *new_vars = _PyHamt_Assoc(
ctx->ctx_vars, (PyObject *)var, val);
context_get_current_vars(ctx), (PyObject *)var, val);
if (new_vars == NULL) {
return -1;
}

Py_SETREF(ctx->ctx_vars, new_vars);
context_set_vars(ctx, new_vars);

#ifndef Py_GIL_DISABLED
var->var_cached = val; /* borrow */
Expand All @@ -823,7 +894,7 @@ contextvar_del(PyContextVar *var)
return -1;
}

PyHamtObject *vars = ctx->ctx_vars;
PyHamtObject *vars = context_get_current_vars(ctx);
PyHamtObject *new_vars = _PyHamt_Without(vars, (PyObject *)var);
if (new_vars == NULL) {
return -1;
Expand All @@ -835,7 +906,7 @@ contextvar_del(PyContextVar *var)
return -1;
}

Py_SETREF(ctx->ctx_vars, new_vars);
context_set_vars(ctx, new_vars);
return 0;
}

Expand Down
Loading