Add attempt finished signal

Author: igalshilman

python/restate/context.py

@@ -72,6 +72,31 @@ class RestateDurableSleepFuture(RestateDurableFuture[None]):
     def __await__(self) -> typing.Generator[Any, Any, None]:
         pass
 
+class AttemptFinishedEvent(abc.ABC):
+    """
+    Represents an attempt finished event.
+
+    This event is used to signal that an attempt has finished (either successfully or with an error), and it is now
+    safe to cleanup any attempt related resources, such as pending ctx.run() 3rd party calls, or any other resources that
+    are only valid for the duration of the attempt.
+
+    An attempt is considered finished when either the connection to the restate server is closed, the invocation is completed, or a transient
+    error occurs.
+    """
+
+    @abc.abstractmethod
+    def is_set(self) -> bool:
+        """
+        Returns True if the event is set, False otherwise.
+        """
+
+
+    @abc.abstractmethod
+    async def wait(self):
+        """
+        Waits for the event to be set.
+        """
+
 
 @dataclass
 class Request:
@@ -83,11 +108,13 @@ class Request:
         headers (dict[str, str]): The headers of the request.
         attempt_headers (dict[str, str]): The attempt headers of the request.
         body (bytes): The body of the request.
+        attempt_finished_event (AttemptFinishedEvent): The teardown event of the request.
     """
     id: str
     headers: Dict[str, str]
     attempt_headers: Dict[str,str]
     body: bytes
+    attempt_finished_event: AttemptFinishedEvent
 
 
 class KeyValueStore(abc.ABC):
@@ -158,6 +185,7 @@ async def cancel_invocation(self) -> None:
             await ctx.cancel_invocation(await f.invocation_id())
         """
 
+
 class Context(abc.ABC):
     """
     Represents the context of the current invocation.

python/restate/discovery.py

@@ -184,7 +184,7 @@ def compute_discovery(endpoint: RestateEndpoint, discovered_as : typing.Literal[
             if handler.handler_io.input_type and handler.handler_io.input_type.is_void:
                 inp = {}
             else:
-                inp =InputPayload(required=False,
+                inp = InputPayload(required=False,
                                    contentType=handler.handler_io.accept,
                                    jsonSchema=json_schema_from_type_hint(handler.handler_io.input_type))
             # output

python/restate/server.py

@@ -126,11 +126,15 @@ async def process_invocation_to_completion(vm: VMWrapper,
     try:
         await context.enter()
     except asyncio.exceptions.CancelledError:
+        context.on_attempt_finished()
         raise
     # pylint: disable=W0718
     except Exception:
         traceback.print_exc()
-    await context.leave()
+    try:
+        await context.leave()
+    finally:
+        context.on_attempt_finished()
 
 class LifeSpanNotImplemented(ValueError):
     """Signal to the asgi server that we didn't implement lifespans"""

python/restate/server_context.py

@@ -24,7 +24,7 @@
 import typing
 import traceback
 
-from restate.context import DurablePromise, ObjectContext, Request, RestateDurableCallFuture, RestateDurableFuture, SendHandle, RestateDurableSleepFuture
+from restate.context import DurablePromise, AttemptFinishedEvent, ObjectContext, Request, RestateDurableCallFuture, RestateDurableFuture, SendHandle, RestateDurableSleepFuture
 from restate.exceptions import TerminalError
 from restate.handler import Handler, handler_from_callable, invoke_handler
 from restate.serde import BytesSerde, DefaultSerde, JsonSerde, Serde
@@ -37,6 +37,24 @@
 I = TypeVar('I')
 O = TypeVar('O')
 
+
+class ServerTeardownEvent(AttemptFinishedEvent):
+    """
+    This class implements the teardown event for the server.
+    """
+
+    def __init__(self, event: asyncio.Event) -> None:
+        super().__init__()
+        self.event = event
+
+    def is_set(self):
+        return self.event.is_set()
+
+    async def wait(self):
+        """Wait for the event to be set."""
+        await self.event.wait()
+
+
 class LazyFuture:
     """
     Creates a task lazily, and allows multiple awaiters to the same coroutine.
@@ -229,6 +247,7 @@ def __init__(self,
         self.receive = receive
         self.run_coros_to_execute: dict[int,  Callable[[], Awaitable[None]]] = {}
         self.sync_point = SyncPoint()
+        self.request_finished_event = asyncio.Event()
 
     async def enter(self):
         """Invoke the user code."""
@@ -285,6 +304,12 @@ async def leave(self):
             'body': b'',
             'more_body': False,
         })
+        # notify to any holder of the abort signal that we are done
+
+    def on_attempt_finished(self):
+        """Notify the attempt finished event."""
+        self.request_finished_event.set()
+
 
     async def take_and_send_output(self):
         """Take output from state machine and send it"""
@@ -407,6 +432,7 @@ def request(self) -> Request:
             headers=dict(self.invocation.headers),
             attempt_headers=self.attempt_headers,
             body=self.invocation.input_buffer,
+            attempt_finished_event=ServerTeardownEvent(self.request_finished_event),
         )
 
     async def create_run_coroutine(self,