feat: wait for actions using ActionsClient.wait_for

This function allows the users to wait for multiple actions in an efficient way. All actions are queried using a single call, which reduce the potential for running into rate limits.

In addition, users may also configure a duration based timeout when waiting for actions using:

    action.wait_until_finished(timeout=10) # 10 seconds
    # or
    client.actions.wait_for(..., timeout=10)

Author: jooola

hcloud/_client.py

@@ -145,6 +145,7 @@ def __init__(
         application_version: str | None = None,
         poll_interval: int | float | BackoffFunction = 1.0,
         poll_max_retries: int = 120,
+        poll_timeout: float | None = None,
         timeout: float | tuple[float, float] | None = None,
         *,
         api_endpoint_hetzner: str = "https://api.hetzner.com/v1",
@@ -161,6 +162,8 @@ def __init__(
             You may pass a function to compute a custom poll interval.
         :param poll_max_retries:
             Max retries before timeout when polling actions from the API.
+        :param poll_timeout:
+            Duration in seconds before timeout when polling actions from the API.
         :param timeout: Requests timeout in seconds
         """
         self._client = ClientBase(
@@ -170,6 +173,7 @@ def __init__(
             application_version=application_version,
             poll_interval=poll_interval,
             poll_max_retries=poll_max_retries,
+            poll_timeout=poll_timeout,
             timeout=timeout,
         )
         self._client_hetzner = ClientBase(
@@ -179,6 +183,7 @@ def __init__(
             application_version=application_version,
             poll_interval=poll_interval,
             poll_max_retries=poll_max_retries,
+            poll_timeout=poll_timeout,
             timeout=timeout,
         )
 
@@ -336,6 +341,7 @@ def __init__(
         application_version: str | None = None,
         poll_interval: int | float | BackoffFunction = 1.0,
         poll_max_retries: int = 120,
+        poll_timeout: float | None = None,
         timeout: float | tuple[float, float] | None = None,
     ):
         self._token = token
@@ -355,6 +361,7 @@ def __init__(
 
         self._poll_interval_func = poll_interval_func
         self._poll_max_retries = poll_max_retries
+        self._poll_timeout = poll_timeout
 
         self._retry_interval_func = exponential_backoff_function(
             base=1.0, multiplier=2, cap=60.0, jitter=True

hcloud/_utils.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import time
+from collections.abc import Callable, Iterable, Iterator
+from itertools import islice
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+def batched(iterable: Iterable[T], size: int) -> Iterator[list[T]]:
+    """
+    Returns a batch of the provided size from the provided iterable.
+    """
+    iterator = iter(iterable)
+    while True:
+        batch = list(islice(iterator, size))
+        if not batch:
+            break
+        yield batch
+
+
+def waiter(timeout: float | None = None) -> Callable[[float], bool]:
+    """
+    Waiter returns a wait function that sleeps the specified amount of seconds, and
+    handles timeouts.
+
+    The wait function returns True if the timeout was reached, False otherwise.
+
+    :param timeout: Timeout in seconds, defaults to None.
+    :return: Wait function.
+    """
+
+    if timeout:
+        deadline = time.time() + timeout
+
+        def wait(seconds: float) -> bool:
+            now = time.time()
+
+            # Timeout if the deadline exceeded.
+            if deadline < now:
+                return True
+
+            # The deadline is not exceeded after the sleep time.
+            if now + seconds < deadline:
+                sleep(seconds)
+                return False
+
+            # The deadline is exceeded after the sleep time, clamp sleep time to
+            # deadline, and allow one last attempt until next wait call.
+            sleep(deadline - now)
+            return False
+
+    else:
+
+        def wait(seconds: float) -> bool:
+            sleep(seconds)
+            return False
+
+    return wait
+
+
+def sleep(seconds: float) -> None:
+    """
+    An interruptable sleep function that does not lock the entire thread.
+
+    :param seconds: Seconds to sleep.
+    """
+    if seconds < 1:
+        time.sleep(seconds)
+    else:
+        for _ in range(int(seconds)):
+            time.sleep(1)
+        time.sleep(seconds - int(seconds))

hcloud/actions/client.py

@@ -1,11 +1,17 @@
 from __future__ import annotations
 
-import time
 import warnings
+from collections.abc import Callable
 from typing import TYPE_CHECKING, Any, Literal, NamedTuple
 
+from .._utils import batched, waiter
 from ..core import BoundModelBase, Meta, ResourceClientBase
-from .domain import Action, ActionFailedException, ActionStatus, ActionTimeoutException
+from .domain import (
+    Action,
+    ActionFailedException,
+    ActionStatus,
+    ActionTimeoutException,
+)
 
 if TYPE_CHECKING:
     from .._client import Client
@@ -25,33 +31,44 @@ class BoundAction(BoundModelBase[Action], Action):
 
     model = Action
 
-    def wait_until_finished(self, max_retries: int | None = None) -> None:
-        """Wait until the specific action has status=finished.
-
-        :param max_retries: int Specify how many retries will be performed before an ActionTimeoutException will be raised.
-        :raises: ActionFailedException when action is finished with status==error
-        :raises: ActionTimeoutException when Action is still in status==running after max_retries is reached.
+    def wait_until_finished(
+        self,
+        max_retries: int | None = None,
+        *,
+        timeout: float | None = None,
+    ) -> None:
         """
-        if max_retries is None:
-            # pylint: disable=protected-access
-            max_retries = self._client._client._poll_max_retries
+        Waits until the Action is finished by polling the API at the interval defined by
+        the client's poll interval and function. An Action is considered as finished
+        when its status is either "success" or "error".
 
-        retries = 0
-        while True:
-            self.reload()
-            if self.status != Action.STATUS_RUNNING:
-                break
+        If the Action fails (its status is "error"), the function will stop waiting
+        and raise ActionFailedException.
 
-            retries += 1
-            if retries < max_retries:
-                # pylint: disable=protected-access
-                time.sleep(self._client._client._poll_interval_func(retries))
-                continue
+        :param timeout:
+            Duration in seconds before an ActionTimeoutException will be raised when polling actions from the API.
+        :param max_retries:
+            Max retries before an ActionTimeoutException will be raised when polling actions from the API.
+
+        :raises: ActionTimeoutException when an Action is still running after max_retries or timeout is reached.
+        :raises: ActionFailedException when an Action failed.
+        """
+
+        def handle_update(update: BoundAction) -> None:
+            self.data_model = update.data_model
 
-            raise ActionTimeoutException(action=self)
+            if update.status == Action.STATUS_ERROR:
+                raise ActionFailedException(action=update)
 
-        if self.status == Action.STATUS_ERROR:
-            raise ActionFailedException(action=self)
+        try:
+            self._client.wait_for_function(
+                handle_update,
+                [self],
+                timeout=timeout,
+                max_retries=max_retries,
+            )
+        except* ActionTimeoutException as group:
+            raise group.exceptions[0]
 
 
 ActionSort = Literal[
@@ -189,6 +206,148 @@ class ActionsClient(ResourceActionsClient):
     def __init__(self, client: Client):
         super().__init__(client, None)
 
+    def _get_list_by_ids(self, ids: list[int]) -> list[BoundAction]:
+        """
+        Get a list of Actions by their IDs.
+
+        :param ids: List of Action IDs to get.
+        :raises ValueError: Raise when Action IDs were not found.
+        :return: List of Actions.
+        """
+        actions: list[BoundAction] = []
+
+        for ids_batch in batched(ids, 25):
+            params: dict[str, Any] = {
+                "id": ids_batch,
+                "sort": ["status", "id"],
+            }
+
+            response = self._client.request(
+                method="GET",
+                url="/actions",
+                params=params,
+            )
+
+            actions.extend(
+                BoundAction(self._parent.actions, o) for o in response["actions"]
+            )
+
+        if len(ids) != len(actions):
+            found_ids = [a.id for a in actions]
+            not_found_ids = list(set(ids) - set(found_ids))
+
+            raise ValueError(
+                f"actions not found: {', '.join(str(o) for o in not_found_ids)}"
+            )
+
+        return actions
+
+    def wait_for_function(
+        self,
+        handle_update: Callable[[BoundAction], None],
+        actions: list[Action | BoundAction],
+        *,
+        timeout: float | None = None,
+        max_retries: int | None = None,
+    ) -> list[BoundAction]:
+        """
+        Waits until all Actions are finished by polling the API at the interval defined
+        by the client's poll interval and function. An Action is considered as finished
+        when its status is either "success" or "error".
+
+        The handle_update callback is called every time an Action is updated.
+
+        :param handle_update:
+            Function called every time an Action is updated.
+        :param actions:
+            List of Actions to wait for.
+        :param timeout:
+            Duration in seconds before an ActionTimeoutException will be raised when polling actions from the API.
+        :param max_retries:
+            Max retries before an ActionTimeoutException will be raised when polling actions from the API.
+
+        :raises: ActionTimeoutException when an Action is still running after max_retries or timeout is reached.
+
+        :return: List of finished Actions.
+        """
+        if timeout is None:
+            # pylint: disable=protected-access
+            timeout = self._client._poll_timeout
+        if max_retries is None:
+            # pylint: disable=protected-access
+            max_retries = self._client._poll_max_retries
+
+        running: list[BoundAction] = actions.copy()  # type: ignore[assignment]
+        completed: list[BoundAction] = []
+
+        retries = 0
+        wait = waiter(timeout)
+        while len(running) > 0:
+            if max_retries is not None and retries > max_retries:
+                raise ExceptionGroup(
+                    "The actions timed out after",
+                    [ActionTimeoutException(action) for action in running],
+                )
+
+            # pylint: disable=protected-access
+            if wait(self._client._poll_interval_func(retries)):
+                raise ExceptionGroup(
+                    "The actions timed out",
+                    [ActionTimeoutException(action) for action in running],
+                )
+
+            retries += 1
+
+            running = self._get_list_by_ids([a.id for a in running])
+
+            for update in running:
+                if update.status != Action.STATUS_RUNNING:
+                    running.remove(update)
+                    completed.append(update)
+
+                handle_update(update)
+
+        return completed
+
+    def wait_for(
+        self,
+        actions: list[Action | BoundAction],
+        *,
+        timeout: float | None = None,
+        max_retries: int | None = None,
+    ) -> list[BoundAction]:
+        """
+        Waits until all Actions are finished by polling the API at the interval defined
+        by the client's poll interval and function. An Action is considered as finished
+        when its status is either "success" or "error".
+
+        If a single Action fails (its status is "error"), the function will stop waiting
+        and raise ActionFailedException.
+
+        :param actions:
+            List of Actions to wait for.
+        :param timeout:
+            Duration in seconds before an ActionTimeoutException will be raised when polling actions from the API.
+        :param max_retries:
+            Max retries before an ActionTimeoutException will be raised when polling actions from the API.
+
+        :raises: ActionTimeoutException when an Action is still running after max_retries or timeout is reached.
+        :raises: ActionFailedException when an Action failed.
+
+        :return: List of succeeded Actions.
+        """
+
+        def handle_update(update: BoundAction) -> None:
+            if update.status == Action.STATUS_ERROR:
+                raise ActionFailedException(action=update)
+
+        return self.wait_for_function(
+            handle_update,
+            actions,
+            timeout=timeout,
+            max_retries=max_retries,
+        )
+
     def get_list(
         self,
         status: list[ActionStatus] | None = None,