Add a timeout in the gRPC function calls

Signed-off-by: camille-bouvy-frequenz <[email protected]>

Author: camille-bouvy-frequenz

src/frequenz/client/electricity_trading/_client.py

@@ -3,12 +3,15 @@
 
 """Module to define the client class."""
 
+# pylint: disable=too-many-lines
+
 from __future__ import annotations
 
+import asyncio
 import logging
 from datetime import datetime, timezone
 from decimal import Decimal, InvalidOperation
-from typing import Awaitable, cast
+from typing import Any, Awaitable, Callable, cast
 
 import grpc
 
@@ -87,7 +90,9 @@ def validate_decimal_places(value: Decimal, decimal_places: int, name: str) -> N
         ) from exc
 
 
-class Client(BaseApiClient[ElectricityTradingServiceStub]):
+class Client(
+    BaseApiClient[ElectricityTradingServiceStub]
+):  # pylint:disable=too-many-instance-attributes
     """Electricity trading client."""
 
     _instances: dict[tuple[str, str | None], "Client"] = {}
@@ -116,15 +121,28 @@ def __new__(
 
         return cls._instances[key]
 
-    def __init__(
-        self, server_url: str, connect: bool = True, auth_key: str | None = None
+    def __init__(  # pylint: disable=too-many-arguments, too-many-positional-arguments
+        self,
+        server_url: str,
+        connect: bool = True,
+        auth_key: str | None = None,
+        initial_timeout: int = 10,
+        max_timeout: int = 300,
+        timeout_increment: int = 20,
+        max_timeout_retries: int = 30,
+        reset_interval: int = 600,
     ) -> None:
         """Initialize the client.
 
         Args:
             server_url: The URL of the Electricity Trading service.
             connect: Whether to connect to the server immediately.
             auth_key: The API key for the authorization.
+            initial_timeout: Initial timeout duration for gRPC calls (in seconds).
+            max_timeout: Maximum timeout duration for gRPC calls (in seconds).
+            timeout_increment: Increment value for timeout on each retry (in seconds).
+            max_timeout_retries: Maximum number of retry attempts when a timeout is reached.
+            reset_interval: Time (in seconds) before resetting timeout if everything behaves fine.
         """
         if not hasattr(
             self, "_initialized"
@@ -159,6 +177,87 @@ def __init__(
 
         self._metadata = (("key", auth_key),) if auth_key else ()
 
+        # Timeout configuration
+        self._initial_timeout = initial_timeout
+        self._timeout = initial_timeout
+        self._timeout_increment = timeout_increment
+        self._max_timeout = max_timeout
+        self._max_timeout_retries = max_timeout_retries
+        self._reset_interval = reset_interval
+        self._last_success_time = datetime.now(timezone.utc)
+
+    def _increase_timeout(self) -> None:
+        """Increase the global timeout within limits."""
+        if self._timeout < self._max_timeout:
+            self._timeout = min(
+                self._max_timeout, self._timeout + self._timeout_increment
+            )
+            _logger.info("Increased timeout to %ds.", self._timeout)
+        else:
+            _logger.warning(
+                "Timeout is already at the maximum (%ds). No further increments.",
+                self._max_timeout,
+            )
+
+    def _reset_timeout_if_needed(self) -> None:
+        """Reset the timeout to the initial value if the reset interval has passed."""
+        now = datetime.now(timezone.utc)
+        if (now - self._last_success_time).total_seconds() >= self._reset_interval:
+            self._timeout = self._initial_timeout
+            _logger.info(
+                "Timeout reset to initial value %ds after %ds seconds without timeout errors.",
+                self._initial_timeout,
+                self._reset_interval,
+            )
+
+    async def grpc_call_with_timeout(
+        self,
+        call: Callable[..., Awaitable[Any]],
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        """
+        Call a gRPC function with a timeout.
+
+        This timeout is increased dynamically on every none-successful reconnection attempts.
+
+        Args:
+            call: The gRPC method to be called.
+            *args: Positional arguments for the gRPC call.
+            **kwargs: Keyword arguments for the gRPC call.
+
+        Returns:
+            The result of the gRPC call.
+
+        Raises:
+            asyncio.TimeoutError: If all retries are exhausted.
+        """
+        for attempt in range(1, self._max_timeout_retries + 1):
+            try:
+                self._reset_timeout_if_needed()
+                result = await asyncio.wait_for(
+                    call(*args, **kwargs), timeout=self._timeout
+                )
+                # Update last success time on successful call
+                self._last_success_time = datetime.now(timezone.utc)
+                return result
+            except asyncio.TimeoutError:
+                if attempt == self._max_timeout_retries:
+                    _logger.error(
+                        "Timeout after %d retries (timeout=%ds): %s",
+                        self._max_timeout_retries,
+                        self._timeout,
+                        call,
+                    )
+                    raise
+                _logger.warning(
+                    "Timeout on attempt %d/%d (timeout=%ds). Retrying with increased timeout.",
+                    attempt,
+                    self._max_timeout_retries,
+                    self._timeout,
+                )
+                self._increase_timeout()
+
     @property
     def stub(self) -> electricity_trading_pb2_grpc.ElectricityTradingServiceAsyncStub:
         """
@@ -493,10 +592,10 @@ async def create_gridpool_order(
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.CreateGridpoolOrderResponse],
-                self.stub.CreateGridpoolOrder(
+                self.grpc_call_with_timeout(
+                    self.stub.CreateGridpoolOrder,
                     electricity_trading_pb2.CreateGridpoolOrderRequest(
-                        gridpool_id=gridpool_id,
-                        order=order.to_pb(),
+                        gridpool_id=gridpool_id, order=order.to_pb()
                     ),
                     metadata=self._metadata,
                 ),
@@ -605,7 +704,8 @@ async def update_gridpool_order(
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.UpdateGridpoolOrderResponse],
-                self.stub.UpdateGridpoolOrder(
+                self.grpc_call_with_timeout(
+                    self.stub.UpdateGridpoolOrder,
                     electricity_trading_pb2.UpdateGridpoolOrderRequest(
                         gridpool_id=gridpool_id,
                         order_id=order_id,
@@ -640,7 +740,8 @@ async def cancel_gridpool_order(
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.CancelGridpoolOrderResponse],
-                self.stub.CancelGridpoolOrder(
+                self.grpc_call_with_timeout(
+                    self.stub.CancelGridpoolOrder,
                     electricity_trading_pb2.CancelGridpoolOrderRequest(
                         gridpool_id=gridpool_id, order_id=order_id
                     ),
@@ -668,7 +769,8 @@ async def cancel_all_gridpool_orders(self, gridpool_id: int) -> int:
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.CancelAllGridpoolOrdersResponse],
-                self.stub.CancelAllGridpoolOrders(
+                self.grpc_call_with_timeout(
+                    self.stub.CancelAllGridpoolOrders,
                     electricity_trading_pb2.CancelAllGridpoolOrdersRequest(
                         gridpool_id=gridpool_id
                     ),
@@ -700,7 +802,8 @@ async def get_gridpool_order(self, gridpool_id: int, order_id: int) -> OrderDeta
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.GetGridpoolOrderResponse],
-                self.stub.GetGridpoolOrder(
+                self.grpc_call_with_timeout(
+                    self.stub.GetGridpoolOrder,
                     electricity_trading_pb2.GetGridpoolOrderRequest(
                         gridpool_id=gridpool_id, order_id=order_id
                     ),
@@ -760,7 +863,8 @@ async def list_gridpool_orders(
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.ListGridpoolOrdersResponse],
-                self.stub.ListGridpoolOrders(
+                self.grpc_call_with_timeout(
+                    self.stub.ListGridpoolOrders,
                     electricity_trading_pb2.ListGridpoolOrdersRequest(
                         gridpool_id=gridpool_id,
                         filter=gridpool_order_filer.to_pb(),
@@ -832,7 +936,8 @@ async def list_gridpool_trades(
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.ListGridpoolTradesResponse],
-                self.stub.ListGridpoolTrades(
+                self.grpc_call_with_timeout(
+                    self.stub.ListGridpoolTrades,
                     electricity_trading_pb2.ListGridpoolTradesRequest(
                         gridpool_id=gridpool_id,
                         filter=gridpool_trade_filter.to_pb(),
@@ -889,7 +994,8 @@ async def list_public_trades(
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.ListPublicTradesResponse],
-                self.stub.ListPublicTrades(
+                self.grpc_call_with_timeout(
+                    self.stub.ListPublicTrades,
                     electricity_trading_pb2.ListPublicTradesRequest(
                         filter=public_trade_filter.to_pb(),
                         pagination_params=pagination_params.to_proto(),