• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 #include "Python.h"
2 #include "pycore_call.h"          // _PyObject_VectorcallTstate()
3 #include "pycore_context.h"
4 #include "pycore_gc.h"            // _PyObject_GC_MAY_BE_TRACKED()
5 #include "pycore_hamt.h"
6 #include "pycore_initconfig.h"    // _PyStatus_OK()
7 #include "pycore_object.h"
8 #include "pycore_pyerrors.h"
9 #include "pycore_pystate.h"       // _PyThreadState_GET()
10 
11 
12 
13 #include "clinic/context.c.h"
14 /*[clinic input]
15 module _contextvars
16 [clinic start generated code]*/
17 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=a0955718c8b8cea6]*/
18 
19 
20 #define ENSURE_Context(o, err_ret)                                  \
21     if (!PyContext_CheckExact(o)) {                                 \
22         PyErr_SetString(PyExc_TypeError,                            \
23                         "an instance of Context was expected");     \
24         return err_ret;                                             \
25     }
26 
27 #define ENSURE_ContextVar(o, err_ret)                               \
28     if (!PyContextVar_CheckExact(o)) {                              \
29         PyErr_SetString(PyExc_TypeError,                            \
30                        "an instance of ContextVar was expected");   \
31         return err_ret;                                             \
32     }
33 
34 #define ENSURE_ContextToken(o, err_ret)                             \
35     if (!PyContextToken_CheckExact(o)) {                            \
36         PyErr_SetString(PyExc_TypeError,                            \
37                         "an instance of Token was expected");       \
38         return err_ret;                                             \
39     }
40 
41 
42 /////////////////////////// Context API
43 
44 
45 static PyContext *
46 context_new_empty(void);
47 
48 static PyContext *
49 context_new_from_vars(PyHamtObject *vars);
50 
51 static inline PyContext *
52 context_get(void);
53 
54 static PyContextToken *
55 token_new(PyContext *ctx, PyContextVar *var, PyObject *val);
56 
57 static PyContextVar *
58 contextvar_new(PyObject *name, PyObject *def);
59 
60 static int
61 contextvar_set(PyContextVar *var, PyObject *val);
62 
63 static int
64 contextvar_del(PyContextVar *var);
65 
66 
67 #ifdef WITH_FREELISTS
68 static struct _Py_context_freelist *
get_context_freelist(void)69 get_context_freelist(void)
70 {
71     struct _Py_object_freelists *freelists = _Py_object_freelists_GET();
72     return &freelists->contexts;
73 }
74 #endif
75 
76 
77 PyObject *
_PyContext_NewHamtForTests(void)78 _PyContext_NewHamtForTests(void)
79 {
80     return (PyObject *)_PyHamt_New();
81 }
82 
83 
84 PyObject *
PyContext_New(void)85 PyContext_New(void)
86 {
87     return (PyObject *)context_new_empty();
88 }
89 
90 
91 PyObject *
PyContext_Copy(PyObject * octx)92 PyContext_Copy(PyObject * octx)
93 {
94     ENSURE_Context(octx, NULL)
95     PyContext *ctx = (PyContext *)octx;
96     return (PyObject *)context_new_from_vars(ctx->ctx_vars);
97 }
98 
99 
100 PyObject *
PyContext_CopyCurrent(void)101 PyContext_CopyCurrent(void)
102 {
103     PyContext *ctx = context_get();
104     if (ctx == NULL) {
105         return NULL;
106     }
107 
108     return (PyObject *)context_new_from_vars(ctx->ctx_vars);
109 }
110 
111 
112 static int
_PyContext_Enter(PyThreadState * ts,PyObject * octx)113 _PyContext_Enter(PyThreadState *ts, PyObject *octx)
114 {
115     ENSURE_Context(octx, -1)
116     PyContext *ctx = (PyContext *)octx;
117 
118     if (ctx->ctx_entered) {
119         _PyErr_Format(ts, PyExc_RuntimeError,
120                       "cannot enter context: %R is already entered", ctx);
121         return -1;
122     }
123 
124     ctx->ctx_prev = (PyContext *)ts->context;  /* borrow */
125     ctx->ctx_entered = 1;
126 
127     ts->context = Py_NewRef(ctx);
128     ts->context_ver++;
129 
130     return 0;
131 }
132 
133 
134 int
PyContext_Enter(PyObject * octx)135 PyContext_Enter(PyObject *octx)
136 {
137     PyThreadState *ts = _PyThreadState_GET();
138     assert(ts != NULL);
139     return _PyContext_Enter(ts, octx);
140 }
141 
142 
143 static int
_PyContext_Exit(PyThreadState * ts,PyObject * octx)144 _PyContext_Exit(PyThreadState *ts, PyObject *octx)
145 {
146     ENSURE_Context(octx, -1)
147     PyContext *ctx = (PyContext *)octx;
148 
149     if (!ctx->ctx_entered) {
150         PyErr_Format(PyExc_RuntimeError,
151                      "cannot exit context: %R has not been entered", ctx);
152         return -1;
153     }
154 
155     if (ts->context != (PyObject *)ctx) {
156         /* Can only happen if someone misuses the C API */
157         PyErr_SetString(PyExc_RuntimeError,
158                         "cannot exit context: thread state references "
159                         "a different context object");
160         return -1;
161     }
162 
163     Py_SETREF(ts->context, (PyObject *)ctx->ctx_prev);
164     ts->context_ver++;
165 
166     ctx->ctx_prev = NULL;
167     ctx->ctx_entered = 0;
168 
169     return 0;
170 }
171 
172 int
PyContext_Exit(PyObject * octx)173 PyContext_Exit(PyObject *octx)
174 {
175     PyThreadState *ts = _PyThreadState_GET();
176     assert(ts != NULL);
177     return _PyContext_Exit(ts, octx);
178 }
179 
180 
181 PyObject *
PyContextVar_New(const char * name,PyObject * def)182 PyContextVar_New(const char *name, PyObject *def)
183 {
184     PyObject *pyname = PyUnicode_FromString(name);
185     if (pyname == NULL) {
186         return NULL;
187     }
188     PyContextVar *var = contextvar_new(pyname, def);
189     Py_DECREF(pyname);
190     return (PyObject *)var;
191 }
192 
193 
194 int
PyContextVar_Get(PyObject * ovar,PyObject * def,PyObject ** val)195 PyContextVar_Get(PyObject *ovar, PyObject *def, PyObject **val)
196 {
197     ENSURE_ContextVar(ovar, -1)
198     PyContextVar *var = (PyContextVar *)ovar;
199 
200     PyThreadState *ts = _PyThreadState_GET();
201     assert(ts != NULL);
202     if (ts->context == NULL) {
203         goto not_found;
204     }
205 
206 #ifndef Py_GIL_DISABLED
207     if (var->var_cached != NULL &&
208             var->var_cached_tsid == ts->id &&
209             var->var_cached_tsver == ts->context_ver)
210     {
211         *val = var->var_cached;
212         goto found;
213     }
214 #endif
215 
216     assert(PyContext_CheckExact(ts->context));
217     PyHamtObject *vars = ((PyContext *)ts->context)->ctx_vars;
218 
219     PyObject *found = NULL;
220     int res = _PyHamt_Find(vars, (PyObject*)var, &found);
221     if (res < 0) {
222         goto error;
223     }
224     if (res == 1) {
225         assert(found != NULL);
226 #ifndef Py_GIL_DISABLED
227         var->var_cached = found;  /* borrow */
228         var->var_cached_tsid = ts->id;
229         var->var_cached_tsver = ts->context_ver;
230 #endif
231 
232         *val = found;
233         goto found;
234     }
235 
236 not_found:
237     if (def == NULL) {
238         if (var->var_default != NULL) {
239             *val = var->var_default;
240             goto found;
241         }
242 
243         *val = NULL;
244         goto found;
245     }
246     else {
247         *val = def;
248         goto found;
249    }
250 
251 found:
252     Py_XINCREF(*val);
253     return 0;
254 
255 error:
256     *val = NULL;
257     return -1;
258 }
259 
260 
261 PyObject *
PyContextVar_Set(PyObject * ovar,PyObject * val)262 PyContextVar_Set(PyObject *ovar, PyObject *val)
263 {
264     ENSURE_ContextVar(ovar, NULL)
265     PyContextVar *var = (PyContextVar *)ovar;
266 
267     if (!PyContextVar_CheckExact(var)) {
268         PyErr_SetString(
269             PyExc_TypeError, "an instance of ContextVar was expected");
270         return NULL;
271     }
272 
273     PyContext *ctx = context_get();
274     if (ctx == NULL) {
275         return NULL;
276     }
277 
278     PyObject *old_val = NULL;
279     int found = _PyHamt_Find(ctx->ctx_vars, (PyObject *)var, &old_val);
280     if (found < 0) {
281         return NULL;
282     }
283 
284     Py_XINCREF(old_val);
285     PyContextToken *tok = token_new(ctx, var, old_val);
286     Py_XDECREF(old_val);
287 
288     if (contextvar_set(var, val)) {
289         Py_DECREF(tok);
290         return NULL;
291     }
292 
293     return (PyObject *)tok;
294 }
295 
296 
297 int
PyContextVar_Reset(PyObject * ovar,PyObject * otok)298 PyContextVar_Reset(PyObject *ovar, PyObject *otok)
299 {
300     ENSURE_ContextVar(ovar, -1)
301     ENSURE_ContextToken(otok, -1)
302     PyContextVar *var = (PyContextVar *)ovar;
303     PyContextToken *tok = (PyContextToken *)otok;
304 
305     if (tok->tok_used) {
306         PyErr_Format(PyExc_RuntimeError,
307                      "%R has already been used once", tok);
308         return -1;
309     }
310 
311     if (var != tok->tok_var) {
312         PyErr_Format(PyExc_ValueError,
313                      "%R was created by a different ContextVar", tok);
314         return -1;
315     }
316 
317     PyContext *ctx = context_get();
318     if (ctx != tok->tok_ctx) {
319         PyErr_Format(PyExc_ValueError,
320                      "%R was created in a different Context", tok);
321         return -1;
322     }
323 
324     tok->tok_used = 1;
325 
326     if (tok->tok_oldval == NULL) {
327         return contextvar_del(var);
328     }
329     else {
330         return contextvar_set(var, tok->tok_oldval);
331     }
332 }
333 
334 
335 /////////////////////////// PyContext
336 
337 /*[clinic input]
338 class _contextvars.Context "PyContext *" "&PyContext_Type"
339 [clinic start generated code]*/
340 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=bdf87f8e0cb580e8]*/
341 
342 
343 static inline PyContext *
_context_alloc(void)344 _context_alloc(void)
345 {
346     PyContext *ctx;
347 #ifdef WITH_FREELISTS
348     struct _Py_context_freelist *context_freelist = get_context_freelist();
349     if (context_freelist->numfree > 0) {
350         context_freelist->numfree--;
351         ctx = context_freelist->items;
352         context_freelist->items = (PyContext *)ctx->ctx_weakreflist;
353         OBJECT_STAT_INC(from_freelist);
354         ctx->ctx_weakreflist = NULL;
355         _Py_NewReference((PyObject *)ctx);
356     }
357     else
358 #endif
359     {
360         ctx = PyObject_GC_New(PyContext, &PyContext_Type);
361         if (ctx == NULL) {
362             return NULL;
363         }
364     }
365 
366     ctx->ctx_vars = NULL;
367     ctx->ctx_prev = NULL;
368     ctx->ctx_entered = 0;
369     ctx->ctx_weakreflist = NULL;
370 
371     return ctx;
372 }
373 
374 
375 static PyContext *
context_new_empty(void)376 context_new_empty(void)
377 {
378     PyContext *ctx = _context_alloc();
379     if (ctx == NULL) {
380         return NULL;
381     }
382 
383     ctx->ctx_vars = _PyHamt_New();
384     if (ctx->ctx_vars == NULL) {
385         Py_DECREF(ctx);
386         return NULL;
387     }
388 
389     _PyObject_GC_TRACK(ctx);
390     return ctx;
391 }
392 
393 
394 static PyContext *
context_new_from_vars(PyHamtObject * vars)395 context_new_from_vars(PyHamtObject *vars)
396 {
397     PyContext *ctx = _context_alloc();
398     if (ctx == NULL) {
399         return NULL;
400     }
401 
402     ctx->ctx_vars = (PyHamtObject*)Py_NewRef(vars);
403 
404     _PyObject_GC_TRACK(ctx);
405     return ctx;
406 }
407 
408 
409 static inline PyContext *
context_get(void)410 context_get(void)
411 {
412     PyThreadState *ts = _PyThreadState_GET();
413     assert(ts != NULL);
414     PyContext *current_ctx = (PyContext *)ts->context;
415     if (current_ctx == NULL) {
416         current_ctx = context_new_empty();
417         if (current_ctx == NULL) {
418             return NULL;
419         }
420         ts->context = (PyObject *)current_ctx;
421     }
422     return current_ctx;
423 }
424 
425 static int
context_check_key_type(PyObject * key)426 context_check_key_type(PyObject *key)
427 {
428     if (!PyContextVar_CheckExact(key)) {
429         // abort();
430         PyErr_Format(PyExc_TypeError,
431                      "a ContextVar key was expected, got %R", key);
432         return -1;
433     }
434     return 0;
435 }
436 
437 static PyObject *
context_tp_new(PyTypeObject * type,PyObject * args,PyObject * kwds)438 context_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
439 {
440     if (PyTuple_Size(args) || (kwds != NULL && PyDict_Size(kwds))) {
441         PyErr_SetString(
442             PyExc_TypeError, "Context() does not accept any arguments");
443         return NULL;
444     }
445     return PyContext_New();
446 }
447 
448 static int
context_tp_clear(PyContext * self)449 context_tp_clear(PyContext *self)
450 {
451     Py_CLEAR(self->ctx_prev);
452     Py_CLEAR(self->ctx_vars);
453     return 0;
454 }
455 
456 static int
context_tp_traverse(PyContext * self,visitproc visit,void * arg)457 context_tp_traverse(PyContext *self, visitproc visit, void *arg)
458 {
459     Py_VISIT(self->ctx_prev);
460     Py_VISIT(self->ctx_vars);
461     return 0;
462 }
463 
464 static void
context_tp_dealloc(PyContext * self)465 context_tp_dealloc(PyContext *self)
466 {
467     _PyObject_GC_UNTRACK(self);
468 
469     if (self->ctx_weakreflist != NULL) {
470         PyObject_ClearWeakRefs((PyObject*)self);
471     }
472     (void)context_tp_clear(self);
473 
474 #ifdef WITH_FREELISTS
475     struct _Py_context_freelist *context_freelist = get_context_freelist();
476     if (context_freelist->numfree >= 0 && context_freelist->numfree < PyContext_MAXFREELIST) {
477         context_freelist->numfree++;
478         self->ctx_weakreflist = (PyObject *)context_freelist->items;
479         context_freelist->items = self;
480         OBJECT_STAT_INC(to_freelist);
481     }
482     else
483 #endif
484     {
485         Py_TYPE(self)->tp_free(self);
486     }
487 }
488 
489 static PyObject *
context_tp_iter(PyContext * self)490 context_tp_iter(PyContext *self)
491 {
492     return _PyHamt_NewIterKeys(self->ctx_vars);
493 }
494 
495 static PyObject *
context_tp_richcompare(PyObject * v,PyObject * w,int op)496 context_tp_richcompare(PyObject *v, PyObject *w, int op)
497 {
498     if (!PyContext_CheckExact(v) || !PyContext_CheckExact(w) ||
499             (op != Py_EQ && op != Py_NE))
500     {
501         Py_RETURN_NOTIMPLEMENTED;
502     }
503 
504     int res = _PyHamt_Eq(
505         ((PyContext *)v)->ctx_vars, ((PyContext *)w)->ctx_vars);
506     if (res < 0) {
507         return NULL;
508     }
509 
510     if (op == Py_NE) {
511         res = !res;
512     }
513 
514     if (res) {
515         Py_RETURN_TRUE;
516     }
517     else {
518         Py_RETURN_FALSE;
519     }
520 }
521 
522 static Py_ssize_t
context_tp_len(PyContext * self)523 context_tp_len(PyContext *self)
524 {
525     return _PyHamt_Len(self->ctx_vars);
526 }
527 
528 static PyObject *
context_tp_subscript(PyContext * self,PyObject * key)529 context_tp_subscript(PyContext *self, PyObject *key)
530 {
531     if (context_check_key_type(key)) {
532         return NULL;
533     }
534     PyObject *val = NULL;
535     int found = _PyHamt_Find(self->ctx_vars, key, &val);
536     if (found < 0) {
537         return NULL;
538     }
539     if (found == 0) {
540         PyErr_SetObject(PyExc_KeyError, key);
541         return NULL;
542     }
543     return Py_NewRef(val);
544 }
545 
546 static int
context_tp_contains(PyContext * self,PyObject * key)547 context_tp_contains(PyContext *self, PyObject *key)
548 {
549     if (context_check_key_type(key)) {
550         return -1;
551     }
552     PyObject *val = NULL;
553     return _PyHamt_Find(self->ctx_vars, key, &val);
554 }
555 
556 
557 /*[clinic input]
558 _contextvars.Context.get
559     key: object
560     default: object = None
561     /
562 
563 Return the value for `key` if `key` has the value in the context object.
564 
565 If `key` does not exist, return `default`. If `default` is not given,
566 return None.
567 [clinic start generated code]*/
568 
569 static PyObject *
_contextvars_Context_get_impl(PyContext * self,PyObject * key,PyObject * default_value)570 _contextvars_Context_get_impl(PyContext *self, PyObject *key,
571                               PyObject *default_value)
572 /*[clinic end generated code: output=0c54aa7664268189 input=c8eeb81505023995]*/
573 {
574     if (context_check_key_type(key)) {
575         return NULL;
576     }
577 
578     PyObject *val = NULL;
579     int found = _PyHamt_Find(self->ctx_vars, key, &val);
580     if (found < 0) {
581         return NULL;
582     }
583     if (found == 0) {
584         return Py_NewRef(default_value);
585     }
586     return Py_NewRef(val);
587 }
588 
589 
590 /*[clinic input]
591 _contextvars.Context.items
592 
593 Return all variables and their values in the context object.
594 
595 The result is returned as a list of 2-tuples (variable, value).
596 [clinic start generated code]*/
597 
598 static PyObject *
_contextvars_Context_items_impl(PyContext * self)599 _contextvars_Context_items_impl(PyContext *self)
600 /*[clinic end generated code: output=fa1655c8a08502af input=00db64ae379f9f42]*/
601 {
602     return _PyHamt_NewIterItems(self->ctx_vars);
603 }
604 
605 
606 /*[clinic input]
607 _contextvars.Context.keys
608 
609 Return a list of all variables in the context object.
610 [clinic start generated code]*/
611 
612 static PyObject *
_contextvars_Context_keys_impl(PyContext * self)613 _contextvars_Context_keys_impl(PyContext *self)
614 /*[clinic end generated code: output=177227c6b63ec0e2 input=114b53aebca3449c]*/
615 {
616     return _PyHamt_NewIterKeys(self->ctx_vars);
617 }
618 
619 
620 /*[clinic input]
621 _contextvars.Context.values
622 
623 Return a list of all variables' values in the context object.
624 [clinic start generated code]*/
625 
626 static PyObject *
_contextvars_Context_values_impl(PyContext * self)627 _contextvars_Context_values_impl(PyContext *self)
628 /*[clinic end generated code: output=d286dabfc8db6dde input=ce8075d04a6ea526]*/
629 {
630     return _PyHamt_NewIterValues(self->ctx_vars);
631 }
632 
633 
634 /*[clinic input]
635 _contextvars.Context.copy
636 
637 Return a shallow copy of the context object.
638 [clinic start generated code]*/
639 
640 static PyObject *
_contextvars_Context_copy_impl(PyContext * self)641 _contextvars_Context_copy_impl(PyContext *self)
642 /*[clinic end generated code: output=30ba8896c4707a15 input=ebafdbdd9c72d592]*/
643 {
644     return (PyObject *)context_new_from_vars(self->ctx_vars);
645 }
646 
647 
648 static PyObject *
context_run(PyContext * self,PyObject * const * args,Py_ssize_t nargs,PyObject * kwnames)649 context_run(PyContext *self, PyObject *const *args,
650             Py_ssize_t nargs, PyObject *kwnames)
651 {
652     PyThreadState *ts = _PyThreadState_GET();
653 
654     if (nargs < 1) {
655         _PyErr_SetString(ts, PyExc_TypeError,
656                          "run() missing 1 required positional argument");
657         return NULL;
658     }
659 
660     if (_PyContext_Enter(ts, (PyObject *)self)) {
661         return NULL;
662     }
663 
664     PyObject *call_result = _PyObject_VectorcallTstate(
665         ts, args[0], args + 1, nargs - 1, kwnames);
666 
667     if (_PyContext_Exit(ts, (PyObject *)self)) {
668         Py_XDECREF(call_result);
669         return NULL;
670     }
671 
672     return call_result;
673 }
674 
675 
676 static PyMethodDef PyContext_methods[] = {
677     _CONTEXTVARS_CONTEXT_GET_METHODDEF
678     _CONTEXTVARS_CONTEXT_ITEMS_METHODDEF
679     _CONTEXTVARS_CONTEXT_KEYS_METHODDEF
680     _CONTEXTVARS_CONTEXT_VALUES_METHODDEF
681     _CONTEXTVARS_CONTEXT_COPY_METHODDEF
682     {"run", _PyCFunction_CAST(context_run), METH_FASTCALL | METH_KEYWORDS, NULL},
683     {NULL, NULL}
684 };
685 
686 static PySequenceMethods PyContext_as_sequence = {
687     0,                                   /* sq_length */
688     0,                                   /* sq_concat */
689     0,                                   /* sq_repeat */
690     0,                                   /* sq_item */
691     0,                                   /* sq_slice */
692     0,                                   /* sq_ass_item */
693     0,                                   /* sq_ass_slice */
694     (objobjproc)context_tp_contains,     /* sq_contains */
695     0,                                   /* sq_inplace_concat */
696     0,                                   /* sq_inplace_repeat */
697 };
698 
699 static PyMappingMethods PyContext_as_mapping = {
700     (lenfunc)context_tp_len,             /* mp_length */
701     (binaryfunc)context_tp_subscript,    /* mp_subscript */
702 };
703 
704 PyTypeObject PyContext_Type = {
705     PyVarObject_HEAD_INIT(&PyType_Type, 0)
706     "_contextvars.Context",
707     sizeof(PyContext),
708     .tp_methods = PyContext_methods,
709     .tp_as_mapping = &PyContext_as_mapping,
710     .tp_as_sequence = &PyContext_as_sequence,
711     .tp_iter = (getiterfunc)context_tp_iter,
712     .tp_dealloc = (destructor)context_tp_dealloc,
713     .tp_getattro = PyObject_GenericGetAttr,
714     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
715     .tp_richcompare = context_tp_richcompare,
716     .tp_traverse = (traverseproc)context_tp_traverse,
717     .tp_clear = (inquiry)context_tp_clear,
718     .tp_new = context_tp_new,
719     .tp_weaklistoffset = offsetof(PyContext, ctx_weakreflist),
720     .tp_hash = PyObject_HashNotImplemented,
721 };
722 
723 
724 /////////////////////////// ContextVar
725 
726 
727 static int
contextvar_set(PyContextVar * var,PyObject * val)728 contextvar_set(PyContextVar *var, PyObject *val)
729 {
730 #ifndef Py_GIL_DISABLED
731     var->var_cached = NULL;
732     PyThreadState *ts = _PyThreadState_GET();
733 #endif
734 
735     PyContext *ctx = context_get();
736     if (ctx == NULL) {
737         return -1;
738     }
739 
740     PyHamtObject *new_vars = _PyHamt_Assoc(
741         ctx->ctx_vars, (PyObject *)var, val);
742     if (new_vars == NULL) {
743         return -1;
744     }
745 
746     Py_SETREF(ctx->ctx_vars, new_vars);
747 
748 #ifndef Py_GIL_DISABLED
749     var->var_cached = val;  /* borrow */
750     var->var_cached_tsid = ts->id;
751     var->var_cached_tsver = ts->context_ver;
752 #endif
753     return 0;
754 }
755 
756 static int
contextvar_del(PyContextVar * var)757 contextvar_del(PyContextVar *var)
758 {
759 #ifndef Py_GIL_DISABLED
760     var->var_cached = NULL;
761 #endif
762 
763     PyContext *ctx = context_get();
764     if (ctx == NULL) {
765         return -1;
766     }
767 
768     PyHamtObject *vars = ctx->ctx_vars;
769     PyHamtObject *new_vars = _PyHamt_Without(vars, (PyObject *)var);
770     if (new_vars == NULL) {
771         return -1;
772     }
773 
774     if (vars == new_vars) {
775         Py_DECREF(new_vars);
776         PyErr_SetObject(PyExc_LookupError, (PyObject *)var);
777         return -1;
778     }
779 
780     Py_SETREF(ctx->ctx_vars, new_vars);
781     return 0;
782 }
783 
784 static Py_hash_t
contextvar_generate_hash(void * addr,PyObject * name)785 contextvar_generate_hash(void *addr, PyObject *name)
786 {
787     /* Take hash of `name` and XOR it with the object's addr.
788 
789        The structure of the tree is encoded in objects' hashes, which
790        means that sufficiently similar hashes would result in tall trees
791        with many Collision nodes.  Which would, in turn, result in slower
792        get and set operations.
793 
794        The XORing helps to ensure that:
795 
796        (1) sequentially allocated ContextVar objects have
797            different hashes;
798 
799        (2) context variables with equal names have
800            different hashes.
801     */
802 
803     Py_hash_t name_hash = PyObject_Hash(name);
804     if (name_hash == -1) {
805         return -1;
806     }
807 
808     Py_hash_t res = _Py_HashPointer(addr) ^ name_hash;
809     return res == -1 ? -2 : res;
810 }
811 
812 static PyContextVar *
contextvar_new(PyObject * name,PyObject * def)813 contextvar_new(PyObject *name, PyObject *def)
814 {
815     if (!PyUnicode_Check(name)) {
816         PyErr_SetString(PyExc_TypeError,
817                         "context variable name must be a str");
818         return NULL;
819     }
820 
821     PyContextVar *var = PyObject_GC_New(PyContextVar, &PyContextVar_Type);
822     if (var == NULL) {
823         return NULL;
824     }
825 
826     var->var_hash = contextvar_generate_hash(var, name);
827     if (var->var_hash == -1) {
828         Py_DECREF(var);
829         return NULL;
830     }
831 
832     var->var_name = Py_NewRef(name);
833 
834     var->var_default = Py_XNewRef(def);
835 
836 #ifndef Py_GIL_DISABLED
837     var->var_cached = NULL;
838     var->var_cached_tsid = 0;
839     var->var_cached_tsver = 0;
840 #endif
841 
842     if (_PyObject_GC_MAY_BE_TRACKED(name) ||
843             (def != NULL && _PyObject_GC_MAY_BE_TRACKED(def)))
844     {
845         PyObject_GC_Track(var);
846     }
847     return var;
848 }
849 
850 
851 /*[clinic input]
852 class _contextvars.ContextVar "PyContextVar *" "&PyContextVar_Type"
853 [clinic start generated code]*/
854 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=445da935fa8883c3]*/
855 
856 
857 static PyObject *
contextvar_tp_new(PyTypeObject * type,PyObject * args,PyObject * kwds)858 contextvar_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
859 {
860     static char *kwlist[] = {"", "default", NULL};
861     PyObject *name;
862     PyObject *def = NULL;
863 
864     if (!PyArg_ParseTupleAndKeywords(
865             args, kwds, "O|$O:ContextVar", kwlist, &name, &def))
866     {
867         return NULL;
868     }
869 
870     return (PyObject *)contextvar_new(name, def);
871 }
872 
873 static int
contextvar_tp_clear(PyContextVar * self)874 contextvar_tp_clear(PyContextVar *self)
875 {
876     Py_CLEAR(self->var_name);
877     Py_CLEAR(self->var_default);
878 #ifndef Py_GIL_DISABLED
879     self->var_cached = NULL;
880     self->var_cached_tsid = 0;
881     self->var_cached_tsver = 0;
882 #endif
883     return 0;
884 }
885 
886 static int
contextvar_tp_traverse(PyContextVar * self,visitproc visit,void * arg)887 contextvar_tp_traverse(PyContextVar *self, visitproc visit, void *arg)
888 {
889     Py_VISIT(self->var_name);
890     Py_VISIT(self->var_default);
891     return 0;
892 }
893 
894 static void
contextvar_tp_dealloc(PyContextVar * self)895 contextvar_tp_dealloc(PyContextVar *self)
896 {
897     PyObject_GC_UnTrack(self);
898     (void)contextvar_tp_clear(self);
899     Py_TYPE(self)->tp_free(self);
900 }
901 
902 static Py_hash_t
contextvar_tp_hash(PyContextVar * self)903 contextvar_tp_hash(PyContextVar *self)
904 {
905     return self->var_hash;
906 }
907 
908 static PyObject *
contextvar_tp_repr(PyContextVar * self)909 contextvar_tp_repr(PyContextVar *self)
910 {
911     _PyUnicodeWriter writer;
912 
913     _PyUnicodeWriter_Init(&writer);
914 
915     if (_PyUnicodeWriter_WriteASCIIString(
916             &writer, "<ContextVar name=", 17) < 0)
917     {
918         goto error;
919     }
920 
921     PyObject *name = PyObject_Repr(self->var_name);
922     if (name == NULL) {
923         goto error;
924     }
925     if (_PyUnicodeWriter_WriteStr(&writer, name) < 0) {
926         Py_DECREF(name);
927         goto error;
928     }
929     Py_DECREF(name);
930 
931     if (self->var_default != NULL) {
932         if (_PyUnicodeWriter_WriteASCIIString(&writer, " default=", 9) < 0) {
933             goto error;
934         }
935 
936         PyObject *def = PyObject_Repr(self->var_default);
937         if (def == NULL) {
938             goto error;
939         }
940         if (_PyUnicodeWriter_WriteStr(&writer, def) < 0) {
941             Py_DECREF(def);
942             goto error;
943         }
944         Py_DECREF(def);
945     }
946 
947     PyObject *addr = PyUnicode_FromFormat(" at %p>", self);
948     if (addr == NULL) {
949         goto error;
950     }
951     if (_PyUnicodeWriter_WriteStr(&writer, addr) < 0) {
952         Py_DECREF(addr);
953         goto error;
954     }
955     Py_DECREF(addr);
956 
957     return _PyUnicodeWriter_Finish(&writer);
958 
959 error:
960     _PyUnicodeWriter_Dealloc(&writer);
961     return NULL;
962 }
963 
964 
965 /*[clinic input]
966 _contextvars.ContextVar.get
967     default: object = NULL
968     /
969 
970 Return a value for the context variable for the current context.
971 
972 If there is no value for the variable in the current context, the method will:
973  * return the value of the default argument of the method, if provided; or
974  * return the default value for the context variable, if it was created
975    with one; or
976  * raise a LookupError.
977 [clinic start generated code]*/
978 
979 static PyObject *
_contextvars_ContextVar_get_impl(PyContextVar * self,PyObject * default_value)980 _contextvars_ContextVar_get_impl(PyContextVar *self, PyObject *default_value)
981 /*[clinic end generated code: output=0746bd0aa2ced7bf input=30aa2ab9e433e401]*/
982 {
983     if (!PyContextVar_CheckExact(self)) {
984         PyErr_SetString(
985             PyExc_TypeError, "an instance of ContextVar was expected");
986         return NULL;
987     }
988 
989     PyObject *val;
990     if (PyContextVar_Get((PyObject *)self, default_value, &val) < 0) {
991         return NULL;
992     }
993 
994     if (val == NULL) {
995         PyErr_SetObject(PyExc_LookupError, (PyObject *)self);
996         return NULL;
997     }
998 
999     return val;
1000 }
1001 
1002 /*[clinic input]
1003 _contextvars.ContextVar.set
1004     value: object
1005     /
1006 
1007 Call to set a new value for the context variable in the current context.
1008 
1009 The required value argument is the new value for the context variable.
1010 
1011 Returns a Token object that can be used to restore the variable to its previous
1012 value via the `ContextVar.reset()` method.
1013 [clinic start generated code]*/
1014 
1015 static PyObject *
_contextvars_ContextVar_set(PyContextVar * self,PyObject * value)1016 _contextvars_ContextVar_set(PyContextVar *self, PyObject *value)
1017 /*[clinic end generated code: output=446ed5e820d6d60b input=c0a6887154227453]*/
1018 {
1019     return PyContextVar_Set((PyObject *)self, value);
1020 }
1021 
1022 /*[clinic input]
1023 _contextvars.ContextVar.reset
1024     token: object
1025     /
1026 
1027 Reset the context variable.
1028 
1029 The variable is reset to the value it had before the `ContextVar.set()` that
1030 created the token was used.
1031 [clinic start generated code]*/
1032 
1033 static PyObject *
_contextvars_ContextVar_reset(PyContextVar * self,PyObject * token)1034 _contextvars_ContextVar_reset(PyContextVar *self, PyObject *token)
1035 /*[clinic end generated code: output=d4ee34d0742d62ee input=ebe2881e5af4ffda]*/
1036 {
1037     if (!PyContextToken_CheckExact(token)) {
1038         PyErr_Format(PyExc_TypeError,
1039                      "expected an instance of Token, got %R", token);
1040         return NULL;
1041     }
1042 
1043     if (PyContextVar_Reset((PyObject *)self, token)) {
1044         return NULL;
1045     }
1046 
1047     Py_RETURN_NONE;
1048 }
1049 
1050 
1051 static PyMemberDef PyContextVar_members[] = {
1052     {"name", _Py_T_OBJECT, offsetof(PyContextVar, var_name), Py_READONLY},
1053     {NULL}
1054 };
1055 
1056 static PyMethodDef PyContextVar_methods[] = {
1057     _CONTEXTVARS_CONTEXTVAR_GET_METHODDEF
1058     _CONTEXTVARS_CONTEXTVAR_SET_METHODDEF
1059     _CONTEXTVARS_CONTEXTVAR_RESET_METHODDEF
1060     {"__class_getitem__", Py_GenericAlias,
1061     METH_O|METH_CLASS,       PyDoc_STR("See PEP 585")},
1062     {NULL, NULL}
1063 };
1064 
1065 PyTypeObject PyContextVar_Type = {
1066     PyVarObject_HEAD_INIT(&PyType_Type, 0)
1067     "_contextvars.ContextVar",
1068     sizeof(PyContextVar),
1069     .tp_methods = PyContextVar_methods,
1070     .tp_members = PyContextVar_members,
1071     .tp_dealloc = (destructor)contextvar_tp_dealloc,
1072     .tp_getattro = PyObject_GenericGetAttr,
1073     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
1074     .tp_traverse = (traverseproc)contextvar_tp_traverse,
1075     .tp_clear = (inquiry)contextvar_tp_clear,
1076     .tp_new = contextvar_tp_new,
1077     .tp_free = PyObject_GC_Del,
1078     .tp_hash = (hashfunc)contextvar_tp_hash,
1079     .tp_repr = (reprfunc)contextvar_tp_repr,
1080 };
1081 
1082 
1083 /////////////////////////// Token
1084 
1085 static PyObject * get_token_missing(void);
1086 
1087 
1088 /*[clinic input]
1089 class _contextvars.Token "PyContextToken *" "&PyContextToken_Type"
1090 [clinic start generated code]*/
1091 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=338a5e2db13d3f5b]*/
1092 
1093 
1094 static PyObject *
token_tp_new(PyTypeObject * type,PyObject * args,PyObject * kwds)1095 token_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1096 {
1097     PyErr_SetString(PyExc_RuntimeError,
1098                     "Tokens can only be created by ContextVars");
1099     return NULL;
1100 }
1101 
1102 static int
token_tp_clear(PyContextToken * self)1103 token_tp_clear(PyContextToken *self)
1104 {
1105     Py_CLEAR(self->tok_ctx);
1106     Py_CLEAR(self->tok_var);
1107     Py_CLEAR(self->tok_oldval);
1108     return 0;
1109 }
1110 
1111 static int
token_tp_traverse(PyContextToken * self,visitproc visit,void * arg)1112 token_tp_traverse(PyContextToken *self, visitproc visit, void *arg)
1113 {
1114     Py_VISIT(self->tok_ctx);
1115     Py_VISIT(self->tok_var);
1116     Py_VISIT(self->tok_oldval);
1117     return 0;
1118 }
1119 
1120 static void
token_tp_dealloc(PyContextToken * self)1121 token_tp_dealloc(PyContextToken *self)
1122 {
1123     PyObject_GC_UnTrack(self);
1124     (void)token_tp_clear(self);
1125     Py_TYPE(self)->tp_free(self);
1126 }
1127 
1128 static PyObject *
token_tp_repr(PyContextToken * self)1129 token_tp_repr(PyContextToken *self)
1130 {
1131     _PyUnicodeWriter writer;
1132 
1133     _PyUnicodeWriter_Init(&writer);
1134 
1135     if (_PyUnicodeWriter_WriteASCIIString(&writer, "<Token", 6) < 0) {
1136         goto error;
1137     }
1138 
1139     if (self->tok_used) {
1140         if (_PyUnicodeWriter_WriteASCIIString(&writer, " used", 5) < 0) {
1141             goto error;
1142         }
1143     }
1144 
1145     if (_PyUnicodeWriter_WriteASCIIString(&writer, " var=", 5) < 0) {
1146         goto error;
1147     }
1148 
1149     PyObject *var = PyObject_Repr((PyObject *)self->tok_var);
1150     if (var == NULL) {
1151         goto error;
1152     }
1153     if (_PyUnicodeWriter_WriteStr(&writer, var) < 0) {
1154         Py_DECREF(var);
1155         goto error;
1156     }
1157     Py_DECREF(var);
1158 
1159     PyObject *addr = PyUnicode_FromFormat(" at %p>", self);
1160     if (addr == NULL) {
1161         goto error;
1162     }
1163     if (_PyUnicodeWriter_WriteStr(&writer, addr) < 0) {
1164         Py_DECREF(addr);
1165         goto error;
1166     }
1167     Py_DECREF(addr);
1168 
1169     return _PyUnicodeWriter_Finish(&writer);
1170 
1171 error:
1172     _PyUnicodeWriter_Dealloc(&writer);
1173     return NULL;
1174 }
1175 
1176 static PyObject *
token_get_var(PyContextToken * self,void * Py_UNUSED (ignored))1177 token_get_var(PyContextToken *self, void *Py_UNUSED(ignored))
1178 {
1179     return Py_NewRef(self->tok_var);;
1180 }
1181 
1182 static PyObject *
token_get_old_value(PyContextToken * self,void * Py_UNUSED (ignored))1183 token_get_old_value(PyContextToken *self, void *Py_UNUSED(ignored))
1184 {
1185     if (self->tok_oldval == NULL) {
1186         return get_token_missing();
1187     }
1188 
1189     return Py_NewRef(self->tok_oldval);
1190 }
1191 
1192 static PyGetSetDef PyContextTokenType_getsetlist[] = {
1193     {"var", (getter)token_get_var, NULL, NULL},
1194     {"old_value", (getter)token_get_old_value, NULL, NULL},
1195     {NULL}
1196 };
1197 
1198 static PyMethodDef PyContextTokenType_methods[] = {
1199     {"__class_getitem__",    Py_GenericAlias,
1200     METH_O|METH_CLASS,       PyDoc_STR("See PEP 585")},
1201     {NULL}
1202 };
1203 
1204 PyTypeObject PyContextToken_Type = {
1205     PyVarObject_HEAD_INIT(&PyType_Type, 0)
1206     "_contextvars.Token",
1207     sizeof(PyContextToken),
1208     .tp_methods = PyContextTokenType_methods,
1209     .tp_getset = PyContextTokenType_getsetlist,
1210     .tp_dealloc = (destructor)token_tp_dealloc,
1211     .tp_getattro = PyObject_GenericGetAttr,
1212     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
1213     .tp_traverse = (traverseproc)token_tp_traverse,
1214     .tp_clear = (inquiry)token_tp_clear,
1215     .tp_new = token_tp_new,
1216     .tp_free = PyObject_GC_Del,
1217     .tp_hash = PyObject_HashNotImplemented,
1218     .tp_repr = (reprfunc)token_tp_repr,
1219 };
1220 
1221 static PyContextToken *
token_new(PyContext * ctx,PyContextVar * var,PyObject * val)1222 token_new(PyContext *ctx, PyContextVar *var, PyObject *val)
1223 {
1224     PyContextToken *tok = PyObject_GC_New(PyContextToken, &PyContextToken_Type);
1225     if (tok == NULL) {
1226         return NULL;
1227     }
1228 
1229     tok->tok_ctx = (PyContext*)Py_NewRef(ctx);
1230 
1231     tok->tok_var = (PyContextVar*)Py_NewRef(var);
1232 
1233     tok->tok_oldval = Py_XNewRef(val);
1234 
1235     tok->tok_used = 0;
1236 
1237     PyObject_GC_Track(tok);
1238     return tok;
1239 }
1240 
1241 
1242 /////////////////////////// Token.MISSING
1243 
1244 
1245 static PyObject *
context_token_missing_tp_repr(PyObject * self)1246 context_token_missing_tp_repr(PyObject *self)
1247 {
1248     return PyUnicode_FromString("<Token.MISSING>");
1249 }
1250 
1251 static void
context_token_missing_tp_dealloc(_PyContextTokenMissing * Py_UNUSED (self))1252 context_token_missing_tp_dealloc(_PyContextTokenMissing *Py_UNUSED(self))
1253 {
1254 #ifdef Py_DEBUG
1255     /* The singleton is statically allocated. */
1256     _Py_FatalRefcountError("deallocating the token missing singleton");
1257 #else
1258     return;
1259 #endif
1260 }
1261 
1262 
1263 PyTypeObject _PyContextTokenMissing_Type = {
1264     PyVarObject_HEAD_INIT(&PyType_Type, 0)
1265     "Token.MISSING",
1266     sizeof(_PyContextTokenMissing),
1267     .tp_dealloc = (destructor)context_token_missing_tp_dealloc,
1268     .tp_getattro = PyObject_GenericGetAttr,
1269     .tp_flags = Py_TPFLAGS_DEFAULT,
1270     .tp_repr = context_token_missing_tp_repr,
1271 };
1272 
1273 
1274 static PyObject *
get_token_missing(void)1275 get_token_missing(void)
1276 {
1277     return (PyObject *)&_Py_SINGLETON(context_token_missing);
1278 }
1279 
1280 
1281 ///////////////////////////
1282 
1283 
1284 void
_PyContext_ClearFreeList(struct _Py_object_freelists * freelists,int is_finalization)1285 _PyContext_ClearFreeList(struct _Py_object_freelists *freelists, int is_finalization)
1286 {
1287 #ifdef WITH_FREELISTS
1288     struct _Py_context_freelist *state = &freelists->contexts;
1289     for (; state->numfree > 0; state->numfree--) {
1290         PyContext *ctx = state->items;
1291         state->items = (PyContext *)ctx->ctx_weakreflist;
1292         ctx->ctx_weakreflist = NULL;
1293         PyObject_GC_Del(ctx);
1294     }
1295     if (is_finalization) {
1296         state->numfree = -1;
1297     }
1298 #endif
1299 }
1300 
1301 
1302 PyStatus
_PyContext_Init(PyInterpreterState * interp)1303 _PyContext_Init(PyInterpreterState *interp)
1304 {
1305     if (!_Py_IsMainInterpreter(interp)) {
1306         return _PyStatus_OK();
1307     }
1308 
1309     PyObject *missing = get_token_missing();
1310     if (PyDict_SetItemString(
1311         _PyType_GetDict(&PyContextToken_Type), "MISSING", missing))
1312     {
1313         Py_DECREF(missing);
1314         return _PyStatus_ERR("can't init context types");
1315     }
1316     Py_DECREF(missing);
1317 
1318     return _PyStatus_OK();
1319 }
1320