Move generator context state to a generator context manager

src/allmydata/util/eliotutil.py

@@ -5,6 +5,8 @@ Tools aimed at the interaction between Tahoe-LAFS implementation and Eliot.
 from sys import exc_info
 from functools import wraps
 from contextlib import contextmanager
+from weakref import WeakKeyDictionary
+
 
 from eliot import (
     Message,
@@ -14,25 +16,44 @@ from twisted.internet.defer import (
     inlineCallbacks,
 )
 
-@contextmanager
-def _substitute_stack(substitute, target):
-    # Save whatever is there to begin with, making a copy ensures we don't get
-    # affected by any mutations that might happen while the substitute is in
-    # place.
-    saved = list(target)
-    # Put the substitute in place.  Preserve the identity of the target for no
-    # concrete reason but maybe it's a good idea.
-    target[:] = substitute
-    try:
-        # Let some code run.
-        yield
-    finally:
-        # Save whatever substitute state we ended up with back to the
-        # substitute.  Copying again, here.
-        substitute[:] = list(target)
-        # Restore the target to its original state.  Again, preserving
-        # identity.
-        target[:] = saved
+class _GeneratorContext(object):
+    def __init__(self, execution_context):
+        self._execution_context = execution_context
+        self._contexts = WeakKeyDictionary()
+        self._current_generator = None
+
+    def init_stack(self, generator):
+        stack = list(self._execution_context._get_stack())
+        self._contexts[generator] = stack
+
+    def get_stack(self):
+        if self._current_generator is None:
+            # If there is no currently active generator then we have no
+            # special stack to supply.  Let the execution context figure out a
+            # different answer on its own.
+            return None
+        # Otherwise, give back the action context stack we've been tracking
+        # for the currently active generator.  It must have been previously
+        # initialized (it's too late to do it now)!
+        return self._contexts[self._current_generator]
+
+    @contextmanager
+    def context(self, generator):
+        previous_generator = self._current_generator
+        try:
+            self._current_generator = generator
+            yield
+        finally:
+            self._current_generator = previous_generator
+
+
+from eliot._action import _context
+_the_generator_context = _GeneratorContext(_context)
+
+
+def use_generator_context():
+    _context.get_sub_context = _the_generator_context.get_stack
+use_generator_context()
 
 
 def eliot_friendly_generator_function(original):
@@ -49,16 +70,17 @@ def eliot_friendly_generator_function(original):
         # Keep track of the next value to deliver to the generator.
         value_in = None
 
-        # Start tracking our desired inward-facing action context stack.  This
-        # really wants some more help from Eliot.
-        from eliot._action import _context
-        context_in = list(_context._get_stack())
-
         # Create the generator with a call to the generator function.  This
         # happens with whatever Eliot action context happens to be active,
         # which is fine and correct and also irrelevant because no code in the
         # generator function can run until we call send or throw on it.
         gen = original(*a, **kw)
+
+        # Initialize the per-generator Eliot action context stack to the
+        # current action stack.  This might be the main stack or, if another
+        # decorated generator is running, it might be the stack for that
+        # generator.  Not our business.
+        _the_generator_context.init_stack(gen)
         try:
             while True:
                 try:
@@ -66,7 +88,7 @@ def eliot_friendly_generator_function(original):
                     # with the Eliot action context stack we've saved for it.
                     # Then the context manager will re-save it and restore the
                     # "outside" stack for us.
-                    with _substitute_stack(context_in, _context._get_stack()):
+                    with _the_generator_context.context(gen):
                         if ok:
                             value_out = gen.send(value_in)
                         else:
@@ -105,6 +127,7 @@ def eliot_friendly_generator_function(original):
 
     return wrapper
 
+
 def inline_callbacks(original):
     """
     Decorate a function like ``inlineCallbacks`` would but in a more