-
Notifications
You must be signed in to change notification settings - Fork 2
Expand file tree
/
Copy pathwait.py
More file actions
354 lines (308 loc) · 11.4 KB
/
wait.py
File metadata and controls
354 lines (308 loc) · 11.4 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
"""
This file implements waiting according to the "wait protocol"
The wait protocol simply means implementing the "add_done_callback"
A WaitSite class is provided, as well as a WaitableTasklet and
ValueTasklet.
Objects that conform to the wait protocol also include stacklesslib.locks.Event and
stacklesslib.futures.future.
"""
import sys
import traceback
import itertools
import stackless
import stacklesslib.util
from stacklesslib.base import atomic
from stacklesslib.errors import TimeoutError, CancelledError
from stacklesslib.weakmethod import WeakMethodProxy
class WaitSite(object):
"""
This class implements the interface for a waitable object, the
'add_done_callback()' and 'remove_done_callback()' methods.
"""
def __init__(self):
self._callbacks = []
def waitsite_signalled(self):
"""
This function is consulted when a callback is added and if it returns
True, an immediate callback is performed, rather than scheduling one later.
If a waitsite is stateful, this should be redefined to reflect its state.
"""
return False
def waitsite_signal(self):
callbacks, self._callbacks = self._callbacks, []
for cb in callbacks[:]:
self._cb(cb)
# This function is synonomous with the function in futures
def add_done_callback(self, cb):
"""
Add a callback. May result in an immediate call of "cb(self)"
if object is already signalled.
"""
# immediate callback if already signalled
self._callbacks.append(cb)
if self.waitsite_signalled():
self._cb(cb)
def remove_done_callback(self, cb):
"""
Remove the callback previously registered with 'add_done_callback()'
"""
try:
self._callbacks.remove(cb)
except ValueError:
pass
def _cb(self, cb):
try:
cb(self)
except Exception:
self.handle_exception(sys.exc_info())
def handle_exception(self, ei):
"""Handle an exception in the callback. Can be overridden."""
traceback.print_exception(*ei)
class Observer(WaitSite):
"""
An observer observes another waitable. With subclassing this can be used
to add state logic. This class acts as a proxy.
"""
def __init__(self, observed):
super(Observer, self).__init__()
self.got_callback = False
self.cb = WeakMethodProxy(self.callback)
# We use a weak method here, so that the callback doesn't keep us alive.
self.observed = observed
observed.add_done_callback(self.cb)
def __del__(self):
self.close()
def close(self):
if self.observed is not None:
self.observed.remove_done_callback(self.cb)
self.observed = None
def callback(self, target):
self.got_callback = True
if self.observed is not None:
if self.filter():
self.waitsite_signal()
def waitsite_signalled(self):
# A callback on add only occurs if a callback already occurred and filter approves
return self.got_callback and self.filter()
def filter(self):
"""
Decides if a callback from the observed should cause a callback
to any waiters.
"""
return True
class WaitableTasklet(stackless.tasklet, WaitSite):
"""A tasklet that implements the wait interface"""
def __new__(cls, *args, **kwds):
# Compatibility with old version of stackless that required
# that no extra arguments were passed to __new__
return stackless.tasklet.__new__(cls)
def __init__(self, func, args=None, kwargs=None):
WaitSite.__init__(self)
self._done = False
if func or args or kwargs:
self.bind(func, args, kwargs)
def bind(self, func, args=None, kwargs=None):
f = func
if func:
def helper(*args, **kwargs):
try:
return func(*args, **kwargs)
finally:
self._done = True
self.waitsite_signal()
f = helper
# compatibility with old stackless
if args or kwargs:
super(WaitableTasklet, self).bind(f, args, kwargs)
else:
super(WaitableTasklet, self).bind(f)
def waitsite_signalled(self):
return self._done
def join(self, timeout=None):
"""
Wait for the tasklet to finish. Returns True unless the timeout expired.
"""
return bool(wait([self], timeout))
class ValueTasklet(WaitableTasklet):
"""A WaitableTasklet that holds on to its return value or exit exception"""
def __init__(self, func, args=None, kwargs=None):
super(ValueTasklet, self).__init__(func, args, kwargs)
self._value = None
self._error = None
def bind(self, func, args=None, kwargs=None):
f = func
if func:
def helper(*args, **kwargs):
try:
self._value = func(*args, **kwargs)
except TaskletExit as e:
# Convert TaskletExit to CancelledError
self._error = CancelledError, CancelledError(e), sys.exc_info()[2] # considered a success
except:
self._error = sys.exc_info()
f = helper
super(ValueTasklet, self).bind(f, args, kwargs)
def get(self, block=True, timeout=None):
"""
Get the result of the tasklet or raise exception
"""
with atomic():
if not self._done and block:
self.join(timeout)
if not self._done:
raise TimeoutError
if self._error:
raise self._error[0], self._error[1], self._error[2]
return self._value
def result(self, timeout=None):
"""
A function conforming to the futures interface
"""
return self.get(timeout=timeout)
class WaitChannelMixin(WaitSite, stacklesslib.util.ChannelSequenceMixin):
"""
A mixin for channels to make them waitable. A wait will succeed whenever
the channel changes state from being blocked to unblocked, in whichever
direction.
"""
def __init__(self):
super(WaitChannelMixin, self).__init__()
# We must temporarily lie about the balance value while making
# wait callback calls.
self._balance_adjust = 0
@property
def balance(self):
return super(WaitChannelMixin, self).balance + self._balance_adjust
def _notify_state_change(self, adjust):
"""
Called when a channel is about to be receivable. Should be called
while the current tasklet is in the atomic mode.
"""
if super(WaitChannelMixin, self).balance == 0:
# The channel is becoming receivable or sendable. Let any interested
# parties know
self._balance_adjust += adjust
try:
self.waitsite_signal()
finally:
self._balance_adjust -= adjust
def send(self, value):
with atomic():
self._notify_state_change(1)
super(WaitChannelMixin, self).send(value)
def send_exception(self, exc, *args):
with atomic():
self._notify_state_change(1)
super(WaitChannelMixin, self).send_exception(exc, *args)
def send_throw(self, exc, value=None, tb=None):
with atomic():
self._notify_state_change(1)
super(WaitChannelMixin, self).send_throw(exc, value, tb)
def receive(self):
with atomic():
self._notify_state_change(-1)
return super(WaitChannelMixin, self).receive()
class Receivable(Observer):
def filter(self):
return self.observed.balance > 0
class Sendable(Observer):
def filter(self):
return self.observed.balance < 0
class WaitChannel(WaitChannelMixin, stackless.channel):
"""
A channel that provides wait semantics for receive
"""
def iwait(objects, timeout=None, raise_timeout=False):
"""
A generator that returns objects as they become ready. The total waiting
time will not exceed "timeout" if provided. Raises TimeoutError if a timeout
occurs.
"""
channel = stacklesslib.util.QueueChannel()
count = 0
callbacks = {}
def get_cb(obj):
def cb(waitable):
waitable.remove_done_callback(callbacks.pop(waitable))
channel.send(waitable)
callbacks[obj] = cb
return cb
try:
with atomic():
for obj in objects:
obj.add_done_callback(get_cb(obj))
count += 1
if timeout is not None:
# handle 0 timeouts by not blocking at all
if timeout == 0:
with atomic():
while channel.balance > 0:
count -= 1
yield channel.receive()
if count:
raise TimeoutError()
else:
timeouts = stacklesslib.util.Timeouts(timeout)
for i in xrange(count):
with timeouts.timeout():
yield channel.receive()
else:
for i in xrange(count):
yield channel.receive()
finally:
for obj, cb in callbacks.items():
obj.remove_done_callback(cb)
def iwait_no_raise(objects, timeout=None):
"""
Like 'iwait' but doesn't raise a TimeoutError, only stops returning results
when a timeout occurs.
"""
try:
for i in iwait(objects, timeout):
yield i
except TimeoutError:
pass
def wait(objects, timeout=None, count=None):
"""
Wait for objects, for at most "timeout" seconds or until "count" objects
are ready. Returns a list containing the ready objects in the order they became ready.
"""
if count is None:
return list(iwait_no_raise(objects, timeout))
return list(itertools.islice(iwait_no_raise(objects, timeout), count))
def swait(waitable, timeout=None):
"""
A simple wait function to wait for a single waitable. Returns the waitable
or raises TimeoutError.
"""
channel = stacklesslib.util.QueueChannel()
with atomic():
waitable.add_done_callback(channel.send)
try:
with stacklesslib.util.timeout(timeout):
return channel.receive()
finally:
waitable.remove_done_callback(channel.send)
def any(iterable):
"""
Returns a waitable object which is done when any of the waitables
in iterable is ready. its "result" method will return the waitable
that was ready
"""
def any_func():
for i in iwait(iterable):
return i
t = ValueTasklet(any_func)()
t.run()
return t
def all(iterable):
"""
Returns a waitable object which is done when all of the waitables
in iterable are ready. its "result" method returns the waitables
in the order in which they became ready.
"""
def all_func():
return list(iwait(iterable))
t = ValueTasklet(all_func)()
t.run()
return t