feat(core): add idempotent flag on RPC API (#10)

The goal of this change is to continue retrying calls in the event
of network loss, while also allowing this feature to be disabled
for calls that are not idempotent and require more intelligence
on the caller side.

Signed-off-by: Ronan Abhamon <ronan.abhamon@vates.tech>

Author: Wescoeur

src/xcp_storage/rpc/client.py

@@ -74,6 +74,12 @@ def call_api(
         *args: P.args,
         **kwargs: P.kwargs
     ) -> JsonValueT:
+        """
+        Call the RPC API `method`. The call is retried if the connection is lost, unless the method is
+        declared with `idempotent=False` (see `JsonRpcDispatcher.method`).
+        See `JsonRpcClient.call_with_timeout` for more details.
+        """
+
         return self.call_api_with_timeout(self._client_timeout, method, *args, **kwargs)
 
     def call_api_with_timeout(
@@ -87,6 +93,7 @@ def call_api_with_timeout(
             method_name = cast(Any, method)._rpc_name # noqa: SLF001
         except AttributeError:
             raise JsonRpcRequestError("Method is not marked as RPC.") from None
+        idempotent = getattr(method, "_rpc_idempotent", True)
 
         if args and kwargs:
             raise JsonRpcRequestError("Positional and named arguments cannot be mixed.")
@@ -99,4 +106,4 @@ def call_api_with_timeout(
         else:
             params = None
 
-        return cast(JsonValueT, self.call_with_timeout(timeout, method_name, params))
+        return cast(JsonValueT, self.call_with_timeout(timeout, method_name, params, retry_on_disconnect=idempotent))

src/xcp_storage/utils/json/rpc/__init__.py

@@ -30,6 +30,7 @@
     Final,
     List,
     Optional,
+    overload,
     override,
     ParamSpec,
     TypeVar,
@@ -516,14 +517,42 @@ def call_method(self, method: str, *args: Any, **kwargs: Any) -> Any: # noqa: AN
 
         return JsonRpcCallResult(result=result)
 
+    @overload
     def method(self, func: Callable[P, T]) -> Callable[P, T]:
-        name = func.__name__
-        if self._use_module_name:
-            name = PurePath(inspect.getfile(func)).stem + "." + name
-        self._name_to_method[name] = func
+        ...
 
-        cast(Any, func)._rpc_name = name # noqa: SLF001
-        return func
+    @overload
+    def method(self, *, idempotent: bool = True) -> Callable[[Callable[P, T]], Callable[P, T]]:
+        ...
+
+    def method(
+        self,
+        func: Optional[Callable[P, T]] = None,
+        *,
+        idempotent: bool = True
+    ) -> Union[Callable[P, T], Callable[[Callable[P, T]], Callable[P, T]]]:
+        """
+        Register a function as RPC method.
+
+        `idempotent=False` declares a method that must not be executed twice: a client calling it
+        through the API never sends the request again after a lost connection
+        (see `JsonRpcClient.call_with_timeout`).
+        """
+
+        def register(target: Callable[P, T]) -> Callable[P, T]:
+            name = target.__name__
+            if self._use_module_name:
+                name = PurePath(inspect.getfile(target)).stem + "." + name
+            self._name_to_method[name] = target
+
+            target_info = cast(Any, target)
+            target_info._rpc_name = name # noqa: SLF001
+            target_info._rpc_idempotent = idempotent # noqa: SLF001
+            return target
+
+        if func is None:
+            return register
+        return register(func)
 
 # ------------------------------------------------------------------------------
 # Request processor.

src/xcp_storage/utils/json/rpc/client.py

@@ -117,23 +117,53 @@ def disconnect(self) -> None:
     def call(
         self,
         method: str,
-        params: Union[JsonList, JsonDict, None] = None
+        params: Union[JsonList, JsonDict, None] = None,
+        *,
+        retry_on_disconnect: bool = True
     ) -> JsonValue:
-        return self.call_with_timeout(self._client_timeout, method, params)
+        """
+        Call `method` with the client timeout. See `call_with_timeout` for the retry semantics.
+        """
+
+        return self.call_with_timeout(self._client_timeout, method, params, retry_on_disconnect=retry_on_disconnect)
 
     def call_with_timeout(
         self,
         timeout: float,
         method: str,
-        params: Union[JsonList, JsonDict, None] = None
+        params: Union[JsonList, JsonDict, None] = None,
+        *,
+        retry_on_disconnect: bool = True
     ) -> JsonValue:
+        """
+        Call `method` and return its result.
+
+        If the connection is lost (`SocketDisconnectedError`), the client reconnects and sends the
+        same request again, until `timeout` is elapsed. The delivery is therefore "at least once":
+        the failure can occur while waiting for the response, after the server has already executed
+        the request. RPC methods exposed through this client MUST be idempotent, or tolerate being
+        executed twice.
+
+        For the rare methods that must not be executed twice, use `retry_on_disconnect=False` (the
+        API client does it for methods declared with `idempotent=False`): the request is sent at most
+        once and `SocketDisconnectedError` is raised on a lost connection.
+        It is then up to the caller to determine if the method has been executed. Note that the
+        client still waits up to `timeout` for the initial connection to the server: nothing has
+        been sent at this point.
+
+        Other errors are not retried: a timeout (`SocketTimeoutError`) or an error response from
+        the server is raised as is.
+        """
+
         remaining_time = timeout
         start_time = time.monotonic()
 
         while True:
             try:
                 return self._call(remaining_time, method, params)
             except SocketDisconnectedError: # noqa: PERF203
+                if not retry_on_disconnect:
+                    raise
                 remaining_time = timeout - (time.monotonic() - start_time)
                 if remaining_time <= 0:
                     raise

tests/rpc/test_rpc_client_server.py

@@ -22,9 +22,11 @@
 from tests.rpc.conftest import CLIENT_TIMEOUT, RpcClientFactory
 from xcp_storage.network.protocol import Protocol, ProtocolError
 from xcp_storage.network.protocol.xcp import XcpProtocol
+from xcp_storage.network.socket import SocketDisconnectedError
 from xcp_storage.network.tcp_client import TcpClientError
 from xcp_storage.rpc.client import RpcApiClient
 from xcp_storage.rpc.server import RpcApiServer
+from xcp_storage.utils.json import JsonDict
 from xcp_storage.utils.json.rpc import JsonRpcRequestError
 
 from xcp_storage.typing import (
@@ -65,7 +67,10 @@ def test_connect_with_untrusted_certificate(self, rpc_server: RpcApiServer) -> N
 
 @over_plain_and_tls
 class TestRpcClientServer:
-    def test_retry_call_on_socket_disconnect(self, rpc_client_factory: RpcClientFactory) -> None:
+    @pytest.mark.parametrize("retry_on_disconnect", [True, False])
+    def test_retry_call_on_socket_disconnect(
+        self, rpc_client_factory: RpcClientFactory, *, retry_on_disconnect: bool
+    ) -> None:
         rpc_client = rpc_client_factory()
 
         rpc_client.connect()
@@ -76,8 +81,18 @@ def test_retry_call_on_socket_disconnect(self, rpc_client_factory: RpcClientFact
         tpc_socket.close()
 
         message = "Bonjour !"
-        response = rpc_client.call("echo.echo", params={"message": message})
-        assert response == message
+        params: JsonDict = {"message": message}
+        with patch.object(rpc_client, "_call", wraps=rpc_client._call) as call_spy: # noqa: SLF001
+            if retry_on_disconnect:
+                assert rpc_client.call("echo.echo", params=params) == message
+            else:
+                with pytest.raises(SocketDisconnectedError):
+                    rpc_client.call("echo.echo", params=params, retry_on_disconnect=False)
+
+        assert call_spy.call_count == (2 if retry_on_disconnect else 1)
+        assert rpc_client.connected == retry_on_disconnect
+
+        assert rpc_client.call("echo.echo", params=params) == message
         assert rpc_client.connected
 
     def test_rpc_packet_last_sequences(self, rpc_client_factory: RpcClientFactory) -> None:

tests/utils/test_json_rpc.py

@@ -23,7 +23,11 @@
     JsonRpcResponse,
 )
 
-from xcp_storage.typing import Union
+from xcp_storage.typing import (
+    Any,
+    cast,
+    Union,
+)
 
 # ==============================================================================
 
@@ -73,6 +77,35 @@ def assert_response_error(
 
 # ------------------------------------------------------------------------------
 
+class TestJsonRpcDispatcherMethod:
+    def test_method_is_idempotent_by_default(self) -> None:
+        dispatcher = JsonRpcDispatcher()
+
+        @dispatcher.method
+        def add(a: int, b: int) -> int:
+            return a + b
+
+        assert add(1, 2) == 3
+
+        method_info = cast(Any, add)
+        assert method_info._rpc_name == "add" # noqa: SLF001
+        assert method_info._rpc_idempotent # noqa: SLF001
+
+    def test_method_with_arguments(self) -> None:
+        dispatcher = JsonRpcDispatcher()
+
+        @dispatcher.method(idempotent=False)
+        def add(a: int, b: int) -> int:
+            return a + b
+
+        assert add(1, 2) == 3
+
+        method_info = cast(Any, add)
+        assert method_info._rpc_name == "add" # noqa: SLF001
+        assert not method_info._rpc_idempotent # noqa: SLF001
+
+# ------------------------------------------------------------------------------
+
 class TestJsonRpcRequest:
     def test_positional_params(self) -> None:
         dispatcher = JsonRpcDispatcher()