Skip to content

Commit 84d1261

Browse files
committed
gh-155519: fix data-race for Context.ctx_vars
The HAMT held by a context is immutable, but the pointer to it is not: contextvar_set() and contextvar_del() install a new HAMT and drop the reference to the previous one. Use critical sections to ensure ordering in the free-threaded build.
1 parent 219768f commit 84d1261

3 files changed

Lines changed: 170 additions & 21 deletions

File tree

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
import contextvars
2+
import unittest
3+
from threading import Event, Thread
4+
5+
from test.support import threading_helper
6+
7+
8+
@threading_helper.requires_working_threading()
9+
class TestContext(unittest.TestCase):
10+
def test_racing_read_write(self):
11+
# gh-154535: reading a Context object from one thread while another
12+
# thread sets variables in it used to crash. The readers looked at
13+
# Context.ctx_vars without owning a reference to it, so the writer
14+
# could deallocate the mapping while a reader was walking it.
15+
ctx = contextvars.Context()
16+
cvars = [contextvars.ContextVar(f"cvar{i}") for i in range(64)]
17+
done = Event()
18+
errors = []
19+
20+
def writer():
21+
def body():
22+
i = 0
23+
while not done.is_set():
24+
cvars[i % len(cvars)].set(i)
25+
i += 1
26+
try:
27+
ctx.run(body)
28+
except BaseException as e:
29+
errors.append(e)
30+
31+
def reader():
32+
try:
33+
for _ in range(200):
34+
ctx.copy()
35+
len(ctx)
36+
list(ctx)
37+
list(ctx.items())
38+
list(ctx.keys())
39+
list(ctx.values())
40+
cvars[0] in ctx
41+
ctx.get(cvars[0])
42+
ctx == ctx
43+
except BaseException as e:
44+
errors.append(e)
45+
finally:
46+
done.set()
47+
48+
threads = [Thread(target=writer)]
49+
threads += [Thread(target=reader) for _ in range(4)]
50+
with threading_helper.start_threads(threads, done.set):
51+
pass
52+
53+
self.assertEqual(errors, [], msg=f"unexpected errors: {errors}")
54+
55+
56+
if __name__ == "__main__":
57+
unittest.main()
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Avoid a data-race in free-threaded builds when reading and writing context
2+
variables from different threads.

Python/context.c

Lines changed: 111 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
#include "Python.h"
22
#include "pycore_call.h" // _PyObject_VectorcallTstate()
33
#include "pycore_context.h"
4+
#include "pycore_critical_section.h" // Py_BEGIN_CRITICAL_SECTION()
45
#include "pycore_freelist.h" // _Py_FREELIST_FREE(), _Py_FREELIST_POP()
56
#include "pycore_gc.h" // _PyObject_GC_MAY_BE_TRACKED()
67
#include "pycore_hamt.h"
78
#include "pycore_initconfig.h" // _PyStatus_OK()
89
#include "pycore_object.h"
10+
#include "pycore_pyatomic_ft_wrappers.h" // FT_ATOMIC_LOAD_INT_RELAXED()
911
#include "pycore_pyerrors.h"
1012
#include "pycore_pystate.h" // _PyThreadState_GET()
1113

@@ -65,6 +67,58 @@ static int
6567
contextvar_del(PyContextVar *var);
6668

6769

70+
/* The HAMT held by a context is immutable, but the pointer to it is not:
71+
contextvar_set() and contextvar_del() install a new HAMT and drop the
72+
reference to the previous one. Only the thread that entered a context can
73+
do that (a context cannot be entered by two threads at once), but any
74+
thread can read a context object at any time -- ctx.copy(), len(ctx),
75+
ctx.items(), etc. Such a reader has to acquire its own reference to the
76+
HAMT under the context's lock; reading ctx_vars unlocked lets the writer
77+
deallocate the HAMT while the reader is walking it.
78+
79+
With that discipline -- ctx_vars written only by the owning thread, and
80+
read by other threads only under the context's lock -- there is no
81+
unsynchronized concurrent access, so plain (non-atomic) loads and stores
82+
are used. */
83+
84+
static inline PyHamtObject *
85+
context_get_vars(PyContext *ctx)
86+
{
87+
PyHamtObject *vars;
88+
Py_BEGIN_CRITICAL_SECTION(ctx);
89+
vars = ctx->ctx_vars;
90+
assert(vars != NULL);
91+
Py_INCREF(vars);
92+
Py_END_CRITICAL_SECTION();
93+
return vars;
94+
}
95+
96+
/* Same, but for the context that is current on this thread. No other thread
97+
can have it as its current context, so this thread is the only one that can
98+
replace its HAMT: neither the lock nor a new reference is needed here.
99+
Returns a borrowed reference. */
100+
static inline PyHamtObject *
101+
context_get_current_vars(PyContext *ctx)
102+
{
103+
PyHamtObject *vars = ctx->ctx_vars;
104+
assert(vars != NULL);
105+
return vars;
106+
}
107+
108+
/* Install a new HAMT in a context. Steals a reference to new_vars. Must
109+
only be called by the thread that has `ctx` as its current context. */
110+
static inline void
111+
context_set_vars(PyContext *ctx, PyHamtObject *new_vars)
112+
{
113+
PyHamtObject *old_vars;
114+
Py_BEGIN_CRITICAL_SECTION(ctx);
115+
old_vars = ctx->ctx_vars;
116+
ctx->ctx_vars = new_vars;
117+
Py_END_CRITICAL_SECTION();
118+
Py_XDECREF(old_vars);
119+
}
120+
121+
68122
PyObject *
69123
_PyContext_NewHamtForTests(void)
70124
{
@@ -84,7 +138,10 @@ PyContext_Copy(PyObject * octx)
84138
{
85139
ENSURE_Context(octx, NULL)
86140
PyContext *ctx = (PyContext *)octx;
87-
return (PyObject *)context_new_from_vars(ctx->ctx_vars);
141+
PyHamtObject *vars = context_get_vars(ctx);
142+
PyObject *res = (PyObject *)context_new_from_vars(vars);
143+
Py_DECREF(vars);
144+
return res;
88145
}
89146

90147

@@ -96,7 +153,7 @@ PyContext_CopyCurrent(void)
96153
return NULL;
97154
}
98155

