Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with
or
.
Download ZIP
Newer
Older
100644 226 lines (195 sloc) 8.22 kb
721e25d Introduce StackContext, a way to automatically manage exception
Ben Darnell authored
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
18 that follows execution as it moves to other execution contexts.
19
20 The motivating examples are to eliminate the need for explicit
21 async_callback wrappers (as in tornado.web.RequestHandler), and to
22 allow some additional context to be kept for logging.
23
24 This is slightly magic, but it's an extension of the idea that an exception
25 handler is a kind of stack-local state and when that stack is suspended
26 and resumed in a new context that state needs to be preserved. StackContext
27 shifts the burden of restoring that state from each call site (e.g.
28 wrapping each AsyncHTTPClient callback in async_callback) to the mechanisms
29 that transfer control from one context to another (e.g. AsyncHTTPClient
30 itself, IOLoop, thread pools, etc).
31
96c2f2f @bdarnell Finish automatic documentation
bdarnell authored
32 Example usage::
33
b57bfd3 @bdarnell Reindent example in docstring, which apparently trips up emacs' inden…
bdarnell authored
34 @contextlib.contextmanager
35 def die_on_error():
36 try:
37 yield
17eed4f @bdarnell Replace all bare "except:" blocks with "except Exception:" so we don't
bdarnell authored
38 except Exception:
b57bfd3 @bdarnell Reindent example in docstring, which apparently trips up emacs' inden…
bdarnell authored
39 logging.error("exception in asynchronous operation",exc_info=True)
40 sys.exit(1)
721e25d Introduce StackContext, a way to automatically manage exception
Ben Darnell authored
41
b57bfd3 @bdarnell Reindent example in docstring, which apparently trips up emacs' inden…
bdarnell authored
42 with StackContext(die_on_error):
43 # Any exception thrown here *or in callback and its desendents*
44 # will cause the process to exit instead of spinning endlessly
45 # in the ioloop.
46 http_client.fetch(url, callback)
47 ioloop.start()
721e25d Introduce StackContext, a way to automatically manage exception
Ben Darnell authored
48 '''
49
50 from __future__ import with_statement
51
52 import contextlib
53 import functools
54 import itertools
cb75ec7 @bdarnell Use our own copy of contextlib.nested (which no longer exists in py3)
bdarnell authored
55 import sys
721e25d Introduce StackContext, a way to automatically manage exception
Ben Darnell authored
56 import threading
57
58 class _State(threading.local):
9d4444e @bdarnell Reindent everything to 4 spaces.
bdarnell authored
59 def __init__(self):
60 self.contexts = ()
721e25d Introduce StackContext, a way to automatically manage exception
Ben Darnell authored
61 _state = _State()
62
6e7c8b9 @bdarnell Avoid use of the @contextlib.contextmanager decorator.
bdarnell authored
63 class StackContext(object):
300d231 @bdarnell Finish this round of doc updates
bdarnell authored
64 '''Establishes the given context as a StackContext that will be transferred.
6e7c8b9 @bdarnell Avoid use of the @contextlib.contextmanager decorator.
bdarnell authored
65
300d231 @bdarnell Finish this round of doc updates
bdarnell authored
66 Note that the parameter is a callable that returns a context
67 manager, not the context itself. That is, where for a
68 non-transferable context manager you would say::
96c2f2f @bdarnell Finish automatic documentation
bdarnell authored
69
300d231 @bdarnell Finish this round of doc updates
bdarnell authored
70 with my_context():
96c2f2f @bdarnell Finish automatic documentation
bdarnell authored
71
300d231 @bdarnell Finish this round of doc updates
bdarnell authored
72 StackContext takes the function itself rather than its result::
96c2f2f @bdarnell Finish automatic documentation
bdarnell authored
73
300d231 @bdarnell Finish this round of doc updates
bdarnell authored
74 with StackContext(my_context):
75 '''
76 def __init__(self, context_factory):
6e7c8b9 @bdarnell Avoid use of the @contextlib.contextmanager decorator.
bdarnell authored
77 self.context_factory = context_factory
78
dd706f9 @bdarnell Implement ExceptionStackContext independently of StackContext to improve
bdarnell authored
79 # Note that some of this code is duplicated in ExceptionStackContext
80 # below. ExceptionStackContext is more common and doesn't need
81 # the full generality of this class.
6e7c8b9 @bdarnell Avoid use of the @contextlib.contextmanager decorator.
bdarnell authored
82 def __enter__(self):
83 self.old_contexts = _state.contexts
dd706f9 @bdarnell Implement ExceptionStackContext independently of StackContext to improve
bdarnell authored
84 # _state.contexts is a tuple of (class, arg) pairs
85 _state.contexts = (self.old_contexts +
86 ((StackContext, self.context_factory),))
6e7c8b9 @bdarnell Avoid use of the @contextlib.contextmanager decorator.
bdarnell authored
87 try:
88 self.context = self.context_factory()
89 self.context.__enter__()
90 except Exception:
91 _state.contexts = self.old_contexts
92 raise
93
94 def __exit__(self, type, value, traceback):
95 try:
96 return self.context.__exit__(type, value, traceback)
97 finally:
98 _state.contexts = self.old_contexts
99
dd706f9 @bdarnell Implement ExceptionStackContext independently of StackContext to improve
bdarnell authored
100 class ExceptionStackContext(object):
300d231 @bdarnell Finish this round of doc updates
bdarnell authored
101 '''Specialization of StackContext for exception handling.
102
103 The supplied exception_handler function will be called in the
104 event of an uncaught exception in this context. The semantics are
105 similar to a try/finally clause, and intended use cases are to log
106 an error, close a socket, or similar cleanup actions. The
107 exc_info triple (type, value, traceback) will be passed to the
108 exception_handler function.
109
110 If the exception handler returns true, the exception will be
111 consumed and will not be propagated to other exception handlers.
112 '''
dd706f9 @bdarnell Implement ExceptionStackContext independently of StackContext to improve
bdarnell authored
113 def __init__(self, exception_handler):
114 self.exception_handler = exception_handler
115
116 def __enter__(self):
117 self.old_contexts = _state.contexts
118 _state.contexts = (self.old_contexts +
119 ((ExceptionStackContext, self.exception_handler),))
120
121 def __exit__(self, type, value, traceback):
122 try:
6e7c8b9 @bdarnell Avoid use of the @contextlib.contextmanager decorator.
bdarnell authored
123 if type is not None:
dd706f9 @bdarnell Implement ExceptionStackContext independently of StackContext to improve
bdarnell authored
124 return self.exception_handler(type, value, traceback)
125 finally:
126 _state.contexts = self.old_contexts
721e25d Introduce StackContext, a way to automatically manage exception
Ben Darnell authored
127
6e7c8b9 @bdarnell Avoid use of the @contextlib.contextmanager decorator.
bdarnell authored
128 class NullContext(object):
9d4444e @bdarnell Reindent everything to 4 spaces.
bdarnell authored
129 '''Resets the StackContext.
721e25d Introduce StackContext, a way to automatically manage exception
Ben Darnell authored
130
9d4444e @bdarnell Reindent everything to 4 spaces.
bdarnell authored
131 Useful when creating a shared resource on demand (e.g. an AsyncHTTPClient)
132 where the stack that caused the creating is not relevant to future
133 operations.
134 '''
6e7c8b9 @bdarnell Avoid use of the @contextlib.contextmanager decorator.
bdarnell authored
135 def __enter__(self):
136 self.old_contexts = _state.contexts
9d4444e @bdarnell Reindent everything to 4 spaces.
bdarnell authored
137 _state.contexts = ()
6e7c8b9 @bdarnell Avoid use of the @contextlib.contextmanager decorator.
bdarnell authored
138
139 def __exit__(self, type, value, traceback):
140 _state.contexts = self.old_contexts
721e25d Introduce StackContext, a way to automatically manage exception
Ben Darnell authored
141
336bad3 @bdarnell StackContext speedups: special-case some common situations,
bdarnell authored
142 class _StackContextWrapper(functools.partial):
143 pass
144
41c14b8 @bdarnell Remove the ability to pre-bind parameters in stack_context.wrap.
bdarnell authored
145 def wrap(fn):
9d4444e @bdarnell Reindent everything to 4 spaces.
bdarnell authored
146 '''Returns a callable object that will resore the current StackContext
147 when executed.
148
149 Use this whenever saving a callback to be executed later in a
150 different execution context (either in a different thread or
151 asynchronously in the same thread).
152 '''
af29183 @bdarnell Small speedups
bdarnell authored
153 if fn is None or fn.__class__ is _StackContextWrapper:
154 return fn
9d4444e @bdarnell Reindent everything to 4 spaces.
bdarnell authored
155 # functools.wraps doesn't appear to work on functools.partial objects
156 #@functools.wraps(fn)
157 def wrapped(callback, contexts, *args, **kwargs):
336bad3 @bdarnell StackContext speedups: special-case some common situations,
bdarnell authored
158 if contexts is _state.contexts or not contexts:
159 callback(*args, **kwargs)
160 return
161 if not _state.contexts:
162 new_contexts = [cls(arg) for (cls, arg) in contexts]
7a7e241 @bdarnell Make StackContext more usable in libraries by reducing the need to
bdarnell authored
163 # If we're moving down the stack, _state.contexts is a prefix
164 # of contexts. For each element of contexts not in that prefix,
165 # create a new StackContext object.
166 # If we're moving up the stack (or to an entirely different stack),
167 # _state.contexts will have elements not in contexts. Use
168 # NullContext to clear the state and then recreate from contexts.
336bad3 @bdarnell StackContext speedups: special-case some common situations,
bdarnell authored
169 elif (len(_state.contexts) > len(contexts) or
dd706f9 @bdarnell Implement ExceptionStackContext independently of StackContext to improve
bdarnell authored
170 any(a[1] is not b[1]
7a7e241 @bdarnell Make StackContext more usable in libraries by reducing the need to
bdarnell authored
171 for a, b in itertools.izip(_state.contexts, contexts))):
172 # contexts have been removed or changed, so start over
173 new_contexts = ([NullContext()] +
dd706f9 @bdarnell Implement ExceptionStackContext independently of StackContext to improve
bdarnell authored
174 [cls(arg) for (cls,arg) in contexts])
7a7e241 @bdarnell Make StackContext more usable in libraries by reducing the need to
bdarnell authored
175 else:
dd706f9 @bdarnell Implement ExceptionStackContext independently of StackContext to improve
bdarnell authored
176 new_contexts = [cls(arg)
177 for (cls, arg) in contexts[len(_state.contexts):]]
6e7c8b9 @bdarnell Avoid use of the @contextlib.contextmanager decorator.
bdarnell authored
178 if len(new_contexts) > 1:
cb75ec7 @bdarnell Use our own copy of contextlib.nested (which no longer exists in py3)
bdarnell authored
179 with _nested(*new_contexts):
9d4444e @bdarnell Reindent everything to 4 spaces.
bdarnell authored
180 callback(*args, **kwargs)
6e7c8b9 @bdarnell Avoid use of the @contextlib.contextmanager decorator.
bdarnell authored
181 elif new_contexts:
182 with new_contexts[0]:
183 callback(*args, **kwargs)
9d4444e @bdarnell Reindent everything to 4 spaces.
bdarnell authored
184 else:
185 callback(*args, **kwargs)
9c17711 @bdarnell Create the StackContext in @asynchronous instead of on all requests,
bdarnell authored
186 if _state.contexts:
187 return _StackContextWrapper(wrapped, fn, _state.contexts)
188 else:
189 return _StackContextWrapper(fn)
6e7c8b9 @bdarnell Avoid use of the @contextlib.contextmanager decorator.
bdarnell authored
190
cb75ec7 @bdarnell Use our own copy of contextlib.nested (which no longer exists in py3)
bdarnell authored
191 @contextlib.contextmanager
192 def _nested(*managers):
193 """Support multiple context managers in a single with-statement.
194
195 Copied from the python 2.6 standard library. It's no longer present
196 in python 3 because the with statement natively supports multiple
197 context managers, but that doesn't help if the list of context
198 managers is not known until runtime.
199 """
200 exits = []
201 vars = []
202 exc = (None, None, None)
203 try:
204 for mgr in managers:
205 exit = mgr.__exit__
206 enter = mgr.__enter__
207 vars.append(enter())
208 exits.append(exit)
209 yield vars
210 except:
211 exc = sys.exc_info()
212 finally:
213 while exits:
214 exit = exits.pop()
215 try:
216 if exit(*exc):
217 exc = (None, None, None)
218 except:
219 exc = sys.exc_info()
220 if exc != (None, None, None):
221 # Don't rely on sys.exc_info() still containing
222 # the right information. Another exception may
223 # have been raised and caught by an exit method
224 raise exc[0], exc[1], exc[2]
225
Something went wrong with that request. Please try again.