Add invocation_id to a call promise

Author: igalshilman

python/restate/context.py

@@ -38,6 +38,19 @@ def __await__(self):
         pass
 
 
+# pylint: disable=R0903
+class RestateDurableCallFuture(RestateDurableFuture[T]):
+    """
+    Represents a durable call future.
+    """
+
+    @abc.abstractmethod
+    async def invocation_id(self) -> str:
+        """
+        Returns the invocation id of the call.
+        """
+
+
 
 @dataclass
 class Request:
@@ -147,7 +160,7 @@ def sleep(self, delta: timedelta) -> RestateDurableFuture[None]:
     def service_call(self,
                      tpe: Callable[[Any, I], Awaitable[O]],
                      arg: I,
-                     idempotency_key: str | None = None) -> RestateDurableFuture[O]:
+                     idempotency_key: str | None = None) -> RestateDurableCallFuture[O]:
         """
         Invokes the given service with the given argument.
         """
@@ -170,7 +183,7 @@ def object_call(self,
                     key: str,
                     arg: I,
                     idempotency_key: str | None = None,
-                    ) -> RestateDurableFuture[O]:
+                    ) -> RestateDurableCallFuture[O]:
         """
         Invokes the given object with the given argument.
         """
@@ -193,7 +206,7 @@ def workflow_call(self,
                     key: str,
                     arg: I,
                     idempotency_key: str | None = None,
-                    ) -> RestateDurableFuture[O]:
+                    ) -> RestateDurableCallFuture[O]:
         """
         Invokes the given workflow with the given argument.
         """
@@ -217,7 +230,7 @@ def generic_call(self,
                      handler: str,
                      arg: bytes,
                      key: Optional[str] = None,
-                     idempotency_key: str | None = None)  -> RestateDurableFuture[bytes]:
+                     idempotency_key: str | None = None)  -> RestateDurableCallFuture[bytes]:
         """
         Invokes the given generic service/handler with the given argument.
         """

python/restate/server_context.py

@@ -18,7 +18,7 @@
 import typing
 import traceback
 
-from restate.context import DurablePromise, ObjectContext, Request, RestateDurableFuture, SendHandle
+from restate.context import DurablePromise, ObjectContext, Request, RestateDurableCallFuture, RestateDurableFuture, SendHandle
 from restate.exceptions import TerminalError
 from restate.handler import Handler, handler_from_callable, invoke_handler
 from restate.serde import BytesSerde, DefaultSerde, JsonSerde, Serde
@@ -45,24 +45,36 @@
 class ServerDurableFuture(RestateDurableFuture[T]):
     """This class implements a durable future API"""
     value: T | None = None
-    metadata: Dict[str, Any] | None = None
 
     def __init__(self, handle: int, factory) -> None:
         super().__init__()
         self.factory = factory
         self.handle = handle
 
-    def with_metadata(self, **metadata) -> 'ServerDurableFuture':
-        """Add metadata to the future."""
-        self.metadata = metadata
-        return self
-
     def __await__(self):
-        print("..........Awaiting............", flush=True)
         task = asyncio.create_task(self.factory())
         return task.__await__()
 
 
+class ServerCallDurableFuture(RestateDurableCallFuture[T], ServerDurableFuture[T]):
+    """This class implements a durable future but for calls"""
+    _invocation_id: typing.Optional[str] = None
+
+    def __init__(self, result_handle: int,
+                 result_factory,
+                 invocation_id_handle: int,
+                 invocation_id_factory) -> None:
+        super().__init__(result_handle, result_factory)
+        self.invocation_id_handle = invocation_id_handle
+        self.invocation_id_factory = invocation_id_factory
+
+    async def invocation_id(self) -> str:
+        """Get the invocation id."""
+        if self._invocation_id is None:
+            self._invocation_id  = await self.invocation_id_factory()
+        return self._invocation_id
+
+
 class ServerSendHandle(SendHandle):
     """This class implements the send API"""
     _invocation_id: typing.Optional[str]
@@ -279,6 +291,21 @@ async def transform():
         return ServerDurableFuture(handle, transform)
 
 
