Stackless issue #239: add support for PEP 567 context variables

Add a private context attribute and appropriate methods to class
tasklet. Document the changes in the manual.
New methods: tasklet.set_context(context), tasklet.context_run(...)
New readonly attribute: tasklet.context_id
Add/improve pickling of the context of tasklets.
New pickle flag "PICKLEFLAGS_PICKLE_CONTEXT", new undocumented
function stackless._tasklet_get_unpicklable_state()

Author: akruis

Doc/library/stackless/pickling.rst

@@ -108,6 +108,25 @@ different address than *t1*, which was displayed earlier.
     objects and frame objects contain code objects. And code objects are
     usually incompatible between different minor versions of |CPY|.
 
+.. note::
+
+    If you pickle a tasklet, its :class:`~contextvars.Context` won't be pickled,
+    because :class:`~contextvars.Context` objects can't be pickled. See
+    :pep:`567` for an explanation.
+
+    It is sometimes possible enable pickling of :class:`~contextvars.Context` objects
+    in an application specific way (see for instance: :func:`copyreg.pickle` or
+    :attr:`pickle.Pickler.dispatch_table` or :attr:`pickle.Pickler.persistent_id`).
+    Such an application can set the pickle flag
+    :const:`~stackless.PICKLEFLAGS_PICKLE_CONTEXT` to include the
+    context in the pickled state of a tasklet.
+
+    Another option is to subclass :class:`tasklet` and overload the methods
+    :meth:`tasklet.__reduce_ex__` and :meth:`tasklet.__setstate__` to
+    pickle the values of particular :class:`~contextvars.ContextVar` objects together
+    with the tasklet.
+
+
 ======================
 Pickling other objects
 ======================

Doc/library/stackless/stackless.rst

@@ -51,6 +51,21 @@ Constants
       These constants have been added on a provisional basis (see :pep:`411`
       for details.)
 
+.. data:: PICKLEFLAGS_PICKLE_CONTEXT
+
+   This constant defines an option flag for the function
+   :func:`pickle_flags`.
+
+   If this flag is set, |SLP| assumes that a :class:`~contextvars.Context` object
+   is pickleable. As a consequence the state information returned by :meth:`tasklet.__reduce_ex__`
+   includes the context of the tasklet.
+
+   .. versionadded:: 3.7.6
+
+   .. note::
+      This constant has been added on a provisional basis (see :pep:`411`
+      for details.)
+
 ---------
 Functions
 ---------

Doc/library/stackless/tasklets.rst

@@ -110,6 +110,13 @@ The ``tasklet`` class
    :meth:`tasklet.setup`.  The difference is that when providing them to
    :meth:`tasklet.bind`, the tasklet is not made runnable yet.
 
+   .. versionadded:: 3.7.6
+
+      If *func* is not :data:`None`, this method also sets the
+      :class:`~contextvars.Context` object of this tasklet to the
+      :class:`~contextvars.Context` object of the current tasklet.
+      Therefore it is usually not required to set the context explicitly.
+
    *func* can be :data:`None` when providing arguments, in which case a previous call
    to :meth:`tasklet.bind` must have provided the function.
 
@@ -344,6 +351,61 @@ The ``tasklet`` class
      # Implement unsafe logic here.
      t.set_ignore_nesting(old_value)
 
+.. method:: tasklet.set_context(context)
+
+   .. versionadded:: 3.7.6
+
+   Set the :class:`~contextvars.Context` object to be used while this tasklet runs.
+
+   Every tasklet has a private context attribute.
+   When the tasklet runs, this context becomes the current context of the thread.
+
+   :param context: the context to be set
+   :type context: :class:`contextvars.Context`
+   :return: the tasklet itself
+   :rtype: :class:`tasklet`
+   :raises RuntimeError: if the tasklet is bound to a foreign thread and is current or scheduled.
+   :raises RuntimeError: if called from within :meth:`contextvars.Context.run`.
+
+   .. note::
+
+      The methods :meth:`__init__`, :meth:`bind` and :meth:`__setstate__` also set the context
+      of the tasklet they are called on to the context of the current tasklet. Therefore it is
+      usually not required to set the context explicitly.
+
+   .. note::
+      This method has been added on a provisional basis (see :pep:`411`
+      for details.)
+
+.. method:: tasklet.context_run(callable, \*args, \*\*kwargs)
+
+   .. versionadded:: 3.7.6
+
+   Execute ``callable(*args, **kwargs)`` in the context object of the tasklet
+   the contest_run method is called on. Return the result of the
+   execution or propagate an exception if one occurred.
+   This method is roughly equivalent following pseudo code::
+
+       def context_run(self, callable, *args, **kwargs):
+           saved_context = stackless.current._internal_get_context()
+           stackless.current.set_context(self._internal_get_context())
+           try:
+               return callable(*args, **kw)
+           finally:
+               stackless.current.set_context(saved_context)
+
+   See also :meth:`contextvars.Context.run` for additional information.
+   Use this method with care, because it lets you manipulate the context of
+   another tasklet. Often it is sufficient to use a copy of the context
+   instead of the original object::
+
+       copied_context = tasklet.context_run(contextvars.copy_context)
+       copied_context.run(...)
+
+   .. note::
+      This method has been added on a provisional basis (see :pep:`411`
+      for details.)
+
 .. method:: tasklet.__del__()
 
    .. versionadded:: 3.7
