Add global start timeout (#138)

Author: davidbrochart

src/fps/_module.py

@@ -6,6 +6,7 @@
 from collections.abc import Callable, Awaitable
 from contextlib import AsyncExitStack
 from inspect import isawaitable, signature, _empty
+from time import time
 from typing import TypeVar, Any, Iterable, cast
 
 import anyio
@@ -56,18 +57,22 @@ def __init__(
         prepare_timeout: float = 1,
         start_timeout: float = 1,
         stop_timeout: float = 1,
+        global_start_timeout: float | None = None,
     ):
         """
         Args:
             name: The name to give to the module.
             prepare_timeout: The time to wait (in seconds) for the "prepare" phase to complete.
             start_timeout: The time to wait (in seconds) for the "start" phase to complete.
             stop_timeout: The time to wait (in seconds) for the "stop" phase to complete.
+            global_start_timeout: The time to wait (in seconds) for the "prepare" and "start"
+                phases to complete.
         """
         self._initialized = False
         self._prepare_timeout = prepare_timeout
         self._start_timeout = start_timeout
         self._stop_timeout = stop_timeout
+        self._global_start_timeout = global_start_timeout
         self._parent: Module | None = None
         self._context = Context()
         self._prepared = Event()
@@ -327,7 +332,12 @@ async def __aenter__(self) -> Module:
             self._task_group = await exit_stack.enter_async_context(create_task_group())
             self._exceptions = []
             self._phase = "preparing"
-            with move_on_after(self._prepare_timeout) as scope:
+            if self._global_start_timeout is None:
+                prepare_timeout = self._prepare_timeout
+            else:
+                prepare_timeout = self._global_start_timeout
+            t0 = time()
+            with move_on_after(prepare_timeout) as scope:
                 self._task_group.start_soon(self._prepare, name=f"{self.path} _prepare")
                 await self._all_prepared()
             if scope.cancelled_caught:
@@ -336,7 +346,12 @@ async def __aenter__(self) -> Module:
                 self._exit.set()
             else:
                 self._phase = "starting"
-                with move_on_after(self._start_timeout) as scope:
+                if self._global_start_timeout is None:
+                    start_timeout = self._start_timeout
+                else:
+                    elapsed_time = time() - t0
+                    start_timeout = max(self._global_start_timeout - elapsed_time, 0)
+                with move_on_after(start_timeout) as scope:
                     self._task_group.start_soon(self._start, name=f"{self.path} start")
                     await self._all_started()
                 if scope.cancelled_caught:

src/fps/cli/_cli.py

@@ -43,6 +43,19 @@
     default="asyncio",
     help="The name of the event loop to use (asyncio or trio).",
 )
+@click.option(
+    "--timeout",
+    type=float,
+    default=None,
+    help="The timeout for starting the module (in seconds).",
+)
+@click.option(
+    "--stop-timeout",
+    type=float,
+    show_default=True,
+    default=1,
+    help="The timeout for stopping the module (in seconds).",
+)
 @click.argument("module", default="")
 def main(
     module: str,
@@ -51,6 +64,8 @@ def main(
     help_all: bool = False,
     set_: list[str] | None = None,
     backend: str = "asyncio",
+    timeout: float | None = None,
+    stop_timeout: float = 1,
 ):
     global CONFIG
     if config is None:
@@ -87,6 +102,8 @@ def main(
         CONFIG = config_dict
         return
     root_module = get_root_module(config_dict)
+    root_module._global_start_timeout = timeout
+    root_module._stop_timeout = stop_timeout
     actual_config = initialize(root_module)
     if help_all:
         click.echo(get_config_description(root_module))

tests/test_start_stop.py

@@ -108,3 +108,30 @@ def __init__(self, name, stop_timeout):
         str(module0.exceptions[0])
         == "Module timed out while stopping: module0.submodule0"
     )
+
+
+async def test_global_start_timeout_in_prepare():
+    class Module0(Module):
+        async def prepare(self):
+            await sleep(1)
+
+    async with Module0("module0", global_start_timeout=0.1) as module0:
+        pass
+
+    assert len(module0.exceptions) == 1
+    assert str(module0.exceptions[0]) == "Module timed out while preparing: module0"
+
+
+async def test_global_start_timeout_in_start():
+    class Module0(Module):
+        async def prepare(self):
+            await sleep(0.1)
+
+        async def start(self):
+            await sleep(1)
+
+    async with Module0("module0", global_start_timeout=0.2) as module0:
+        pass
+
+    assert len(module0.exceptions) == 1
+    assert str(module0.exceptions[0]) == "Module timed out while starting: module0"