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

@@ -5,10 +5,11 @@
 
 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
 
@@ -91,6 +92,31 @@ def validate_decimal_places(value: Decimal, decimal_places: int, name: str) -> N
         ) from exc
 
 
+async def grpc_call_with_timeout(
+    call: Callable[..., Awaitable[Any]], *args: Any, timeout: float = 300, **kwargs: Any
+) -> Any:
+    """
+    Call a gRPC function with a timeout (in seconds).
+
+    Args:
+        call: The gRPC method to be called.
+        *args: Positional arguments for the gRPC call.
+        timeout: Timeout duration in seconds. Defaults to 300.
+        **kwargs: Keyword arguments for the gRPC call.
+
+    Returns:
+        The result of the gRPC call.
+
+    Raises:
+        asyncio.TimeoutError: If the call exceeds the timeout.
+    """
+    try:
+        return await asyncio.wait_for(call(*args, **kwargs), timeout=timeout)
+    except asyncio.TimeoutError:
+        _logger.error("Timeout while calling %s", call)
+        raise
+
+
 class Client(BaseApiClient[ElectricityTradingServiceStub]):
     """Electricity trading client."""
 
@@ -503,10 +529,10 @@ async def create_gridpool_order(
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.CreateGridpoolOrderResponse],
-                self.stub.CreateGridpoolOrder(
+                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,
                 ),
@@ -615,7 +641,8 @@ async def update_gridpool_order(
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.UpdateGridpoolOrderResponse],
-                self.stub.UpdateGridpoolOrder(
+                grpc_call_with_timeout(
+                    self.stub.UpdateGridpoolOrder,
                     electricity_trading_pb2.UpdateGridpoolOrderRequest(
                         gridpool_id=gridpool_id,
                         order_id=order_id,
@@ -650,7 +677,8 @@ async def cancel_gridpool_order(
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.CancelGridpoolOrderResponse],
-                self.stub.CancelGridpoolOrder(
+                grpc_call_with_timeout(
+                    self.stub.CancelGridpoolOrder,
                     electricity_trading_pb2.CancelGridpoolOrderRequest(
                         gridpool_id=gridpool_id, order_id=order_id
                     ),
@@ -678,7 +706,8 @@ async def cancel_all_gridpool_orders(self, gridpool_id: int) -> int:
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.CancelAllGridpoolOrdersResponse],
-                self.stub.CancelAllGridpoolOrders(
+                grpc_call_with_timeout(
+                    self.stub.CancelAllGridpoolOrders,
                     electricity_trading_pb2.CancelAllGridpoolOrdersRequest(
                         gridpool_id=gridpool_id
                     ),
@@ -710,7 +739,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(
+                grpc_call_with_timeout(
+                    self.stub.GetGridpoolOrder,
                     electricity_trading_pb2.GetGridpoolOrderRequest(
                         gridpool_id=gridpool_id, order_id=order_id
                     ),
@@ -770,7 +800,8 @@ async def list_gridpool_orders(
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.ListGridpoolOrdersResponse],
-                self.stub.ListGridpoolOrders(
+                grpc_call_with_timeout(
+                    self.stub.ListGridpoolOrders,
                     electricity_trading_pb2.ListGridpoolOrdersRequest(
                         gridpool_id=gridpool_id,
                         filter=gridpool_order_filer.to_pb(),
@@ -842,7 +873,8 @@ async def list_gridpool_trades(
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.ListGridpoolTradesResponse],
-                self.stub.ListGridpoolTrades(
+                grpc_call_with_timeout(
+                    self.stub.ListGridpoolTrades,
                     electricity_trading_pb2.ListGridpoolTradesRequest(
                         gridpool_id=gridpool_id,
                         filter=gridpool_trade_filter.to_pb(),
@@ -899,7 +931,8 @@ async def list_public_trades(
         try:
             response = await cast(
                 Awaitable[electricity_trading_pb2.ListPublicTradesResponse],
-                self.stub.ListPublicTrades(
+                grpc_call_with_timeout(
+                    self.stub.ListPublicTrades,
                     electricity_trading_pb2.ListPublicTradesRequest(
                         filter=public_trade_filter.to_pb(),
                         pagination_params=pagination_params.to_proto(),