@@ -365,6 +427,13 @@ The ``tasklet`` class
 
    See :meth:`object.__setstate__`.
 
+   .. versionadded:: 3.7.6
+
+      If the tasklet becomes alive through this call and if *state* does not contain
+      a :class:`~contextvars.Context` object, then :meth:`~__setstate__` also sets
+      the :class:`~contextvars.Context` object of the
+      tasklet to the :class:`~contextvars.Context` object of the current tasklet.
+
    :param state: the state as given by ``__reduce_ex__(...)[2]``
    :type state: :class:`tuple`
    :return: self
@@ -423,6 +492,18 @@ The following attributes allow checking of user set situations:
    This attribute is ``True`` while this tasklet is within a
    :meth:`tasklet.set_ignore_nesting` block
 
+.. attribute:: tasklet.context_id
+
+   .. versionadded:: 3.7.6
+
+   This attribute is the :func:`id` of the :class:`~contextvars.Context` object to be used while this tasklet runs.
+   It is intended mostly for debugging.
+
+   .. note::
+      This attribute has been added on a provisional basis (see :pep:`411`
+      for details.)
+
+
 The following attributes allow identification of tasklet place:
 
 .. attribute:: tasklet.is_current
@@ -511,3 +592,72 @@ state transitions these functions are roughly equivalent to the following
 
    def schedule_remove():
        stackless.current.next.switch()