99-
return (PyObject *)context_new_from_vars(ctx->ctx_vars);
156+
return (PyObject *)context_new_from_vars(context_get_current_vars(ctx));
100157
}
101158

102159
static const char *
@@ -298,7 +355,7 @@ PyContextVar_Get(PyObject *ovar, PyObject *def, PyObject **val)
298355
#endif
299356

300357
assert(PyContext_CheckExact(ts->context));
301-
PyHamtObject *vars = ((PyContext *)ts->context)->ctx_vars;
358+
PyHamtObject *vars = context_get_current_vars((PyContext *)ts->context);
302359

303360
PyObject *found = NULL;
304361
int res = _PyHamt_Find(vars, (PyObject*)var, &found);
@@ -354,7 +411,8 @@ PyContextVar_Set(PyObject *ovar, PyObject *val)
354411
}
355412

356413
PyObject *old_val = NULL;
357-
int found = _PyHamt_Find(ctx->ctx_vars, (PyObject *)var, &old_val);
414+
int found = _PyHamt_Find(context_get_current_vars(ctx), (PyObject *)var,
415+
&old_val);
358416
if (found < 0) {
359417
return NULL;
360418
}
@@ -552,7 +610,10 @@ static PyObject *
552610
context_tp_iter(PyObject *op)
553611
{
554612
PyContext *self = _PyContext_CAST(op);
555-
return _PyHamt_NewIterKeys(self->ctx_vars);
613+
PyHamtObject *vars = context_get_vars(self);
614+
PyObject *res = _PyHamt_NewIterKeys(vars);
615+
Py_DECREF(vars);
616+
return res;
556617
}
557618

558619
static PyObject *
@@ -564,8 +625,11 @@ context_tp_richcompare(PyObject *v, PyObject *w, int op)
564625
Py_RETURN_NOTIMPLEMENTED;
565626
}
566627

567-
int res = _PyHamt_Eq(
568-
((PyContext *)v)->ctx_vars, ((PyContext *)w)->ctx_vars);
628+
PyHamtObject *v_vars = context_get_vars((PyContext *)v);
629+
PyHamtObject *w_vars = context_get_vars((PyContext *)w);
630+
int res = _PyHamt_Eq(v_vars, w_vars);
631+
Py_DECREF(v_vars);
632+
Py_DECREF(w_vars);
569633
if (res < 0) {
570634
return NULL;
571635
}
@@ -586,7 +650,10 @@ static Py_ssize_t
586650
context_tp_len(PyObject *op)
587651
{
588652
PyContext *self = _PyContext_CAST(op);
589-
return _PyHamt_Len(self->ctx_vars);
653+
PyHamtObject *vars = context_get_vars(self);
654+
Py_ssize_t res = _PyHamt_Len(vars);
655+
Py_DECREF(vars);
656+
return res;
590657
}
591658

592659
static PyObject *
@@ -597,15 +664,19 @@ context_tp_subscript(PyObject *op, PyObject *key)
597664
}
598665
PyObject *val = NULL;
599666
PyContext *self = _PyContext_CAST(op);
600-
int found = _PyHamt_Find(self->ctx_vars, key, &val);
667+
PyHamtObject *vars = context_get_vars(self);
668+
int found = _PyHamt_Find(vars, key, &val);
669+
/* `val` is borrowed from `vars`, take a reference before dropping it. */
670+
Py_XINCREF(val);
671+
Py_DECREF(vars);
601672
if (found < 0) {
602673
return NULL;
603674
}
604675
if (found == 0) {
605676
PyErr_SetObject(PyExc_KeyError, key);
606677
return NULL;
607678
}
608-
return Py_NewRef(val);
679+
return val;
609680
}
610681

