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
|
#!/usr/bin/env python3
# © Copyright 2021-2022, Scott Gasch
"""
A future that can be treated as a substutute for the result that it
contains and will not block until it is used. At that point, if the
underlying value is not yet available yet, it will block until the
internal result actually becomes available.
"""
from __future__ import annotations
import concurrent
import concurrent.futures as fut
import logging
from typing import Callable, List, Set, TypeVar
from overrides import overrides
import id_generator
# This module is commonly used by others in here and should avoid
# taking any unnecessary dependencies back on them.
from deferred_operand import DeferredOperand
logger = logging.getLogger(__name__)
T = TypeVar('T')
def wait_any(
futures: List[SmartFuture],
*,
callback: Callable = None,
log_exceptions: bool = True,
):
"""Await the completion of any of a collection of SmartFutures and
invoke callback each time one completes, repeatedly, until they are
all finished.
Args:
futures: A collection of SmartFutures to wait on
callback: An optional callback to invoke whenever one of the
futures completes
log_exceptions: Should we log (warning + exception) any
underlying exceptions raised during future processing or
silently ignore then?
"""
real_futures = []
smart_future_by_real_future = {}
completed_futures: Set[fut.Future] = set()
for x in futures:
assert isinstance(x, SmartFuture)
real_futures.append(x.wrapped_future)
smart_future_by_real_future[x.wrapped_future] = x
while len(completed_futures) != len(real_futures):
newly_completed_futures = concurrent.futures.as_completed(real_futures)
for f in newly_completed_futures:
if callback is not None:
callback()
completed_futures.add(f)
if log_exceptions and not f.cancelled():
exception = f.exception()
if exception is not None:
logger.warning('Future 0x%x raised an unhandled exception and exited.', id(f))
logger.exception(exception)
raise exception
yield smart_future_by_real_future[f]
if callback is not None:
callback()
def wait_all(
futures: List[SmartFuture],
*,
log_exceptions: bool = True,
) -> None:
"""Wait for all of the SmartFutures in the collection to finish before
returning.
Args:
futures: A collection of futures that we're waiting for
log_exceptions: Should we log (warning + exception) any
underlying exceptions raised during future processing or
silently ignore then?
"""
real_futures = []
for x in futures:
assert isinstance(x, SmartFuture)
real_futures.append(x.wrapped_future)
(done, not_done) = concurrent.futures.wait(
real_futures, timeout=None, return_when=concurrent.futures.ALL_COMPLETED
)
if log_exceptions:
for f in real_futures:
if not f.cancelled():
exception = f.exception()
if exception is not None:
logger.warning('Future 0x%x raised an unhandled exception and exited.', id(f))
logger.exception(exception)
raise exception
assert len(done) == len(real_futures)
assert len(not_done) == 0
class SmartFuture(DeferredOperand):
"""This is a SmartFuture, a class that wraps a normal :class:`Future`
and can then be used, mostly, like a normal (non-Future)
identifier of the type of that SmartFuture's result.
Using a FutureWrapper in expressions will block and wait until
the result of the deferred operation is known.
"""
def __init__(self, wrapped_future: fut.Future) -> None:
assert isinstance(wrapped_future, fut.Future)
self.wrapped_future = wrapped_future
self.id = id_generator.get("smart_future_id")
def get_id(self) -> int:
return self.id
def is_ready(self) -> bool:
return self.wrapped_future.done()
# You shouldn't have to call this; instead, have a look at defining a
# method on DeferredOperand base class.
@overrides
def _resolve(self, timeout=None) -> T:
return self.wrapped_future.result(timeout)
|
