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