1b2732e9dSimarom#!/usr/bin/env python
2b2732e9dSimarom#
3b2732e9dSimarom# Copyright 2010 Facebook
4b2732e9dSimarom#
5b2732e9dSimarom# Licensed under the Apache License, Version 2.0 (the "License"); you may
6b2732e9dSimarom# not use this file except in compliance with the License. You may obtain
7b2732e9dSimarom# a copy of the License at
8b2732e9dSimarom#
9b2732e9dSimarom#     http://www.apache.org/licenses/LICENSE-2.0
10b2732e9dSimarom#
11b2732e9dSimarom# Unless required by applicable law or agreed to in writing, software
12b2732e9dSimarom# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
13b2732e9dSimarom# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
14b2732e9dSimarom# License for the specific language governing permissions and limitations
15b2732e9dSimarom# under the License.
16b2732e9dSimarom
17b2732e9dSimarom"""`StackContext` allows applications to maintain threadlocal-like state
18b2732e9dSimaromthat follows execution as it moves to other execution contexts.
19b2732e9dSimarom
20b2732e9dSimaromThe motivating examples are to eliminate the need for explicit
21b2732e9dSimarom``async_callback`` wrappers (as in `tornado.web.RequestHandler`), and to
22b2732e9dSimaromallow some additional context to be kept for logging.
23b2732e9dSimarom
24b2732e9dSimaromThis is slightly magic, but it's an extension of the idea that an
25b2732e9dSimaromexception handler is a kind of stack-local state and when that stack
26b2732e9dSimaromis suspended and resumed in a new context that state needs to be
27b2732e9dSimarompreserved.  `StackContext` shifts the burden of restoring that state
28b2732e9dSimaromfrom each call site (e.g.  wrapping each `.AsyncHTTPClient` callback
29b2732e9dSimaromin ``async_callback``) to the mechanisms that transfer control from
30b2732e9dSimaromone context to another (e.g. `.AsyncHTTPClient` itself, `.IOLoop`,
31b2732e9dSimaromthread pools, etc).
32b2732e9dSimarom
33b2732e9dSimaromExample usage::
34b2732e9dSimarom
35b2732e9dSimarom    @contextlib.contextmanager
36b2732e9dSimarom    def die_on_error():
37b2732e9dSimarom        try:
38b2732e9dSimarom            yield
39b2732e9dSimarom        except Exception:
40b2732e9dSimarom            logging.error("exception in asynchronous operation",exc_info=True)
41b2732e9dSimarom            sys.exit(1)
42b2732e9dSimarom
43b2732e9dSimarom    with StackContext(die_on_error):
44b2732e9dSimarom        # Any exception thrown here *or in callback and its desendents*
45b2732e9dSimarom        # will cause the process to exit instead of spinning endlessly
46b2732e9dSimarom        # in the ioloop.
47b2732e9dSimarom        http_client.fetch(url, callback)
48b2732e9dSimarom    ioloop.start()
49b2732e9dSimarom
50b2732e9dSimaromMost applications shouln't have to work with `StackContext` directly.
51b2732e9dSimaromHere are a few rules of thumb for when it's necessary:
52b2732e9dSimarom
53b2732e9dSimarom* If you're writing an asynchronous library that doesn't rely on a
54b2732e9dSimarom  stack_context-aware library like `tornado.ioloop` or `tornado.iostream`
55b2732e9dSimarom  (for example, if you're writing a thread pool), use
56b2732e9dSimarom  `.stack_context.wrap()` before any asynchronous operations to capture the
57b2732e9dSimarom  stack context from where the operation was started.
58b2732e9dSimarom
59b2732e9dSimarom* If you're writing an asynchronous library that has some shared
60b2732e9dSimarom  resources (such as a connection pool), create those shared resources
61b2732e9dSimarom  within a ``with stack_context.NullContext():`` block.  This will prevent
62b2732e9dSimarom  ``StackContexts`` from leaking from one request to another.
63b2732e9dSimarom
64b2732e9dSimarom* If you want to write something like an exception handler that will
65b2732e9dSimarom  persist across asynchronous calls, create a new `StackContext` (or
66b2732e9dSimarom  `ExceptionStackContext`), and make your asynchronous calls in a ``with``
67b2732e9dSimarom  block that references your `StackContext`.
68b2732e9dSimarom"""
69b2732e9dSimarom
70b2732e9dSimaromfrom __future__ import absolute_import, division, print_function, with_statement
71b2732e9dSimarom
72b2732e9dSimaromimport sys
73b2732e9dSimaromimport threading
74b2732e9dSimarom
75b2732e9dSimaromfrom .util import raise_exc_info
76b2732e9dSimarom
77b2732e9dSimarom
78b2732e9dSimaromclass StackContextInconsistentError(Exception):
79b2732e9dSimarom    pass
80b2732e9dSimarom
81b2732e9dSimarom
82b2732e9dSimaromclass _State(threading.local):
83b2732e9dSimarom    def __init__(self):
84b2732e9dSimarom        self.contexts = (tuple(), None)
85b2732e9dSimarom_state = _State()
86b2732e9dSimarom
87b2732e9dSimarom
88b2732e9dSimaromclass StackContext(object):
89b2732e9dSimarom    """Establishes the given context as a StackContext that will be transferred.
90b2732e9dSimarom
91b2732e9dSimarom    Note that the parameter is a callable that returns a context
92b2732e9dSimarom    manager, not the context itself.  That is, where for a
93b2732e9dSimarom    non-transferable context manager you would say::
94b2732e9dSimarom
95b2732e9dSimarom      with my_context():
96b2732e9dSimarom
97b2732e9dSimarom    StackContext takes the function itself rather than its result::
98b2732e9dSimarom
99b2732e9dSimarom      with StackContext(my_context):
100b2732e9dSimarom
101b2732e9dSimarom    The result of ``with StackContext() as cb:`` is a deactivation
102b2732e9dSimarom    callback.  Run this callback when the StackContext is no longer
103b2732e9dSimarom    needed to ensure that it is not propagated any further (note that
104b2732e9dSimarom    deactivating a context does not affect any instances of that
105b2732e9dSimarom    context that are currently pending).  This is an advanced feature
106b2732e9dSimarom    and not necessary in most applications.
107b2732e9dSimarom    """
108b2732e9dSimarom    def __init__(self, context_factory):
109b2732e9dSimarom        self.context_factory = context_factory
110b2732e9dSimarom        self.contexts = []
111b2732e9dSimarom        self.active = True
112b2732e9dSimarom
113b2732e9dSimarom    def _deactivate(self):
114b2732e9dSimarom        self.active = False
115b2732e9dSimarom
116b2732e9dSimarom    # StackContext protocol
117b2732e9dSimarom    def enter(self):
118b2732e9dSimarom        context = self.context_factory()
119b2732e9dSimarom        self.contexts.append(context)
120b2732e9dSimarom        context.__enter__()
121b2732e9dSimarom
122b2732e9dSimarom    def exit(self, type, value, traceback):
123b2732e9dSimarom        context = self.contexts.pop()
124b2732e9dSimarom        context.__exit__(type, value, traceback)
125b2732e9dSimarom
126b2732e9dSimarom    # Note that some of this code is duplicated in ExceptionStackContext
127b2732e9dSimarom    # below.  ExceptionStackContext is more common and doesn't need
128b2732e9dSimarom    # the full generality of this class.
129b2732e9dSimarom    def __enter__(self):
130b2732e9dSimarom        self.old_contexts = _state.contexts
131b2732e9dSimarom        self.new_contexts = (self.old_contexts[0] + (self,), self)
132b2732e9dSimarom        _state.contexts = self.new_contexts
133b2732e9dSimarom
134b2732e9dSimarom        try:
135b2732e9dSimarom            self.enter()
136b2732e9dSimarom        except:
137b2732e9dSimarom            _state.contexts = self.old_contexts
138b2732e9dSimarom            raise
139b2732e9dSimarom
140b2732e9dSimarom        return self._deactivate
141b2732e9dSimarom
142b2732e9dSimarom    def __exit__(self, type, value, traceback):
143b2732e9dSimarom        try:
144b2732e9dSimarom            self.exit(type, value, traceback)
145b2732e9dSimarom        finally:
146b2732e9dSimarom            final_contexts = _state.contexts
147b2732e9dSimarom            _state.contexts = self.old_contexts
148b2732e9dSimarom
149b2732e9dSimarom            # Generator coroutines and with-statements with non-local
150b2732e9dSimarom            # effects interact badly.  Check here for signs of
151b2732e9dSimarom            # the stack getting out of sync.
152b2732e9dSimarom            # Note that this check comes after restoring _state.context
153b2732e9dSimarom            # so that if it fails things are left in a (relatively)
154b2732e9dSimarom            # consistent state.
155b2732e9dSimarom            if final_contexts is not self.new_contexts:
156b2732e9dSimarom                raise StackContextInconsistentError(
157b2732e9dSimarom                    'stack_context inconsistency (may be caused by yield '
158b2732e9dSimarom                    'within a "with StackContext" block)')
159b2732e9dSimarom
160b2732e9dSimarom            # Break up a reference to itself to allow for faster GC on CPython.
161b2732e9dSimarom            self.new_contexts = None
162b2732e9dSimarom
163b2732e9dSimarom
164b2732e9dSimaromclass ExceptionStackContext(object):
165b2732e9dSimarom    """Specialization of StackContext for exception handling.
166b2732e9dSimarom
167b2732e9dSimarom    The supplied ``exception_handler`` function will be called in the
168b2732e9dSimarom    event of an uncaught exception in this context.  The semantics are
169b2732e9dSimarom    similar to a try/finally clause, and intended use cases are to log
170b2732e9dSimarom    an error, close a socket, or similar cleanup actions.  The
171b2732e9dSimarom    ``exc_info`` triple ``(type, value, traceback)`` will be passed to the
172b2732e9dSimarom    exception_handler function.
173b2732e9dSimarom
174b2732e9dSimarom    If the exception handler returns true, the exception will be
175b2732e9dSimarom    consumed and will not be propagated to other exception handlers.
176b2732e9dSimarom    """
177b2732e9dSimarom    def __init__(self, exception_handler):
178b2732e9dSimarom        self.exception_handler = exception_handler
179b2732e9dSimarom        self.active = True
180b2732e9dSimarom
181b2732e9dSimarom    def _deactivate(self):
182b2732e9dSimarom        self.active = False
183b2732e9dSimarom
184b2732e9dSimarom    def exit(self, type, value, traceback):
185b2732e9dSimarom        if type is not None:
186b2732e9dSimarom            return self.exception_handler(type, value, traceback)
187b2732e9dSimarom
188b2732e9dSimarom    def __enter__(self):
189b2732e9dSimarom        self.old_contexts = _state.contexts
190b2732e9dSimarom        self.new_contexts = (self.old_contexts[0], self)
191b2732e9dSimarom        _state.contexts = self.new_contexts
192b2732e9dSimarom
193b2732e9dSimarom        return self._deactivate
194b2732e9dSimarom
195b2732e9dSimarom    def __exit__(self, type, value, traceback):
196b2732e9dSimarom        try:
197b2732e9dSimarom            if type is not None:
198b2732e9dSimarom                return self.exception_handler(type, value, traceback)
199b2732e9dSimarom        finally:
200b2732e9dSimarom            final_contexts = _state.contexts
201b2732e9dSimarom            _state.contexts = self.old_contexts
202b2732e9dSimarom
203b2732e9dSimarom            if final_contexts is not self.new_contexts:
204b2732e9dSimarom                raise StackContextInconsistentError(
205b2732e9dSimarom                    'stack_context inconsistency (may be caused by yield '
206b2732e9dSimarom                    'within a "with StackContext" block)')
207b2732e9dSimarom
208b2732e9dSimarom            # Break up a reference to itself to allow for faster GC on CPython.
209b2732e9dSimarom            self.new_contexts = None
210b2732e9dSimarom
211b2732e9dSimarom
212b2732e9dSimaromclass NullContext(object):
213b2732e9dSimarom    """Resets the `StackContext`.
214b2732e9dSimarom
215b2732e9dSimarom    Useful when creating a shared resource on demand (e.g. an
216b2732e9dSimarom    `.AsyncHTTPClient`) where the stack that caused the creating is
217b2732e9dSimarom    not relevant to future operations.
218b2732e9dSimarom    """
219b2732e9dSimarom    def __enter__(self):
220b2732e9dSimarom        self.old_contexts = _state.contexts
221b2732e9dSimarom        _state.contexts = (tuple(), None)
222b2732e9dSimarom
223b2732e9dSimarom    def __exit__(self, type, value, traceback):
224b2732e9dSimarom        _state.contexts = self.old_contexts
225b2732e9dSimarom
226b2732e9dSimarom
227b2732e9dSimaromdef _remove_deactivated(contexts):
228b2732e9dSimarom    """Remove deactivated handlers from the chain"""
229b2732e9dSimarom    # Clean ctx handlers
230b2732e9dSimarom    stack_contexts = tuple([h for h in contexts[0] if h.active])
231b2732e9dSimarom
232b2732e9dSimarom    # Find new head
233b2732e9dSimarom    head = contexts[1]
234b2732e9dSimarom    while head is not None and not head.active:
235b2732e9dSimarom        head = head.old_contexts[1]
236b2732e9dSimarom
237b2732e9dSimarom    # Process chain
238b2732e9dSimarom    ctx = head
239b2732e9dSimarom    while ctx is not None:
240b2732e9dSimarom        parent = ctx.old_contexts[1]
241b2732e9dSimarom
242b2732e9dSimarom        while parent is not None:
243b2732e9dSimarom            if parent.active:
244b2732e9dSimarom                break
245b2732e9dSimarom            ctx.old_contexts = parent.old_contexts
246b2732e9dSimarom            parent = parent.old_contexts[1]
247b2732e9dSimarom
248b2732e9dSimarom        ctx = parent
249b2732e9dSimarom
250b2732e9dSimarom    return (stack_contexts, head)
251b2732e9dSimarom
252b2732e9dSimarom
253b2732e9dSimaromdef wrap(fn):
254b2732e9dSimarom    """Returns a callable object that will restore the current `StackContext`
255b2732e9dSimarom    when executed.
256b2732e9dSimarom
257b2732e9dSimarom    Use this whenever saving a callback to be executed later in a
258b2732e9dSimarom    different execution context (either in a different thread or
259b2732e9dSimarom    asynchronously in the same thread).
260b2732e9dSimarom    """
261b2732e9dSimarom    # Check if function is already wrapped
262b2732e9dSimarom    if fn is None or hasattr(fn, '_wrapped'):
263b2732e9dSimarom        return fn
264b2732e9dSimarom
265b2732e9dSimarom    # Capture current stack head
266b2732e9dSimarom    # TODO: Any other better way to store contexts and update them in wrapped function?
267b2732e9dSimarom    cap_contexts = [_state.contexts]
268b2732e9dSimarom
269b2732e9dSimarom    def wrapped(*args, **kwargs):
270b2732e9dSimarom        ret = None
271b2732e9dSimarom        try:
272b2732e9dSimarom            # Capture old state
273b2732e9dSimarom            current_state = _state.contexts
274b2732e9dSimarom
275b2732e9dSimarom            # Remove deactivated items
276b2732e9dSimarom            cap_contexts[0] = contexts = _remove_deactivated(cap_contexts[0])
277b2732e9dSimarom
278b2732e9dSimarom            # Force new state
279b2732e9dSimarom            _state.contexts = contexts
280b2732e9dSimarom
281b2732e9dSimarom            # Current exception
282b2732e9dSimarom            exc = (None, None, None)
283b2732e9dSimarom            top = None
284b2732e9dSimarom
285b2732e9dSimarom            # Apply stack contexts
286b2732e9dSimarom            last_ctx = 0
287b2732e9dSimarom            stack = contexts[0]
288b2732e9dSimarom
289b2732e9dSimarom            # Apply state
290b2732e9dSimarom            for n in stack:
291b2732e9dSimarom                try:
292b2732e9dSimarom                    n.enter()
293b2732e9dSimarom                    last_ctx += 1
294b2732e9dSimarom                except:
295b2732e9dSimarom                    # Exception happened. Record exception info and store top-most handler
296b2732e9dSimarom                    exc = sys.exc_info()
297b2732e9dSimarom                    top = n.old_contexts[1]
298b2732e9dSimarom
299b2732e9dSimarom            # Execute callback if no exception happened while restoring state
300b2732e9dSimarom            if top is None:
301b2732e9dSimarom                try:
302b2732e9dSimarom                    ret = fn(*args, **kwargs)
303b2732e9dSimarom                except:
304b2732e9dSimarom                    exc = sys.exc_info()
305b2732e9dSimarom                    top = contexts[1]
306b2732e9dSimarom
307b2732e9dSimarom            # If there was exception, try to handle it by going through the exception chain
308b2732e9dSimarom            if top is not None:
309b2732e9dSimarom                exc = _handle_exception(top, exc)
310b2732e9dSimarom            else:
311b2732e9dSimarom                # Otherwise take shorter path and run stack contexts in reverse order
312b2732e9dSimarom                while last_ctx > 0:
313b2732e9dSimarom                    last_ctx -= 1
314b2732e9dSimarom                    c = stack[last_ctx]
315b2732e9dSimarom
316b2732e9dSimarom                    try:
317b2732e9dSimarom                        c.exit(*exc)
318b2732e9dSimarom                    except:
319b2732e9dSimarom                        exc = sys.exc_info()
320b2732e9dSimarom                        top = c.old_contexts[1]
321b2732e9dSimarom                        break
322b2732e9dSimarom                else:
323b2732e9dSimarom                    top = None
324b2732e9dSimarom
325b2732e9dSimarom                # If if exception happened while unrolling, take longer exception handler path
326b2732e9dSimarom                if top is not None:
327b2732e9dSimarom                    exc = _handle_exception(top, exc)
328b2732e9dSimarom
329b2732e9dSimarom            # If exception was not handled, raise it
330b2732e9dSimarom            if exc != (None, None, None):
331b2732e9dSimarom                raise_exc_info(exc)
332b2732e9dSimarom        finally:
333b2732e9dSimarom            _state.contexts = current_state
334b2732e9dSimarom        return ret
335b2732e9dSimarom
336b2732e9dSimarom    wrapped._wrapped = True
337b2732e9dSimarom    return wrapped
338b2732e9dSimarom
339b2732e9dSimarom
340b2732e9dSimaromdef _handle_exception(tail, exc):
341b2732e9dSimarom    while tail is not None:
342b2732e9dSimarom        try:
343b2732e9dSimarom            if tail.exit(*exc):
344b2732e9dSimarom                exc = (None, None, None)
345b2732e9dSimarom        except:
346b2732e9dSimarom            exc = sys.exc_info()
347b2732e9dSimarom
348b2732e9dSimarom        tail = tail.old_contexts[1]
349b2732e9dSimarom
350b2732e9dSimarom    return exc
351b2732e9dSimarom
352b2732e9dSimarom
353b2732e9dSimaromdef run_with_stack_context(context, func):
354b2732e9dSimarom    """Run a coroutine ``func`` in the given `StackContext`.
355b2732e9dSimarom
356b2732e9dSimarom    It is not safe to have a ``yield`` statement within a ``with StackContext``
357b2732e9dSimarom    block, so it is difficult to use stack context with `.gen.coroutine`.
358b2732e9dSimarom    This helper function runs the function in the correct context while
359b2732e9dSimarom    keeping the ``yield`` and ``with`` statements syntactically separate.
360b2732e9dSimarom
361b2732e9dSimarom    Example::
362b2732e9dSimarom
363b2732e9dSimarom        @gen.coroutine
364b2732e9dSimarom        def incorrect():
365b2732e9dSimarom            with StackContext(ctx):
366b2732e9dSimarom                # ERROR: this will raise StackContextInconsistentError
367b2732e9dSimarom                yield other_coroutine()
368b2732e9dSimarom
369b2732e9dSimarom        @gen.coroutine
370b2732e9dSimarom        def correct():
371b2732e9dSimarom            yield run_with_stack_context(StackContext(ctx), other_coroutine)
372b2732e9dSimarom
373b2732e9dSimarom    .. versionadded:: 3.1
374b2732e9dSimarom    """
375b2732e9dSimarom    with context:
376b2732e9dSimarom        return func()
377