stack_context.py 14 KB

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