+    
+    def create_call_df(self, handle: int, invocation_id_handle: int, serde: Serde[T] | None = None) -> ServerCallDurableFuture[T]:
+        """Create a durable future."""
+
+        async def transform():
+            res = await self.create_poll_or_cancel_coroutine(handle)
+            if res is None or serde is None:
+                return res
+            return serde.deserialize(res)
+
+        def inv_id_factory():
+            return self.create_poll_or_cancel_coroutine(invocation_id_handle)
+
+        return ServerCallDurableFuture(handle, transform, invocation_id_handle, inv_id_factory)
+
 
     def get(self, name: str, serde: Serde[T] = JsonSerde()) -> Awaitable[Optional[T]]:
         handle = self.vm.sys_get_state(name)
@@ -366,7 +393,7 @@ def do_call(self,
                 send: bool = False,
                 idempotency_key: str | None = None,
                 headers: typing.List[typing.Tuple[str, str]] | None = None
-                ) -> RestateDurableFuture[O] | SendHandle:
+                ) -> RestateDurableCallFuture[O] | SendHandle:
         """Make an RPC call to the given handler"""
         target_handler = handler_from_callable(tpe)
         service=target_handler.service_tag.name
@@ -387,7 +414,7 @@ def do_raw_call(self,
                  send: bool = False,
                  idempotency_key: str | None = None,
                  headers: typing.List[typing.Tuple[str, str]] | None = None
-                 ) -> RestateDurableFuture[O] | SendHandle:
+                 ) -> RestateDurableCallFuture[O] | SendHandle:
         """Make an RPC call to the given handler"""
         parameter = input_serde.serialize(input_param)
         if send_delay:
@@ -405,15 +432,16 @@ def do_raw_call(self,
                                   idempotency_key=idempotency_key,
                                   headers=headers)
 
-        # TODO: specialize this future for calls!
-        return self.create_df(handle=handle.result_handle, serde=output_serde).with_metadata(invocation_id=handle.invocation_id_handle)
+        return self.create_call_df(handle=handle.result_handle,
+                                   invocation_id_handle=handle.invocation_id_handle,
+                                   serde=output_serde)
 
     def service_call(self,
                      tpe: Callable[[Any, I], Awaitable[O]],
                      arg: I,
                      idempotency_key: str | None = None,
                      headers: typing.List[typing.Tuple[str, str]] | None = None
-                     ) -> RestateDurableFuture[O]:
+                     ) -> RestateDurableCallFuture[O]:
         coro = self.do_call(tpe, arg, idempotency_key=idempotency_key, headers=headers)
         assert not isinstance(coro, SendHandle)
         return coro
@@ -429,7 +457,7 @@ def object_call(self,
                     arg: I,
                     idempotency_key: str | None = None,
                     headers: typing.List[typing.Tuple[str, str]] | None = None
-                    ) -> RestateDurableFuture[O]:
+                    ) -> RestateDurableCallFuture[O]:
         coro = self.do_call(tpe, arg, key, idempotency_key=idempotency_key, headers=headers)
         assert not isinstance(coro, SendHandle)
         return coro
@@ -445,15 +473,15 @@ def workflow_call(self,
                         arg: I,
                         idempotency_key: str | None = None,
                         headers: typing.List[typing.Tuple[str, str]] | None = None
-                        ) -> RestateDurableFuture[O]:
+                        ) -> RestateDurableCallFuture[O]:
         return self.object_call(tpe, key, arg, idempotency_key=idempotency_key, headers=headers)
 
     def workflow_send(self, tpe: Callable[[Any, I], Awaitable[O]], key: str, arg: I, send_delay: timedelta | None = None, idempotency_key: str | None = None, headers: typing.List[typing.Tuple[str, str]] | None = None) -> SendHandle:
         send = self.object_send(tpe, key, arg, send_delay, idempotency_key=idempotency_key, headers=headers)
         assert isinstance(send, SendHandle)
         return send
 
-    def generic_call(self, service: str, handler: str, arg: bytes, key: str | None = None, idempotency_key: str | None = None, headers: typing.List[typing.Tuple[str, str]] | None = None) -> RestateDurableFuture[bytes]:
+    def generic_call(self, service: str, handler: str, arg: bytes, key: str | None = None, idempotency_key: str | None = None, headers: typing.List[typing.Tuple[str, str]] | None = None) -> RestateDurableCallFuture[bytes]:
         serde = BytesSerde()
         call_handle = self.do_raw_call(service=service,
                                 handler=handler,