611682
static int
@@ -616,7 +687,10 @@ context_tp_contains(PyObject *op, PyObject *key)
616687
}
617688
PyObject *val = NULL;
618689
PyContext *self = _PyContext_CAST(op);
619-
return _PyHamt_Find(self->ctx_vars, key, &val);
690+
PyHamtObject *vars = context_get_vars(self);
691+
int res = _PyHamt_Find(vars, key, &val);
692+
Py_DECREF(vars);
693+
return res;
620694
}
621695

622696

@@ -643,14 +717,18 @@ _contextvars_Context_get_impl(PyContext *self, PyObject *key,
643717
}
644718

645719
PyObject *val = NULL;
646-
int found = _PyHamt_Find(self->ctx_vars, key, &val);
720+
PyHamtObject *vars = context_get_vars(self);
721+
int found = _PyHamt_Find(vars, key, &val);
722+
/* `val` is borrowed from `vars`, take a reference before dropping it. */
723+
Py_XINCREF(val);
724+
Py_DECREF(vars);
647725
if (found < 0) {
648726
return NULL;
649727
}
650728
if (found == 0) {
651729
return Py_NewRef(default_value);
652730
}
653-
return Py_NewRef(val);
731+
return val;
654732
}
655733

656734

@@ -666,7 +744,10 @@ static PyObject *
666744
_contextvars_Context_items_impl(PyContext *self)
667745
/*[clinic end generated code: output=fa1655c8a08502af input=00db64ae379f9f42]*/
668746
{
669-
return _PyHamt_NewIterItems(self->ctx_vars);
747+
PyHamtObject *vars = context_get_vars(self);
748+
PyObject *res = _PyHamt_NewIterItems(vars);
749+
Py_DECREF(vars);
750+
return res;
670751
}
671752

672753

@@ -680,7 +761,10 @@ static PyObject *
680761
_contextvars_Context_keys_impl(PyContext *self)
681762
/*[clinic end generated code: output=177227c6b63ec0e2 input=114b53aebca3449c]*/
682763
{
683-
return _PyHamt_NewIterKeys(self->ctx_vars);
764+
PyHamtObject *vars = context_get_vars(self);
765+
PyObject *res = _PyHamt_NewIterKeys(vars);
766+
Py_DECREF(vars);
767+
return res;
684768
}
685769

686770

@@ -694,7 +778,10 @@ static PyObject *
694778
_contextvars_Context_values_impl(PyContext *self)
695779
/*[clinic end generated code: output=d286dabfc8db6dde input=ce8075d04a6ea526]*/
696780
{
697-
return _PyHamt_NewIterValues(self->ctx_vars);
781+
PyHamtObject *vars = context_get_vars(self);
782+
PyObject *res = _PyHamt_NewIterValues(vars);
783+
Py_DECREF(vars);
784+
return res;
698785
}
699786

700787

@@ -708,7 +795,10 @@ static PyObject *
708795
_contextvars_Context_copy_impl(PyContext *self)
709796
/*[clinic end generated code: output=30ba8896c4707a15 input=ebafdbdd9c72d592]*/
710797
{
711-
return (PyObject *)context_new_from_vars(self->ctx_vars);
798+
PyHamtObject *vars = context_get_vars(self);
799+
PyObject *res = (PyObject *)context_new_from_vars(vars);
800+
Py_DECREF(vars);
801+
return res;
712802
}
713803

714804

@@ -796,12 +886,12 @@ contextvar_set(PyContextVar *var, PyObject *val)
796886
}
797887

798888
PyHamtObject *new_vars = _PyHamt_Assoc(
799-
ctx->ctx_vars, (PyObject *)var, val);
889+
context_get_current_vars(ctx), (PyObject *)var, val);
800890
if (new_vars == NULL) {
801891
return -1;
802892
}
803893

804-
Py_SETREF(ctx->ctx_vars, new_vars);
894+
context_set_vars(ctx, new_vars);
805895

806896
#ifndef Py_GIL_DISABLED
807897
var->var_cached = val; /* borrow */
@@ -823,7 +913,7 @@ contextvar_del(PyContextVar *var)
823913
return -1;
824914
}
825915

826-
PyHamtObject *vars = ctx->ctx_vars;
916+
PyHamtObject *vars = context_get_current_vars(ctx);
827917
PyHamtObject *new_vars = _PyHamt_Without(vars, (PyObject *)var);
828918
if (new_vars == NULL) {
829919
return -1;
@@ -835,7 +925,7 @@ contextvar_del(PyContextVar *var)
835925
return -1;
836926
}
837927

838-
Py_SETREF(ctx->ctx_vars, new_vars);
928+
context_set_vars(ctx, new_vars);
839929
return 0;
840930
}
841931

0 commit comments

Comments
 (0)