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