Mirror of https://github.com/roostorg/osprey
github.com/roostorg/osprey
Add configurable timeouts to output sinks (#88)
+8
-2
+5
-2
osprey_worker/src/osprey/worker/sinks/sink/output_sink.py
+5
-2
osprey_worker/src/osprey/worker/sinks/sink/output_sink.py
···
19
19
20
20
logger = get_logger()
21
21
22
-
GEVENT_TIMEOUT = 2
22
+
DEFAULT_GEVENT_TIMEOUT = 2
23
23
24
24
25
25
class BaseOutputSink(abc.ABC):
26
+
# Default timeout for sink operations. Subclasses can override this.
27
+
timeout: float = DEFAULT_GEVENT_TIMEOUT
28
+
26
29
@abc.abstractmethod
27
30
def will_do_work(self, result: ExecutionResult) -> bool:
28
31
"""A quick way to determine if this sink needs to do anything for this result."""
···
59
62
with (
60
63
trace(f'{sink_name}.push'),
61
64
metrics.timed('handled_message_output', tags=[f'sink:{sink_name}'], use_ms=True),
62
-
gevent.Timeout(GEVENT_TIMEOUT),
65
+
gevent.Timeout(sink.timeout),
63
66
):
64
67
sink.push(result)
65
68
except gevent.Timeout as timeout_exc:
+3
osprey_worker/src/osprey/worker/sinks/sink/stored_execution_result_output_sink.py
+3
osprey_worker/src/osprey/worker/sinks/sink/stored_execution_result_output_sink.py
···
6
6
class StoredExecutionResultOutputSink(BaseOutputSink):
7
7
"""An output sink that persists the execution result to an EventRecord."""
8
8
9
+
# BigTable writes can take longer than the default 2s, especially on first connection
10
+
timeout: float = 10.0
11
+
9
12
def __init__(self):
10
13
self._service = bootstrap_execution_result_storage_service()
11
14