Repository: itamarst/eliot

Issue:
Current mechanism for sub-contexts is buggy and hard to extend
Currently we support `asyncio` sub-contexts to the thread-local context. But there's a bug: it only enables e.g. asyncio contexts in the _current thread_ if you call `use_asyncio_context()`.

Additionally, we want to support a different kind of sub-context, for Twisted's `inlineCallbacks`. There's no good way to support both, or even to make sure only one gets registered.

Discussion:

Patch:
diff --git a/docs/source/news.rst b/docs/source/news.rst
index 886bb3d..40e7d75 100644
--- a/docs/source/news.rst
+++ b/docs/source/news.rst
@@ -20,6 +20,7 @@ Features:
 
 Bug fixes:
 
+* Fixed bug in the ``asyncio`` coroutine support where only the thread where ``use_asyncio_context()`` was called supported coroutine-specific contexts. Fixes issue #388.
 * ``ILogger.write`` is now explicitly thread-safe. The ``MemoryLogger`` (as used
   by tests) implementation of this method which was previously not thread-safe
   is now thread-safe. Thanks to Jean-Paul Calderone for the patch. Fixes issue
diff --git a/eliot/__init__.py b/eliot/__init__.py
index 359ce53..da6b7c0 100644
--- a/eliot/__init__.py
+++ b/eliot/__init__.py
@@ -7,7 +7,7 @@ from warnings import warn
 from ._message import Message
 from ._action import (
     start_action, startTask, Action, preserve_context, current_action,
-    use_asyncio_context, log_call
+    log_call
 )
 from ._output import (
     ILogger,
@@ -20,7 +20,7 @@ from ._validation import Field, fields, MessageType, ActionType, ValidationError
 from ._traceback import write_traceback, writeFailure
 from ._errors import register_exception_extractor
 from ._version import get_versions
-
+from ._asyncio import use_asyncio_context
 
 # Backwards compatibility:
 def add_destination(destination):
diff --git a/eliot/_action.py b/eliot/_action.py
index 09ed9b8..9cb2733 100644
--- a/eliot/_action.py
+++ b/eliot/_action.py
@@ -49,15 +49,23 @@ class _ExecutionContext(threading.local):
 
     The context is thread-specific, but can be made e.g. coroutine-specific by
     overriding C{get_sub_context}.
+
+    If we ever change this again we might want to make a ``threading.local``
+    subclass a sub-object?  Or come up with better API for thread-locals.
     """
 
     def __init__(self):
+        """
+        Every time an attribute is looked up from a new thread, this will be
+        called again for that thread, because this is a ``threading.local``
+        subclass.
+        """
         self._main_stack = []
         self.get_sub_context = lambda: None
 
     def _get_stack(self):
         """
-        Get the stack for the current asyncio Task.
+        Get the stack for the current context.
         """
         stack = self.get_sub_context()
         if stack is None:
@@ -91,20 +99,51 @@ class _ExecutionContext(threading.local):
         return stack[-1]
 
 
-_context = _ExecutionContext()
-current_action = _context.current
+class _ECOwner(object):
+    """Owner of the global execution context singleton.
 
+    It allows setting-once-only a replacement class for the default
+    L{_ExecutionContext}, so different sub-contexts (e.g. asyncio and
+    generators) don't stomp on each other.
 
-def use_asyncio_context():
+    @ivar context: The current global L{_ExecutionContext}.  Don't set it
+        directly, only get it!  You can use C{set} to set it.
     """
-    Use a logging context that is tied to the current asyncio coroutine.
 
-    Call this first thing, before doing any other logging.
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        """Reset to default context, to be used by tests only."""
+        self.context = _ExecutionContext()
+        self._set = False
+
+    def set(self, context_class):
+        """Set a new context of the given class.
 
-    Does not currently support event loops other than asyncio.
+        If the same class as current one, no changes are made.
+
+        @raises C{RuntimeError}: If the context has already been set to a
+            different class.
+        """
+        if self.context.__class__ == context_class:
+            return
+        if self._set:
+            raise RuntimeError(
+                "Context class already set to " + str(self.context.__class__)
+            )
+        self.context = context_class()
+        self._set = True
+
+
+_context_owner = _ECOwner()
+
+
+def current_action():
+    """
+    @return: The current C{Action} in context, or C{None} if none were set.
     """
-    from ._asyncio import AsyncioContext
-    _context.get_sub_context = AsyncioContext().get_stack
+    return _context_owner.context.current()
 
 
 class TaskLevel(object):
@@ -422,11 +461,11 @@ class Action(object):
         """
         Run the given function with this L{Action} as its execution context.
         """
-        _context.push(self)
+        _context_owner.context.push(self)
         try:
             return f(*args, **kwargs)
         finally:
-            _context.pop()
+            _context_owner.context.pop()
 
     def addSuccessFields(self, **fields):
         """
@@ -447,25 +486,25 @@ class Action(object):
 
         The action does NOT finish when the context is exited.
         """
-        _context.push(self)
+        _context_owner.context.push(self)
         try:
             yield self
         finally:
-            _context.pop()
+            _context_owner.context.pop()
 
     # Python context manager implementation:
     def __enter__(self):
         """
         Push this action onto the execution context.
         """
-        _context.push(self)
+        _context_owner.context.push(self)
         return self
 
     def __exit__(self, type, exception, traceback):
         """
         Pop this action off the execution context, log finish message.
         """
-        _context.pop()
+        _context_owner.context.pop()
         self.finish(exception)
 
 
diff --git a/eliot/_asyncio.py b/eliot/_asyncio.py
index 86c5900..1e4b3a5 100644
--- a/eliot/_asyncio.py
+++ b/eliot/_asyncio.py
@@ -2,11 +2,16 @@
 Support for asyncio coroutines.
 """
 
-from asyncio import Task
+try:
+    from asyncio import Task
+except ImportError:
+    Task = None  # Python 2
 from weakref import WeakKeyDictionary
 
+from ._action import _ExecutionContext
 
-class AsyncioContext:
+
+class AsyncioSubContext:
     """
     Per-Task context, allowing different coroutines to have different logging
     context.
@@ -31,3 +36,23 @@ class AsyncioContext:
         if task not in self._per_task:
             self._per_task[task] = []
         return self._per_task[task]
+
+
+class AsyncioExecutionContext(_ExecutionContext):
+    """ExecutionContext that supports asyncio sub-contexts."""
+
+    def __init__(self):
+        _ExecutionContext.__init__(self)
+        self.get_sub_context = AsyncioSubContext().get_stack
+
+
+def use_asyncio_context():
+    """
+    Use a logging context that is tied to the current asyncio coroutine.
+
+    Call this first thing, before doing any other logging.
+
+    Does not currently support event loops other than asyncio.
+    """
+    from ._action import _context_owner
+    _context_owner.set(AsyncioExecutionContext)


Test patch:
diff --git a/eliot/tests/corotests.py b/eliot/tests/corotests.py
index 78e3501..ec0a112 100644
--- a/eliot/tests/corotests.py
+++ b/eliot/tests/corotests.py
@@ -12,8 +12,8 @@ from unittest import TestCase
 from ..testing import capture_logging
 from ..parse import Parser
 from .. import start_action
-from .._action import _ExecutionContext, _context, use_asyncio_context
-from .._asyncio import AsyncioContext
+from .._action import _context_owner
+from .._asyncio import AsyncioExecutionContext, use_asyncio_context
 
 
 async def standalone_coro():
@@ -51,12 +51,9 @@ class CoroutineTests(TestCase):
     Tests for coroutines.
     """
     def setUp(self):
+        self.addCleanup(_context_owner.reset)
         use_asyncio_context()
 
-        def cleanup():
-            _context.get_sub_context = lambda: None
-        self.addCleanup(cleanup)
-
     @capture_logging(None)
     def test_coroutine_vs_main_thread_context(self, logger):
         """
@@ -108,8 +105,7 @@ class ContextTests(TestCase):
         Each thread gets its own execution context even when using asyncio
         contexts.
         """
-        ctx = _ExecutionContext()
-        ctx.get_sub_context = AsyncioContext().get_stack
+        ctx = AsyncioExecutionContext()
         first = object()
         ctx.push(first)
 
@@ -132,8 +128,7 @@ class ContextTests(TestCase):
         A coroutine has a different Eliot context than the thread that runs the
         event loop.
         """
-        ctx = _ExecutionContext()
-        ctx.get_sub_context = AsyncioContext().get_stack
+        ctx = AsyncioExecutionContext()
         current_context = []
 
         async def coro():
@@ -149,12 +144,44 @@ class ContextTests(TestCase):
             current_context,
             [("main", "B"), ("coro", None), ("coro", "A"), ("main", "B")])
 
+    def test_coroutine_vs_main_thread_context_different_thread(self):
+        """
+        A coroutine has a different Eliot context than the thread that runs the
+        event loop, when run in different thread than the one where the context
+        was created.
+        """
+        # Create context in one thread:
+        ctx = AsyncioExecutionContext()
+        current_context = []
+
+        # Run asyncio event loop and coroutines in a different thread:
+        def run():
+            event_loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(event_loop)
+
+            async def coro():
+                current_context.append(("coro", ctx.current()))
+                ctx.push("A")
+                current_context.append(("coro", ctx.current()))
+
+            ctx.push("B")
+            current_context.append(("main", ctx.current()))
+            run_coroutines(coro)
+            current_context.append(("main", ctx.current()))
+
+        thread = Thread(target=run)
+        thread.start()
+        thread.join()
+
+        self.assertEqual(
+            current_context,
+            [("main", "B"), ("coro", None), ("coro", "A"), ("main", "B")])
+
     def test_multiple_coroutines_contexts(self):
         """
         Each top-level ("Task") coroutine has its own Eliot separate context.
         """
-        ctx = _ExecutionContext()
-        ctx.get_sub_context = AsyncioContext().get_stack
+        ctx = AsyncioExecutionContext()
         current_context = []
 
         async def coro2():
@@ -181,8 +208,7 @@ class ContextTests(TestCase):
         A sub-coroutine (scheduled with await) inherits the parent coroutine's
         context.
         """
-        ctx = _ExecutionContext()
-        ctx.get_sub_context = AsyncioContext().get_stack
+        ctx = AsyncioExecutionContext()
         current_context = []
 
         async def coro2():
diff --git a/eliot/tests/test_action.py b/eliot/tests/test_action.py
index 500004d..7b42404 100644
--- a/eliot/tests/test_action.py
+++ b/eliot/tests/test_action.py
@@ -31,7 +31,7 @@ from .._action import (
     ACTION_STATUS_FIELD, ACTION_TYPE_FIELD, FAILED_STATUS, STARTED_STATUS,
     SUCCEEDED_STATUS, DuplicateChild, InvalidStartMessage, InvalidStatus,
     TaskLevel, WrittenAction, WrongActionType, WrongTask, WrongTaskLevel,
-    TooManyCalls, log_call
+    TooManyCalls, log_call, _context_owner
 )
 from .._message import (
     EXCEPTION_FIELD,
@@ -46,7 +46,6 @@ from .._validation import ActionType, Field, _ActionSerializers
 from ..testing import assertContainsFields, capture_logging
 from ..parse import Parser
 from .. import (
-    _action,
     add_destination,
     remove_destination,
     register_exception_extractor,
@@ -66,7 +65,6 @@ from .strategies import (
     union,
     written_from_pmap,
 )
-import eliot
 
 
 class ExecutionContextTests(TestCase):
@@ -178,14 +176,76 @@ class ExecutionContextTests(TestCase):
         self.assertEqual(valuesInThread, [second])
         self.assertIs(ctx.current(), first)
 
-    def test_globalInstance(self):
+
+class ContextOwnerTests(TestCase):
+    """Tests for the shared global context."""
+
+    def test_reset(self):
         """
-        A global L{_ExecutionContext} is exposed in the L{eliot._action}
-        module.
+        Resetting resets to a normal C{_ExecutionContext}.
+        """
+        _context_owner.context = object()
+        _context_owner.reset()
+        self.assertIsInstance(_context_owner.context, _ExecutionContext)
+
+    def test_current_action(self):
+        """
+        L{current_action} returns the current value of the current global
+        context.
+        """
+        self.addCleanup(_context_owner.reset)
+        _context_owner.context.push("A")
+        self.assertEqual(current_action(), "A")
+        _context_owner.context.push("B")
+        self.assertEqual(current_action(), "B")
+        _context_owner.reset()
+        self.assertEqual(current_action(), None)
+        _context_owner.context.push("C")
+        self.assertEqual(current_action(), "C")
+
+    def test_set(self):
+        """
+        It's possible to set a new C{_ExecutionContext} once.
+        """
+        self.addCleanup(_context_owner.reset)
+
+        class A(_ExecutionContext):
+            pass
+
+        _context_owner.set(A)
+        self.assertIsInstance(_context_owner.context, A)
+
+    def test_set_twice(self):
         """
-        self.assertIsInstance(_action._context, _ExecutionContext)
-        self.assertEqual(_action.current_action, _action._context.current)
-        self.assertEqual(eliot.current_action, _action._context.current)
+        Setting the same new C{_ExecutionContext} twice is no-op.
+        """
+        self.addCleanup(_context_owner.reset)
+
+        class A(_ExecutionContext):
+            pass
+
+        _context_owner.set(A)
+        a = _context_owner.context
+        _context_owner.set(A)
+        # Same object as before:
+        self.assertIs(_context_owner.context, a)
+
+    def test_set_twice_different_class(self):
+        """
+        Setting twice with different C{_ExecutionContext} classes raises a
+        C{RuntimeError}.
+        """
+        self.addCleanup(_context_owner.reset)
+
+        class A(_ExecutionContext):
+            pass
+
+        class B(_ExecutionContext):
+            pass
+
+        _context_owner.set(A)
+        with self.assertRaises(RuntimeError):
+            _context_owner.set(B)
 
 
 class ActionTests(TestCase):