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
@@ -64,6 +66,41 @@ contextvar_set(PyContextVar *var, PyObject *val);
6466static int
6567contextvar_del (PyContextVar * var );
6668
69+ static inline PyHamtObject *
70+ context_get_vars (PyContext * ctx )
71+ {
72+ PyHamtObject * vars ;
73+ Py_BEGIN_CRITICAL_SECTION (ctx );
74+ vars = ctx -> ctx_vars ;
75+ assert (vars != NULL );
76+ Py_INCREF (vars );
77+ Py_END_CRITICAL_SECTION ();
78+ return vars ;
79+ }
80+
81+ static inline PyHamtObject *
82+ context_get_current_vars (PyContext * ctx )
83+ {
84+ // ctx_vars written only by the owning thread, and read by other threads
85+ // only under the context's lock, a plain (non-atomic) load is okay
86+ PyHamtObject * vars = ctx -> ctx_vars ;
87+ assert (vars != NULL );
88+ return vars ;
89+ }
90+
91+ // Note: steals a reference to new_vars and must only be called by the thread
92+ // that has `ctx` as its current context.
93+ static inline void
94+ context_set_vars (PyContext * ctx , PyHamtObject * new_vars )
95+ {
96+ PyHamtObject * old_vars ;
97+ Py_BEGIN_CRITICAL_SECTION (ctx );
98+ old_vars = ctx -> ctx_vars ;
99+ ctx -> ctx_vars = new_vars ;
100+ Py_END_CRITICAL_SECTION ();
101+ Py_XDECREF (old_vars );
102+ }
103+
67104
68105PyObject *
69106_PyContext_NewHamtForTests (void )
@@ -84,7 +121,10 @@ PyContext_Copy(PyObject * octx)
84121{
85122 ENSURE_Context (octx , NULL )
86123 PyContext * ctx = (PyContext * )octx ;
87- return (PyObject * )context_new_from_vars (ctx -> ctx_vars );
124+ PyHamtObject * vars = context_get_vars (ctx );
125+ PyObject * res = (PyObject * )context_new_from_vars (vars );
126+ Py_DECREF (vars );
127+ return res ;
88128}
89129
90130
@@ -96,7 +136,7 @@ PyContext_CopyCurrent(void)
96136 return NULL ;
97137 }
98138
99- return (PyObject * )context_new_from_vars (ctx -> ctx_vars );
139+ return (PyObject * )context_new_from_vars (context_get_current_vars ( ctx ) );
100140}
101141
102142static const char *
@@ -293,7 +333,7 @@ PyContextVar_Get(PyObject *ovar, PyObject *def, PyObject **val)
293333#endif
294334
295335 assert (PyContext_CheckExact (ts -> context ));
296- PyHamtObject * vars = ((PyContext * )ts -> context )-> ctx_vars ;
336+ PyHamtObject * vars = context_get_current_vars ((PyContext * )ts -> context );
297337
298338 PyObject * found = NULL ;
299339 int res = _PyHamt_Find (vars , (PyObject * )var , & found );
@@ -349,7 +389,8 @@ PyContextVar_Set(PyObject *ovar, PyObject *val)
349389 }
350390
351391 PyObject * old_val = NULL ;
352- int found = _PyHamt_Find (ctx -> ctx_vars , (PyObject * )var , & old_val );
392+ int found = _PyHamt_Find (context_get_current_vars (ctx ), (PyObject * )var ,
393+ & old_val );
353394 if (found < 0 ) {
354395 return NULL ;
355396 }
@@ -547,7 +588,10 @@ static PyObject *
547588context_tp_iter (PyObject * op )
548589{
549590 PyContext * self = _PyContext_CAST (op );
550- return _PyHamt_NewIterKeys (self -> ctx_vars );
591+ PyHamtObject * vars = context_get_vars (self );
592+ PyObject * res = _PyHamt_NewIterKeys (vars );
593+ Py_DECREF (vars );
594+ return res ;
551595}
552596
553597static PyObject *
@@ -559,8 +603,11 @@ context_tp_richcompare(PyObject *v, PyObject *w, int op)
559603 Py_RETURN_NOTIMPLEMENTED ;
560604 }
561605
562- int res = _PyHamt_Eq (
563- ((PyContext * )v )-> ctx_vars , ((PyContext * )w )-> ctx_vars );
606+ PyHamtObject * v_vars = context_get_vars ((PyContext * )v );
607+ PyHamtObject * w_vars = context_get_vars ((PyContext * )w );
608+ int res = _PyHamt_Eq (v_vars , w_vars );
609+ Py_DECREF (v_vars );
610+ Py_DECREF (w_vars );
564611 if (res < 0 ) {
565612 return NULL ;
566613 }
@@ -581,7 +628,10 @@ static Py_ssize_t
581628context_tp_len (PyObject * op )
582629{
583630 PyContext * self = _PyContext_CAST (op );
584- return _PyHamt_Len (self -> ctx_vars );
631+ PyHamtObject * vars = context_get_vars (self );
632+ Py_ssize_t res = _PyHamt_Len (vars );
633+ Py_DECREF (vars );
634+ return res ;
585635}
586636
587637static PyObject *
@@ -592,15 +642,18 @@ context_tp_subscript(PyObject *op, PyObject *key)
592642 }
593643 PyObject * val = NULL ;
594644 PyContext * self = _PyContext_CAST (op );
595- int found = _PyHamt_Find (self -> ctx_vars , key , & val );
645+ PyHamtObject * vars = context_get_vars (self );
646+ int found = _PyHamt_Find (vars , key , & val );
647+ Py_XINCREF (val );
648+ Py_DECREF (vars );
596649 if (found < 0 ) {
597650 return NULL ;
598651 }
599652 if (found == 0 ) {
600653 PyErr_SetObject (PyExc_KeyError , key );
601654 return NULL ;
602655 }
603- return Py_NewRef ( val ) ;
656+ return val ;
604657}
605658
606659static int
@@ -611,7 +664,10 @@ context_tp_contains(PyObject *op, PyObject *key)
611664 }
612665 PyObject * val = NULL ;
613666 PyContext * self = _PyContext_CAST (op );
614- return _PyHamt_Find (self -> ctx_vars , key , & val );
667+ PyHamtObject * vars = context_get_vars (self );
668+ int res = _PyHamt_Find (vars , key , & val );
669+ Py_DECREF (vars );
670+ return res ;
615671}
616672
617673
@@ -637,14 +693,17 @@ _contextvars_Context_get_impl(PyContext *self, PyObject *key,
637693 }
638694
639695 PyObject * val = NULL ;
640- int found = _PyHamt_Find (self -> ctx_vars , key , & val );
696+ PyHamtObject * vars = context_get_vars (self );
697+ int found = _PyHamt_Find (vars , key , & val );
698+ Py_XINCREF (val );
699+ Py_DECREF (vars );
641700 if (found < 0 ) {
642701 return NULL ;
643702 }
644703 if (found == 0 ) {
645704 return Py_NewRef (default_value );
646705 }
647- return Py_NewRef ( val ) ;
706+ return val ;
648707}
649708
650709
@@ -660,7 +719,10 @@ static PyObject *
660719_contextvars_Context_items_impl (PyContext * self )
661720/*[clinic end generated code: output=fa1655c8a08502af input=00db64ae379f9f42]*/
662721{
663- return _PyHamt_NewIterItems (self -> ctx_vars );
722+ PyHamtObject * vars = context_get_vars (self );
723+ PyObject * res = _PyHamt_NewIterItems (vars );
724+ Py_DECREF (vars );
725+ return res ;
664726}
665727
666728
@@ -674,7 +736,10 @@ static PyObject *
674736_contextvars_Context_keys_impl (PyContext * self )
675737/*[clinic end generated code: output=177227c6b63ec0e2 input=114b53aebca3449c]*/
676738{
677- return _PyHamt_NewIterKeys (self -> ctx_vars );
739+ PyHamtObject * vars = context_get_vars (self );
740+ PyObject * res = _PyHamt_NewIterKeys (vars );
741+ Py_DECREF (vars );
742+ return res ;
678743}
679744
680745
@@ -688,7 +753,10 @@ static PyObject *
688753_contextvars_Context_values_impl (PyContext * self )
689754/*[clinic end generated code: output=d286dabfc8db6dde input=ce8075d04a6ea526]*/
690755{
691- return _PyHamt_NewIterValues (self -> ctx_vars );
756+ PyHamtObject * vars = context_get_vars (self );
757+ PyObject * res = _PyHamt_NewIterValues (vars );
758+ Py_DECREF (vars );
759+ return res ;
692760}
693761
694762
@@ -702,7 +770,10 @@ static PyObject *
702770_contextvars_Context_copy_impl (PyContext * self )
703771/*[clinic end generated code: output=30ba8896c4707a15 input=ebafdbdd9c72d592]*/
704772{
705- return (PyObject * )context_new_from_vars (self -> ctx_vars );
773+ PyHamtObject * vars = context_get_vars (self );
774+ PyObject * res = (PyObject * )context_new_from_vars (vars );
775+ Py_DECREF (vars );
776+ return res ;
706777}
707778
708779
@@ -790,12 +861,12 @@ contextvar_set(PyContextVar *var, PyObject *val)
790861 }
791862
792863 PyHamtObject * new_vars = _PyHamt_Assoc (
793- ctx -> ctx_vars , (PyObject * )var , val );
864+ context_get_current_vars ( ctx ) , (PyObject * )var , val );
794865 if (new_vars == NULL ) {
795866 return -1 ;
796867 }
797868
798- Py_SETREF (ctx -> ctx_vars , new_vars );
869+ context_set_vars (ctx , new_vars );
799870
800871#ifndef Py_GIL_DISABLED
801872 var -> var_cached = val ; /* borrow */
@@ -817,7 +888,7 @@ contextvar_del(PyContextVar *var)
817888 return -1 ;
818889 }
819890
820- PyHamtObject * vars = ctx -> ctx_vars ;
891+ PyHamtObject * vars = context_get_current_vars ( ctx ) ;
821892 PyHamtObject * new_vars = _PyHamt_Without (vars , (PyObject * )var );
822893 if (new_vars == NULL ) {
823894 return -1 ;
@@ -829,7 +900,7 @@ contextvar_del(PyContextVar *var)
829900 return -1 ;
830901 }
831902
832- Py_SETREF (ctx -> ctx_vars , new_vars );
903+ context_set_vars (ctx , new_vars );
833904 return 0 ;
834905}
835906
0 commit comments