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