+
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Tasklets and Context Variables
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. versionadded:: 3.7.6
+
+Version 3.7 of the |PPL| adds context variables, see module :mod:`contextvars`.
+Usually they are used in connection with
+:mod:`asyncio`, but they are a useful concept for |SLP| too.
+Using context variables and multiple tasklets together didn't work well in |SLP| versions 3.7.0 to
+3.7.5, because all tasklets of a given thread shared the same context.
+
+Starting with version 3.7.6 |SLP| adds explicit support for context variables.
+Design requirements were:
+
+1. Be fully compatible with |CPY| and its design decisions.
+2. Be fully compatible with previous applications of |SLP|, which are unaware of context variables.
+3. Automatically share a context between related tasklets. This way a tasklet, that needs to set
+   a context variable, can delegate this duty to a sub-tasklet without the need to manage the
+   context of the sub-tasklet manually.
+4. Enable the integration of tasklet-based co-routines into the :mod:`asyncio` framework.
+   This is an obvious application which involves context variables and tasklets. See
+   `slp-coroutine <https://pypi.org/project/slp-coroutine>`_ for an example.
+
+Now each tasklet object has it own private context attribute. The design goals have some consequences:
+
+* The active :class:`~contextvars.Context` object of a thread (as defined by the |PPL|)
+  is the context of the :attr:`~stackless.current` tasklet. This implies that a tasklet switch,
+  switches the active context of the thread.
+
+* In accordance with the design decisions made in :pep:`567` the context of a tasklet can't be
+  accessed directly [#f1]_, but you can use the method :meth:`tasklet.context_run` to run arbitrary code
+  in this context. For instance ``tasklet.context_run(contextvars.copy_context())`` returns a copy
+  of the context.
+  The attribute :attr:`tasklet.context_id` can be used to test, if two tasklets share the context.
+
+* If you use the C-API, the context attribute of a tasklet is stored in the field *context* of the structure
+  :c:type:`PyTaskletObject` or :c:type:`PyThreadState`. This field is is either undefined (``NULL``) or a pointer to a
+  :class:`~contextvars.Context` object.
+  A tasklet, whose *context* is ``NULL`` **must** behave identically to a tasklet, whose context is an
+  empty :class:`~contextvars.Context` object [#f2]_. Therefore the |PY| API provides no way to distinguish
+  both states. Whenever the context of a tasklet is to be shared with another tasklet and `tasklet->context`
+  is initially `NULL`, it must be set to a newly created :class:`~contextvars.Context` object beforehand.
+  This affects the methods :meth:`~tasklet.context_run`, :meth:`~tasklet.__init__`, :meth:`~tasklet.bind`
+  and :meth:`~tasklet.__setstate__` as well as the attribute :attr:`tasklet.context_id`.
+
+* If the state of a tasklet changes from *not alive* to *bound* or to *alive* (methods :meth:`~tasklet.__init__`,
+  :meth:`~tasklet.bind` or :meth:`~tasklet.__setstate__`), the context
+  of the tasklet is set to the currently active context. This way a newly initialized tasklet automatically
+  shares the context of its creator.
+
+* The :mod:`contextvars` implementation of |CPY| imposes several restrictions on |SLP|. Especially the sanity checks in
+  :c:func:`PyContext_Enter` and :c:func:`PyContext_Exit` make it impossible to replace the current context within
+  the execution of the method :meth:`contextvars.Context.run`. In that case |SLP| raises :exc:`RuntimeError`.
+
+.. note::
+   Context support has been added on a provisional basis (see :pep:`411` for details.)
+
+.. rubric:: Footnotes
+
+.. [#f1]   Not exactly true. The return value of :meth:`tasklet.__reduce_ex__` can contain references to class
+           :class:`contextvars.Context`, but it is strongly discouraged, to use them for any other purpose
+           than pickling.
+
+.. [#f2]   Setting a context variable to a non default value changes the value of the field *context* from ``NULL``
+           to a pointer to a newly created :class:`~contextvars.Context` object. This can happen anytime in a
+           library call. Therefore any difference between an undefined context and an empty context causes ill defined
+           behavior.

Include/internal/slp_prickelpit.h

@@ -49,7 +49,8 @@ Py_ssize_t slp_from_tuple_with_nulls(PyObject **start, PyObject *tup);
 #define SLP_PICKLEFLAGS_PRESERVE_TRACING_STATE (1U)
 #define SLP_PICKLEFLAGS_PRESERVE_AG_FINALIZER  (1U<<1)
 #define SLP_PICKLEFLAGS_RESET_AG_FINALIZER     (1U<<2)
-#define SLP_PICKLEFLAGS__MAX_VALUE             ((1<<3)-1) /* must be a signed value */
+#define SLP_PICKLEFLAGS_PICKLE_CONTEXT         (1U<<3)
+#define SLP_PICKLEFLAGS__MAX_VALUE             ((1<<4)-1) /* must be a signed value */
 
 /* helper functions for module dicts */
 

Include/internal/stackless_impl.h

@@ -761,6 +761,11 @@ PyTaskletTStateStruc * slp_get_saved_tstate(PyTaskletObject *task);
 PyObject * slp_channel_seq_callback(struct _frame *f,  int throwflag, PyObject *retval);
 PyObject * slp_get_channel_callback(void);
 
+/*
+ * contextvars related prototypes
+ */
+PyObject* slp_context_run_callback(PyFrameObject *f, int exc, PyObject *result);
+
 /* macro for use when interrupting tasklets from watchdog */
 #define TASKLET_NESTING_OK(task) \
     (ts->st.nesting_level == 0 || \

Include/slp_structs.h

@@ -134,6 +134,9 @@ typedef struct _tasklet {
     int recursion_depth;
     PyObject *def_globals;
     PyObject *tsk_weakreflist;
+    /* If the tasklet is current: NULL. (The context of a current tasklet is always in ts->tasklet.)
+     * If the tasklet is not current: the context for the tasklet */
+    PyObject *context;
 } PyTaskletObject;
 
 

Lib/stackless.py

@@ -25,6 +25,7 @@ def __reduce_ex__(*args):
 PICKLEFLAGS_PRESERVE_TRACING_STATE = 1
 PICKLEFLAGS_PRESERVE_AG_FINALIZER = 2
 PICKLEFLAGS_RESET_AG_FINALIZER = 4
+PICKLEFLAGS_PICKLE_CONTEXT = 8
 
 # Backwards support for unpickling older pickles, even from 2.7
 from _stackless import _wrap
@@ -73,6 +74,37 @@ def __iter__(self):
 # expressions like "stackless.current" as well defined.
 current = runcount = main = debug = uncollectables = threads = pickle_with_tracing_state = None
 
+def _tasklet_get_unpicklable_state(tasklet):
+    """Get a dict with additional state, that can't be pickled
+
+    The method tasklet.__reduce_ex__() returns the picklable state and this
+    function returns a tuple containing the rest.
+
+    The items in the return value are:
+    'context': the context of the tasklet
+
+    Additional items may be added later.
+
+    Note: this function has been added on a provisional basis (see :pep:`411` for details.)
+    """
+    if not isinstance(tasklet, _stackless.tasklet):
+        raise TypeError("Argument must be a tasklet")
+
+    with atomic():
+        if tasklet.is_current or tasklet.thread_id != _stackless.current.thread_id:
+            # - A current tasklet can't be reduced.
+            # - We can't set pickle_flags for a foreign thread
+            # To mitigate these problems, we copy the context to a new tasklet
+            # (implicit copy by tasklet.__init__(callable, ...)) and reduce the new
+            # context instead
+            tasklet = tasklet.context_run(_stackless.tasklet, id)  # "id" is just an arbitrary callable
+
+        flags = pickle_flags(PICKLEFLAGS_PICKLE_CONTEXT, PICKLEFLAGS_PICKLE_CONTEXT)
+        try:
+            return {'context': tasklet.__reduce__()[2][8]}
+        finally:
+            pickle_flags(flags, PICKLEFLAGS_PICKLE_CONTEXT)
+
 def transmogrify():
     """
     this function creates a subclass of the ModuleType with properties.

Python/context.c

@@ -604,23 +604,32 @@ _contextvars_Context_copy_impl(PyContext *self)
 
 
 #ifdef STACKLESS
-static PyObject* context_run_callback(PyFrameObject *f, int exc, PyObject *result)
+PyObject* slp_context_run_callback(PyFrameObject *f, int exc, PyObject *result)
 {
     PyCFrameObject *cf = (PyCFrameObject *)f;
-    assert(PyContext_CheckExact(cf->ob1));
     PyObject *context = cf->ob1;
     cf->ob1 = NULL;
 
-    if (PyContext_Exit(context)) {
-        Py_CLEAR(result);
+    if (cf->i) {
+        /* called by tasklet.context_run(...) */
+        PyThreadState *ts = PyThreadState_GET();
+        assert(ts);
+        assert(NULL == context || PyContext_CheckExact(context));
+        Py_XSETREF(ts->context, context);
+        ts->context_ver++;
+    } else {
+        assert(PyContext_CheckExact(context));
+        if (PyContext_Exit(context)) {
+            Py_CLEAR(result);
+        }
+        Py_DECREF(context);
     }
 
-    Py_DECREF(context);
     SLP_STORE_NEXT_FRAME(PyThreadState_GET(), cf->f_back);
     return result;
 }
 
-SLP_DEF_INVALID_EXEC(context_run_callback)
+SLP_DEF_INVALID_EXEC(slp_context_run_callback)
 #endif
 
 
@@ -629,6 +638,7 @@ context_run(PyContext *self, PyObject *const *args,
             Py_ssize_t nargs, PyObject *kwnames)
 {
     STACKLESS_GETARG();
+    assert(NULL != self);
 
     if (nargs < 1) {
         PyErr_SetString(PyExc_TypeError,
@@ -644,11 +654,12 @@ context_run(PyContext *self, PyObject *const *args,
     PyThreadState *ts = PyThreadState_GET();
     PyCFrameObject *f = NULL;
     if (stackless) {
-        f = slp_cframe_new(context_run_callback, 1);
+        f = slp_cframe_new(slp_context_run_callback, 1);
         if (f == NULL)
             return NULL;
         Py_INCREF(self);
         f->ob1 = (PyObject *)self;
+        assert(f->i == 0);
         SLP_SET_CURRENT_FRAME(ts, (PyFrameObject *)f);
         /* f contains the only counted reference to current frame. This reference
          * keeps the fame alive during the following _PyObject_FastCallKeywords().
@@ -1326,7 +1337,7 @@ _PyContext_Init(void)
 
 #ifdef STACKLESS
     if (slp_register_execute(&PyCFrame_Type, "context_run_callback",
-        context_run_callback, SLP_REF_INVALID_EXEC(context_run_callback)) != 0)
+        slp_context_run_callback, SLP_REF_INVALID_EXEC(slp_context_run_callback)) != 0)
     {
         return 0;
     }

Stackless/changelog.txt

@@ -9,6 +9,16 @@ What's New in Stackless 3.X.X?
 
 *Release date: 20XX-XX-XX*
 
+- https://github.com/stackless-dev/stackless/issues/239
+  Add support for PEP 567 context variables to tasklets.
+  Each tasklet now has a contextvars.Context object, that becomes active during
+  the execution of the tasklet.
+  New tasklet methods "set_context()" and "context_run()".
+  New read only attribute "tasklet.context_id".
+  New constant "stackless.PICKLEFLAGS_PICKLE_CONTEXT".
+  And an intentionally undocumented function
+  "stackless._tasklet_get_unpicklable_state()"
+
 - https://github.com/stackless-dev/stackless/issues/245
   Prevent a crash, if you call tasklet.__init__() or tasklet.bind() with wrong
   argument types.