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 descendants*
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 shouldn'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# pylint: skip-file
70
71from __future__ import absolute_import, division, print_function
72
73import sys
74import threading
75
76from salt.ext.tornado.util import raise_exc_info
77
78
79class StackContextInconsistentError(Exception):
80    pass
81
82
83class _State(threading.local):
84    def __init__(self):
85        self.contexts = (tuple(), None)
86
87
88_state = _State()
89
90
91class StackContext(object):
92    """Establishes the given context as a StackContext that will be transferred.
93
94    Note that the parameter is a callable that returns a context
95    manager, not the context itself.  That is, where for a
96    non-transferable context manager you would say::
97
98      with my_context():
99
100    StackContext takes the function itself rather than its result::
101
102      with StackContext(my_context):
103
104    The result of ``with StackContext() as cb:`` is a deactivation
105    callback.  Run this callback when the StackContext is no longer
106    needed to ensure that it is not propagated any further (note that
107    deactivating a context does not affect any instances of that
108    context that are currently pending).  This is an advanced feature
109    and not necessary in most applications.
110    """
111    def __init__(self, context_factory):
112        self.context_factory = context_factory
113        self.contexts = []
114        self.active = True
115
116    def _deactivate(self):
117        self.active = False
118
119    # StackContext protocol
120    def enter(self):
121        context = self.context_factory()
122        self.contexts.append(context)
123        context.__enter__()
124
125    def exit(self, type, value, traceback):
126        context = self.contexts.pop()
127        context.__exit__(type, value, traceback)
128
129    # Note that some of this code is duplicated in ExceptionStackContext
130    # below.  ExceptionStackContext is more common and doesn't need
131    # the full generality of this class.
132    def __enter__(self):
133        self.old_contexts = _state.contexts
134        self.new_contexts = (self.old_contexts[0] + (self,), self)
135        _state.contexts = self.new_contexts
136
137        try:
138            self.enter()
139        except:
140            _state.contexts = self.old_contexts
141            raise
142
143        return self._deactivate
144
145    def __exit__(self, type, value, traceback):
146        try:
147            self.exit(type, value, traceback)
148        finally:
149            final_contexts = _state.contexts
150            _state.contexts = self.old_contexts
151
152            # Generator coroutines and with-statements with non-local
153            # effects interact badly.  Check here for signs of
154            # the stack getting out of sync.
155            # Note that this check comes after restoring _state.context
156            # so that if it fails things are left in a (relatively)
157            # consistent state.
158            if final_contexts is not self.new_contexts:
159                raise StackContextInconsistentError(
160                    'stack_context inconsistency (may be caused by yield '
161                    'within a "with StackContext" block)')
162
163            # Break up a reference to itself to allow for faster GC on CPython.
164            self.new_contexts = None
165
166
167class ExceptionStackContext(object):
168    """Specialization of StackContext for exception handling.
169
170    The supplied ``exception_handler`` function will be called in the
171    event of an uncaught exception in this context.  The semantics are
172    similar to a try/finally clause, and intended use cases are to log
173    an error, close a socket, or similar cleanup actions.  The
174    ``exc_info`` triple ``(type, value, traceback)`` will be passed to the
175    exception_handler function.
176
177    If the exception handler returns true, the exception will be
178    consumed and will not be propagated to other exception handlers.
179    """
180    def __init__(self, exception_handler):
181        self.exception_handler = exception_handler
182        self.active = True
183
184    def _deactivate(self):
185        self.active = False
186
187    def exit(self, type, value, traceback):
188        if type is not None:
189            return self.exception_handler(type, value, traceback)
190
191    def __enter__(self):
192        self.old_contexts = _state.contexts
193        self.new_contexts = (self.old_contexts[0], self)
194        _state.contexts = self.new_contexts
195
196        return self._deactivate
197
198    def __exit__(self, type, value, traceback):
199        try:
200            if type is not None:
201                return self.exception_handler(type, value, traceback)
202        finally:
203            final_contexts = _state.contexts
204            _state.contexts = self.old_contexts
205
206            if final_contexts is not self.new_contexts:
207                raise StackContextInconsistentError(
208                    'stack_context inconsistency (may be caused by yield '
209                    'within a "with StackContext" block)')
210
211            # Break up a reference to itself to allow for faster GC on CPython.
212            self.new_contexts = None
213
214
215class NullContext(object):
216    """Resets the `StackContext`.
217
218    Useful when creating a shared resource on demand (e.g. an
219    `.AsyncHTTPClient`) where the stack that caused the creating is
220    not relevant to future operations.
221    """
222    def __enter__(self):
223        self.old_contexts = _state.contexts
224        _state.contexts = (tuple(), None)
225
226    def __exit__(self, type, value, traceback):
227        _state.contexts = self.old_contexts
228
229
230def _remove_deactivated(contexts):
231    """Remove deactivated handlers from the chain"""
232    # Clean ctx handlers
233    stack_contexts = tuple([h for h in contexts[0] if h.active])
234
235    # Find new head
236    head = contexts[1]
237    while head is not None and not head.active:
238        head = head.old_contexts[1]
239
240    # Process chain
241    ctx = head
242    while ctx is not None:
243        parent = ctx.old_contexts[1]
244
245        while parent is not None:
246            if parent.active:
247                break
248            ctx.old_contexts = parent.old_contexts
249            parent = parent.old_contexts[1]
250
251        ctx = parent
252
253    return (stack_contexts, head)
254
255
256def wrap(fn):
257    """Returns a callable object that will restore the current `StackContext`
258    when executed.
259
260    Use this whenever saving a callback to be executed later in a
261    different execution context (either in a different thread or
262    asynchronously in the same thread).
263    """
264    # Check if function is already wrapped
265    if fn is None or hasattr(fn, '_wrapped'):
266        return fn
267
268    # Capture current stack head
269    # TODO: Any other better way to store contexts and update them in wrapped function?
270    cap_contexts = [_state.contexts]
271
272    if not cap_contexts[0][0] and not cap_contexts[0][1]:
273        # Fast path when there are no active contexts.
274        def null_wrapper(*args, **kwargs):
275            try:
276                current_state = _state.contexts
277                _state.contexts = cap_contexts[0]
278                return fn(*args, **kwargs)
279            finally:
280                _state.contexts = current_state
281        null_wrapper._wrapped = True
282        return null_wrapper
283
284    def wrapped(*args, **kwargs):
285        ret = None
286        try:
287            # Capture old state
288            current_state = _state.contexts
289
290            # Remove deactivated items
291            cap_contexts[0] = contexts = _remove_deactivated(cap_contexts[0])
292
293            # Force new state
294            _state.contexts = contexts
295
296            # Current exception
297            exc = (None, None, None)
298            top = None
299
300            # Apply stack contexts
301            last_ctx = 0
302            stack = contexts[0]
303
304            # Apply state
305            for n in stack:
306                try:
307                    n.enter()
308                    last_ctx += 1
309                except:
310                    # Exception happened. Record exception info and store top-most handler
311                    exc = sys.exc_info()
312                    top = n.old_contexts[1]
313
314            # Execute callback if no exception happened while restoring state
315            if top is None:
316                try:
317                    ret = fn(*args, **kwargs)
318                except:
319                    exc = sys.exc_info()
320                    top = contexts[1]
321
322            # If there was exception, try to handle it by going through the exception chain
323            if top is not None:
324                exc = _handle_exception(top, exc)
325            else:
326                # Otherwise take shorter path and run stack contexts in reverse order
327                while last_ctx > 0:
328                    last_ctx -= 1
329                    c = stack[last_ctx]
330
331                    try:
332                        c.exit(*exc)
333                    except:
334                        exc = sys.exc_info()
335                        top = c.old_contexts[1]
336                        break
337                else:
338                    top = None
339
340                # If if exception happened while unrolling, take longer exception handler path
341                if top is not None:
342                    exc = _handle_exception(top, exc)
343
344            # If exception was not handled, raise it
345            if exc != (None, None, None):
346                raise_exc_info(exc)
347        finally:
348            _state.contexts = current_state
349        return ret
350
351    wrapped._wrapped = True
352    return wrapped
353
354
355def _handle_exception(tail, exc):
356    while tail is not None:
357        try:
358            if tail.exit(*exc):
359                exc = (None, None, None)
360        except:
361            exc = sys.exc_info()
362
363        tail = tail.old_contexts[1]
364
365    return exc
366
367
368def run_with_stack_context(context, func):
369    """Run a coroutine ``func`` in the given `StackContext`.
370
371    It is not safe to have a ``yield`` statement within a ``with StackContext``
372    block, so it is difficult to use stack context with `.gen.coroutine`.
373    This helper function runs the function in the correct context while
374    keeping the ``yield`` and ``with`` statements syntactically separate.
375
376    Example::
377
378        @gen.coroutine
379        def incorrect():
380            with StackContext(ctx):
381                # ERROR: this will raise StackContextInconsistentError
382                yield other_coroutine()
383
384        @gen.coroutine
385        def correct():
386            yield run_with_stack_context(StackContext(ctx), other_coroutine)
387
388    .. versionadded:: 3.1
389    """
390    with context:
391        return func()
392