enhancement: Add Lifespan Events Handling In ASGI app (#113)

Author: objecthuman

.github/workflows/unit-test.yaml

@@ -0,0 +1,37 @@
+name: Unit Tests
+
+on:
+  pull_request:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+
+permissions:
+  contents: read
+  checks: write
+  pull-requests: write
+
+jobs:
+  restate-sdk-unit-tests:
+    name: "Run Unit Tests (Python ${{ matrix.python }})"
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python: ["3.11", "3.12"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: extractions/setup-just@v2
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+      - name: Build Rust module
+        uses: PyO3/maturin-action@v1
+        with:
+          args: --out dist --interpreter python${{ matrix.python }}
+          sccache: "true"
+          container: off
+      - run: pip install -r requirements.txt
+      - run: pip install dist/*
+      - name: Run Unit Tests
+        run: python -m pytest tests/ -v

examples/example.py

@@ -12,24 +12,27 @@
 # pylint: disable=C0116
 # pylint: disable=W0613
 
+import asyncio
 import restate
 
+from concurrent_greeter import concurrent_greeter
 from greeter import greeter
+from pydantic_greeter import pydantic_greeter
 from virtual_object import counter
 from workflow import payment
-from pydantic_greeter import pydantic_greeter
-from concurrent_greeter import concurrent_greeter
+
+
 
 app = restate.app(services=[greeter,
                             counter,
                             payment,
                             pydantic_greeter,
-                            concurrent_greeter])
+                            concurrent_greeter,
+                            ],)
 
 if __name__ == "__main__":
     import hypercorn
     import hypercorn.asyncio
-    import asyncio
     conf = hypercorn.Config()
     conf.bind = ["0.0.0.0:9080"]
     asyncio.run(hypercorn.asyncio.serve(app, conf))

pyproject.toml

@@ -21,7 +21,7 @@ Source = "https://github.com/restatedev/sdk-python"
 "Bug Tracker" = "https://github.com/restatedev/sdk-python/issues"
 
 [project.optional-dependencies]
-test = ["pytest", "hypercorn"]
+test = ["pytest", "hypercorn", "pytest-asyncio==1.1.0"]
 lint = ["mypy", "pylint"]
 harness = ["testcontainers", "hypercorn", "httpx"]
 serde = ["dacite", "pydantic"]

python/restate/aws_lambda.py

@@ -21,7 +21,7 @@
                                   Receive,
                                   HTTPResponseStartEvent,
                                   HTTPResponseBodyEvent,
-                                  HTTPRequestEvent)
+                                  HTTPRequestEvent, Send)
 
 class RestateLambdaRequest(TypedDict):
     """
@@ -162,8 +162,9 @@ def lambda_handler(event: RestateLambdaRequest, _context: Any) -> RestateLambdaR
         scope = create_scope(event)
         recv = request_to_receive(event)
         send = ResponseCollector()
+        send_typed = cast(Send, send)
 
-        asgi_instance = asgi_app(scope, recv, send)
+        asgi_instance = asgi_app(scope, recv, send_typed)
         asgi_task = loop.create_task(asgi_instance)  # type: ignore[var-annotated, arg-type]
         loop.run_until_complete(asgi_task)
 

python/restate/context.py

@@ -220,8 +220,8 @@ def request(self) -> Request:
         """
 
 
-    @typing_extensions.deprecated("`run` is deprecated, use `run_typed` instead for better type safety")
     @overload
+    @typing_extensions.deprecated("`run` is deprecated, use `run_typed` instead for better type safety")
     @abc.abstractmethod
     def run(self,
             name: str,
@@ -250,8 +250,8 @@ def run(self,
 
         """
 
-    @typing_extensions.deprecated("`run` is deprecated, use `run_typed` instead for better type safety")
     @overload
+    @typing_extensions.deprecated("`run` is deprecated, use `run_typed` instead for better type safety")
     @abc.abstractmethod
     def run(self,
             name: str,
@@ -280,6 +280,7 @@ def run(self,
 
         """
 
+    @overload
     @typing_extensions.deprecated("`run` is deprecated, use `run_typed` instead for better type safety")
     @abc.abstractmethod
     def run(self,

python/restate/endpoint.py

@@ -15,6 +15,7 @@
 
 import typing
 
+from restate.server_types import LifeSpan
 from restate.service import Service
 from restate.object import VirtualObject
 from restate.workflow import Workflow
@@ -80,7 +81,7 @@ def identity_key(self, identity_key: str):
         """Add an identity key to this endpoint."""
         self.identity_keys.append(identity_key)
 
-    def app(self):
+    def app(self, lifespan: typing.Optional[LifeSpan] = None):
         """
         Returns the ASGI application for this endpoint.
 
@@ -94,12 +95,13 @@ def app(self):
         # pylint: disable=C0415
         # pylint: disable=R0401
         from restate.server import asgi_app
-        return asgi_app(self)
+        return asgi_app(self, lifespan)
 
 def app(
     services: typing.Iterable[typing.Union[Service, VirtualObject, Workflow]],
     protocol: typing.Optional[typing.Literal["bidi", "request_response"]] = None,
-    identity_keys: typing.Optional[typing.List[str]] = None):
+    identity_keys: typing.Optional[typing.List[str]] = None,
+    lifespan: typing.Optional[LifeSpan] = None):
     """A restate ASGI application that hosts the given services."""
     endpoint = Endpoint()
     if protocol == "bidi":
@@ -111,4 +113,4 @@ def app(
     if identity_keys:
         for key in identity_keys:
             endpoint.identity_key(key)
-    return endpoint.app()
+    return endpoint.app(lifespan)

python/restate/server.py

@@ -11,12 +11,14 @@
 """This module contains the ASGI server for the restate framework."""
 
 import asyncio
-from typing import Dict, TypedDict, Literal
+import inspect
+from typing import Any, Dict, TypedDict, Literal
 import traceback
+import typing
 from restate.discovery import compute_discovery_json
 from restate.endpoint import Endpoint
 from restate.server_context import ServerInvocationContext, DisconnectedException
-from restate.server_types import Receive, ReceiveChannel, Scope, Send, binary_to_header, header_to_binary # pylint: disable=line-too-long
+from restate.server_types import Receive, ReceiveChannel, Scope, Send, binary_to_header, header_to_binary, LifeSpan # pylint: disable=line-too-long
 from restate.vm import VMWrapper
 from restate._internal import PyIdentityVerifier, IdentityVerificationException # pylint: disable=import-error,no-name-in-module
 from restate._internal import SDK_VERSION # pylint: disable=import-error,no-name-in-module
@@ -186,10 +188,6 @@ async def process_invocation_to_completion(vm: VMWrapper,
     finally:
         context.on_attempt_finished()
 
-class LifeSpanNotImplemented(ValueError):
-    """Signal to the asgi server that we didn't implement lifespans"""
-
-
 class ParsedPath(TypedDict):
     """Parsed path from the request."""
     type: Literal["invocation", "health", "discover", "unknown"]
@@ -216,8 +214,55 @@ def parse_path(request: str) -> ParsedPath:
     # anything other than invoke is 404
     return { "type": "unknown" , "service": None, "handler": None }
 
+def is_async_context_manager(obj: Any):
+    """check if passed object is an async context manager"""
+    return (hasattr(obj, '__aenter__') and
+            hasattr(obj, '__aexit__') and
+            inspect.iscoroutinefunction(obj.__aenter__) and
+            inspect.iscoroutinefunction(obj.__aexit__))
 
-def asgi_app(endpoint: Endpoint):
+
+async def lifespan_processor(
+    scope: Scope,
+    receive: Receive,
+    send: Send,
+    lifespan: LifeSpan
+) -> None:
+    """Process lifespan context manager."""
+    started = False
+    assert scope["type"] in ["lifespan", "lifespan.startup", "lifespan.shutdown"]
+    assert is_async_context_manager(lifespan()), "lifespan must be an async context manager"
+    await receive()
+    try:
+        async with lifespan() as maybe_state:
+            if maybe_state is not None:
+                if "state" not in scope:
+                    raise RuntimeError("The server does not support state in lifespan")
+                scope["state"] = maybe_state
+            await send({
+                "type": "lifespan.startup.complete", # type: ignore
+            })
+            started = True
+            await receive()
+    except Exception:
+        exc_text = traceback.format_exc()
+        if started:
+            await send({
+                "type": "lifespan.shutdown.failed",
+                "message": exc_text
+            })
+        else:
+            await send({
+                "type": "lifespan.startup.failed",
+                "message": exc_text
+            })
+        raise
+    await send({
+        "type": "lifespan.shutdown.complete" # type: ignore
+    })
+
+# pylint: disable=too-many-return-statements
+def asgi_app(endpoint: Endpoint, lifespan: typing.Optional[LifeSpan] = None):
     """Create an ASGI-3 app for the given endpoint."""
 
     # Prepare request signer
@@ -226,14 +271,17 @@ def asgi_app(endpoint: Endpoint):
     async def app(scope: Scope, receive: Receive, send: Send):
         try:
             if scope['type'] == 'lifespan':
-                raise LifeSpanNotImplemented()
+                if lifespan is not None:
+                    await lifespan_processor(scope, receive, send, lifespan)
+                    return
+                return
+
             if scope['type'] != 'http':
                 raise NotImplementedError(f"Unknown scope type {scope['type']}")
 
             request_path = scope['path']
             assert isinstance(request_path, str)
             request: ParsedPath = parse_path(request_path)
-
             # Health check
             if request['type'] == 'health':
                 await send_health_check(send)
@@ -249,7 +297,6 @@ async def app(scope: Scope, receive: Receive, send: Send):
                 # Identify verification failed, send back unauthorized and close
                 await send_status(send, receive, 401)
                 return
-
             # might be a discovery request
             if request['type'] == 'discover':
                 await send_discovery(scope, send, endpoint)
@@ -283,8 +330,6 @@ async def app(scope: Scope, receive: Receive, send: Send):
                                                            send)
             finally:
                 await receive_channel.close()
-        except LifeSpanNotImplemented as e:
-            raise e
         except Exception as e:
             traceback.print_exc()
             raise e

python/restate/server_types.py

@@ -16,7 +16,8 @@
 
 import asyncio
 from typing import (Awaitable, Callable, Dict, Iterable, List,
-                    Tuple, Union, TypedDict, Literal, Optional, NotRequired, Any)
+                    Tuple, Union, TypedDict, Literal, Optional,
+                    NotRequired, Any, AsyncContextManager)
 
 class ASGIVersions(TypedDict):
     """ASGI Versions"""
@@ -25,7 +26,7 @@ class ASGIVersions(TypedDict):
 
 class Scope(TypedDict):
     """ASGI Scope"""
-    type: Literal["http"]
+    type: Literal["http", "lifespan", "lifespan.startup", "lifespan.shutdown"]
     asgi: ASGIVersions
     http_version: str
     method: str
@@ -64,18 +65,29 @@ class HTTPResponseBodyEvent(TypedDict):
     body: bytes
     more_body: bool
 
+class LifeSpanEvent(TypedDict):
+    """ASGI LifeSpan event"""
+    type: Literal["lifespan.startup.complete",
+                "lifespan.shutdown.complete",
+                "lifespan.startup.failed",
+                "lifespan.shutdown.failed"]
+    message: Optional[str]
 
-ASGIReceiveEvent = HTTPRequestEvent
+
+ASGIReceiveEvent = Union[HTTPRequestEvent, LifeSpanEvent]
 
 
 ASGISendEvent = Union[
     HTTPResponseStartEvent,
-    HTTPResponseBodyEvent
+    HTTPResponseBodyEvent,
+    LifeSpanEvent
 ]
 
 Receive = Callable[[], Awaitable[ASGIReceiveEvent]]
 Send = Callable[[ASGISendEvent], Awaitable[None]]
 
+LifeSpan = Callable[[], AsyncContextManager[Any]]
+
 ASGIApp = Callable[
     [
         Scope,
@@ -84,7 +96,6 @@ class HTTPResponseBodyEvent(TypedDict):
     ],
     Awaitable[None],
 ]
-
 def header_to_binary(headers: Iterable[Tuple[str, str]]) -> List[Tuple[bytes, bytes]]:
     """Convert a list of headers to a list of binary headers."""
     return [ (k.encode('utf-8'), v.encode('utf-8')) for k,v in headers ]

requirements.txt

@@ -7,3 +7,4 @@ pydantic
 httpx
 testcontainers
 typing-extensions>=4.14.0
+pytest-asyncio==1.1.0