starpc 0.52.0 → 0.52.2

package/echo/echo_srpc.py

@@ -0,0 +1,309 @@
+from __future__ import annotations
+
+import contextlib
+from collections.abc import AsyncIterable, AsyncIterator
+from typing import Protocol
+
+from google.protobuf import empty_pb2 as _google_protobuf_empty_pb2
+
+from echo import (
+    echo_pb2 as _github_com_aperturerobotics_starpc_echo_echo_pb2,
+)
+from rpcstream import (
+    rpcstream_pb2 as _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2,
+)
+from starpc.call import Call, CallProtocolError
+from starpc.client import Client
+from starpc.server import ServiceRegistry
+from starpc.service import MethodDescriptor, ServiceDescriptor, bidirectional_bytes
+
+ECHOER_SERVICE = ServiceDescriptor(
+    "echo.Echoer",
+    (
+        MethodDescriptor(
+            "Echo",
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg,
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg,
+            False,
+            False,
+        ),
+        MethodDescriptor(
+            "EchoServerStream",
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg,
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg,
+            False,
+            True,
+        ),
+        MethodDescriptor(
+            "EchoClientStream",
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg,
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg,
+            True,
+            False,
+        ),
+        MethodDescriptor(
+            "EchoBidiStream",
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg,
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg,
+            True,
+            True,
+        ),
+        MethodDescriptor(
+            "RpcStream",
+            _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2.RpcStreamPacket,
+            _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2.RpcStreamPacket,
+            True,
+            True,
+        ),
+        MethodDescriptor(
+            "DoNothing",
+            _google_protobuf_empty_pb2.Empty,
+            _google_protobuf_empty_pb2.Empty,
+            False,
+            False,
+        ),
+    ),
+)
+
+
+class EchoerClient:
+    def __init__(self, client: Client, service: str | None = None) -> None:
+        self._client = client
+        self._service = service or "echo.Echoer"
+
+    async def echo(
+        self, request: _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+    ) -> _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg:
+        call = await self._client.open_call(
+            self._service, "Echo", request.SerializeToString(deterministic=True)
+        )
+        try:
+            data = await call.receive()
+            if data is None:
+                raise CallProtocolError("missing unary response")
+            response = _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg()
+            response.ParseFromString(data)
+            if await call.receive() is not None:
+                raise CallProtocolError("extra unary response")
+            return response
+        finally:
+            await call.aclose()
+
+    async def echo_server_stream(
+        self, request: _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+    ) -> AsyncIterator[_github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg]:
+        call = await self._client.open_call(
+            self._service,
+            "EchoServerStream",
+            request.SerializeToString(deterministic=True),
+        )
+        try:
+            while True:
+                data = await call.receive()
+                if data is None:
+                    return
+                response = _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg()
+                response.ParseFromString(data)
+                yield response
+        finally:
+            await call.aclose()
+
+    async def echo_client_stream(
+        self,
+        requests: AsyncIterable[
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+        ],
+    ) -> _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg:
+        call = await self._client.open_call(self._service, "EchoClientStream")
+        try:
+            async for request in requests:
+                await call.send(request.SerializeToString(deterministic=True))
+            # A failed half-close means the call ended; receive reports its outcome.
+            with contextlib.suppress(OSError):
+                await call.finish()
+            data = await call.receive()
+            if data is None:
+                raise CallProtocolError("missing client-stream response")
+            response = _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg()
+            response.ParseFromString(data)
+            if await call.receive() is not None:
+                raise CallProtocolError("extra client-stream response")
+            return response
+        finally:
+            await call.aclose()
+
+    async def echo_bidi_stream(
+        self,
+        requests: AsyncIterable[
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+        ],
+    ) -> AsyncIterator[_github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg]:
+        call = await self._client.open_call(self._service, "EchoBidiStream")
+
+        async def encoded() -> AsyncIterator[bytes]:
+            async for request in requests:
+                yield request.SerializeToString(deterministic=True)
+
+        async for data in bidirectional_bytes(call, encoded()):
+            response = _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg()
+            response.ParseFromString(data)
+            yield response
+
+    async def rpc_stream(
+        self,
+        requests: AsyncIterable[
+            _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2.RpcStreamPacket
+        ],
+    ) -> AsyncIterator[
+        _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2.RpcStreamPacket
+    ]:
+        call = await self._client.open_call(self._service, "RpcStream")
+
+        async def encoded() -> AsyncIterator[bytes]:
+            async for request in requests:
+                yield request.SerializeToString(deterministic=True)
+
+        async for data in bidirectional_bytes(call, encoded()):
+            response = _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2.RpcStreamPacket()
+            response.ParseFromString(data)
+            yield response
+
+    async def do_nothing(
+        self, request: _google_protobuf_empty_pb2.Empty
+    ) -> _google_protobuf_empty_pb2.Empty:
+        call = await self._client.open_call(
+            self._service, "DoNothing", request.SerializeToString(deterministic=True)
+        )
+        try:
+            data = await call.receive()
+            if data is None:
+                raise CallProtocolError("missing unary response")
+            response = _google_protobuf_empty_pb2.Empty()
+            response.ParseFromString(data)
+            if await call.receive() is not None:
+                raise CallProtocolError("extra unary response")
+            return response
+        finally:
+            await call.aclose()
+
+
+class EchoerServer(Protocol):
+    async def echo(
+        self, request: _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+    ) -> _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg: ...
+    def echo_server_stream(
+        self, request: _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+    ) -> AsyncIterator[_github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg]: ...
+    async def echo_client_stream(
+        self,
+        requests: AsyncIterator[
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+        ],
+    ) -> _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg: ...
+    def echo_bidi_stream(
+        self,
+        requests: AsyncIterator[
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+        ],
+    ) -> AsyncIterator[_github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg]: ...
+    def rpc_stream(
+        self,
+        requests: AsyncIterator[
+            _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2.RpcStreamPacket
+        ],
+    ) -> AsyncIterator[
+        _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2.RpcStreamPacket
+    ]: ...
+    async def do_nothing(
+        self, request: _google_protobuf_empty_pb2.Empty
+    ) -> _google_protobuf_empty_pb2.Empty: ...
+
+
+def register_echoer(
+    registry: ServiceRegistry,
+    implementation: EchoerServer,
+    service: str = "echo.Echoer",
+) -> None:
+    async def echo_handler(call: Call) -> None:
+        first = await call.receive()
+        if first is None:
+            raise CallProtocolError("missing initial request")
+        request = _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg()
+        request.ParseFromString(first)
+        response = await implementation.echo(request)
+        await call.send(response.SerializeToString(deterministic=True))
+
+    registry.register(service, "Echo", echo_handler)
+
+    async def echo_server_stream_handler(call: Call) -> None:
+        first = await call.receive()
+        if first is None:
+            raise CallProtocolError("missing initial request")
+        request = _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg()
+        request.ParseFromString(first)
+        async for response in implementation.echo_server_stream(request):
+            await call.send(response.SerializeToString(deterministic=True))
+
+    registry.register(service, "EchoServerStream", echo_server_stream_handler)
+
+    async def echo_client_stream_handler(call: Call) -> None:
+        async def requests() -> AsyncIterator[
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+        ]:
+            while True:
+                data = await call.receive()
+                if data is None:
+                    return
+                request = _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg()
+                request.ParseFromString(data)
+                yield request
+
+        response = await implementation.echo_client_stream(requests())
+        await call.send(response.SerializeToString(deterministic=True))
+
+    registry.register(service, "EchoClientStream", echo_client_stream_handler)
+
+    async def echo_bidi_stream_handler(call: Call) -> None:
+        async def requests() -> AsyncIterator[
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+        ]:
+            while True:
+                data = await call.receive()
+                if data is None:
+                    return
+                request = _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg()
+                request.ParseFromString(data)
+                yield request
+
+        async for response in implementation.echo_bidi_stream(requests()):
+            await call.send(response.SerializeToString(deterministic=True))
+
+    registry.register(service, "EchoBidiStream", echo_bidi_stream_handler)
+
+    async def rpc_stream_handler(call: Call) -> None:
+        async def requests() -> AsyncIterator[
+            _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2.RpcStreamPacket
+        ]:
+            while True:
+                data = await call.receive()
+                if data is None:
+                    return
+                request = _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2.RpcStreamPacket()
+                request.ParseFromString(data)
+                yield request
+
+        async for response in implementation.rpc_stream(requests()):
+            await call.send(response.SerializeToString(deterministic=True))
+
+    registry.register(service, "RpcStream", rpc_stream_handler)
+
+    async def do_nothing_handler(call: Call) -> None:
+        first = await call.receive()
+        if first is None:
+            raise CallProtocolError("missing initial request")
+        request = _google_protobuf_empty_pb2.Empty()
+        request.ParseFromString(first)
+        response = await implementation.do_nothing(request)
+        await call.send(response.SerializeToString(deterministic=True))
+
+    registry.register(service, "DoNothing", do_nothing_handler)

package/echo/echo_srpc.pyi

@@ -0,0 +1,85 @@
+from collections.abc import AsyncIterable, AsyncIterator
+from typing import Protocol
+
+from google.protobuf import empty_pb2 as _google_protobuf_empty_pb2
+
+from echo import (
+    echo_pb2 as _github_com_aperturerobotics_starpc_echo_echo_pb2,
+)
+from rpcstream import (
+    rpcstream_pb2 as _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2,
+)
+from starpc.client import Client
+from starpc.server import ServiceRegistry
+from starpc.service import ServiceDescriptor
+
+ECHOER_SERVICE: ServiceDescriptor
+
+class EchoerClient:
+    def __init__(self, client: Client, service: str | None = None) -> None: ...
+    async def echo(
+        self, request: _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+    ) -> _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg: ...
+    def echo_server_stream(
+        self, request: _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+    ) -> AsyncIterator[_github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg]: ...
+    async def echo_client_stream(
+        self,
+        requests: AsyncIterable[
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+        ],
+    ) -> _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg: ...
+    def echo_bidi_stream(
+        self,
+        requests: AsyncIterable[
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+        ],
+    ) -> AsyncIterator[_github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg]: ...
+    def rpc_stream(
+        self,
+        requests: AsyncIterable[
+            _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2.RpcStreamPacket
+        ],
+    ) -> AsyncIterator[
+        _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2.RpcStreamPacket
+    ]: ...
+    async def do_nothing(
+        self, request: _google_protobuf_empty_pb2.Empty
+    ) -> _google_protobuf_empty_pb2.Empty: ...
+
+class EchoerServer(Protocol):
+    async def echo(
+        self, request: _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+    ) -> _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg: ...
+    def echo_server_stream(
+        self, request: _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+    ) -> AsyncIterator[_github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg]: ...
+    async def echo_client_stream(
+        self,
+        requests: AsyncIterator[
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+        ],
+    ) -> _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg: ...
+    def echo_bidi_stream(
+        self,
+        requests: AsyncIterator[
+            _github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg
+        ],
+    ) -> AsyncIterator[_github_com_aperturerobotics_starpc_echo_echo_pb2.EchoMsg]: ...
+    def rpc_stream(
+        self,
+        requests: AsyncIterator[
+            _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2.RpcStreamPacket
+        ],
+    ) -> AsyncIterator[
+        _github_com_aperturerobotics_starpc_rpcstream_rpcstream_pb2.RpcStreamPacket
+    ]: ...
+    async def do_nothing(
+        self, request: _google_protobuf_empty_pb2.Empty
+    ) -> _google_protobuf_empty_pb2.Empty: ...
+
+def register_echoer(
+    registry: ServiceRegistry,
+    implementation: EchoerServer,
+    service: str = "echo.Echoer",
+) -> None: ...

package/echo/half_close_test.rs

@@ -0,0 +1,178 @@
+//! Half-close tests for the generated Rust callers.
+//!
+//! The half-close is only a notice once the request is written: a caller that
+//! half-closes after the remote ended the call must still read the buffered
+//! responses and the call's outcome.
+
+#[allow(dead_code)]
+mod gen;
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use futures::StreamExt as _;
+use prost::Message;
+use starpc::client::OpenStream;
+use starpc::testing::{create_pipe_default, create_test_pair, SingleInMemoryOpener};
+use starpc::{Client, Error, Mux, PacketCodec, Result, Server, SrpcClient, Stream, StreamExt};
+use tokio::task::JoinHandle;
+use tokio_util::codec::FramedRead;
+
+use gen::{EchoMsg, EchoerClient, EchoerClientImpl, EchoerHandler, EchoerServer};
+
+const BODY_TXT: &str = "hello world via starpc half-close test";
+
+/// Answers each streaming call once and completes it.
+struct OnceEchoServer;
+
+#[async_trait]
+impl EchoerServer for OnceEchoServer {
+    async fn echo(&self, _context: &starpc::Context, request: EchoMsg) -> Result<EchoMsg> {
+        Ok(request)
+    }
+
+    async fn echo_server_stream(&self, request: EchoMsg, stream: Box<dyn Stream>) -> Result<()> {
+        stream.msg_send(&request).await
+    }
+
+    async fn echo_client_stream(&self, stream: &dyn Stream) -> Result<EchoMsg> {
+        // Reply to the first message without waiting for the half-close.
+        stream.msg_recv().await
+    }
+
+    async fn echo_bidi_stream(&self, _stream: Box<dyn Stream>) -> Result<()> {
+        Err(Error::Unimplemented)
+    }
+
+    async fn rpc_stream(&self, _stream: Box<dyn Stream>) -> Result<()> {
+        Err(Error::Unimplemented)
+    }
+
+    async fn do_nothing(
+        &self,
+        _context: &starpc::Context,
+        request: gen::Empty,
+    ) -> Result<gen::Empty> {
+        Ok(request)
+    }
+}
+
+/// Returns each new stream only after its transport has closed, so the
+/// generated caller half-closes after the remote already ended the call.
+struct EndedClient<T: OpenStream> {
+    client: SrpcClient<T>,
+}
+
+#[async_trait]
+impl<T: OpenStream + 'static> Client for EndedClient<T> {
+    async fn exec_call<I, O>(&self, service: &str, method: &str, input: &I) -> Result<O>
+    where
+        I: Message + Send + Sync,
+        O: Message + Default,
+    {
+        self.client.exec_call(service, method, input).await
+    }
+
+    async fn new_stream(
+        &self,
+        service: &str,
+        method: &str,
+        first_msg: Option<&[u8]>,
+    ) -> Result<Box<dyn Stream>> {
+        let stream = self.client.new_stream(service, method, first_msg).await?;
+        // The transport close handler cancels the context after it settles
+        // the call's outcome and closes the writer.
+        stream.context().cancelled().await;
+        Ok(stream)
+    }
+}
+
+/// Returns one transport opener and the server task that its caller must join.
+fn serve_once_echo() -> (SingleInMemoryOpener, JoinHandle<Result<()>>) {
+    let mux = Arc::new(Mux::new());
+    mux.register(Arc::new(EchoerHandler::new(OnceEchoServer)))
+        .unwrap();
+    let (opener, server_stream) = create_test_pair();
+    let server = Server::with_arc(mux);
+    let task = tokio::spawn(async move { server.handle_stream(server_stream).await });
+    (opener, task)
+}
+
+/// Returns a transport that closes after CallStart and the peer task to join after observation.
+fn drop_after_request() -> (SingleInMemoryOpener, JoinHandle<()>) {
+    // Use the actual framing boundary to order peer closure after request admission.
+    let (client_stream, server_stream) = create_pipe_default();
+    let task = tokio::spawn(async move {
+        let mut framed = FramedRead::new(server_stream, PacketCodec::new());
+        let _ = framed.next().await;
+    });
+    (SingleInMemoryOpener::new(client_stream), task)
+}
+
+#[tokio::test]
+async fn server_stream_completed_before_close_send() {
+    // Complete the server stream before the generated caller half-closes.
+    let (opener, server) = serve_once_echo();
+    let client = EchoerClientImpl::new(EndedClient {
+        client: SrpcClient::new(opener),
+    });
+    let request = EchoMsg {
+        body: BODY_TXT.into(),
+    };
+    let stream = client
+        .echo_server_stream(&request)
+        .await
+        .expect("server stream discarded after remote completion");
+
+    // Read the buffered response, then the clean completion.
+    let msg = stream.recv().await.expect("buffered response lost");
+    assert_eq!(msg.body, BODY_TXT);
+    let end = stream.recv().await;
+    assert!(
+        matches!(end, Err(Error::StreamClosed)),
+        "expected StreamClosed after completion, got {end:?}"
+    );
+    server.await.unwrap().unwrap();
+}
+
+#[tokio::test]
+async fn client_stream_completed_before_close_send() {
+    let (opener, server) = serve_once_echo();
+    let client = EchoerClientImpl::new(SrpcClient::new(opener));
+    let stream = client.echo_client_stream().await.unwrap();
+
+    // Send the one request and wait for the call to end.
+    let request = EchoMsg {
+        body: BODY_TXT.into(),
+    };
+    stream.send(&request).await.unwrap();
+    stream.context().cancelled().await;
+
+    // The half-close fails, but the reply is still readable.
+    let msg = stream.close_and_recv().await.expect("buffered reply lost");
+    assert_eq!(msg.body, BODY_TXT);
+    server.await.unwrap().unwrap();
+}
+
+#[tokio::test]
+async fn server_stream_transport_failure_before_close_send() {
+    let (opener, server) = drop_after_request();
+    let client = EchoerClientImpl::new(EndedClient {
+        client: SrpcClient::new(opener),
+    });
+    let request = EchoMsg {
+        body: BODY_TXT.into(),
+    };
+    let stream = client
+        .echo_server_stream(&request)
+        .await
+        .expect("server stream discarded after transport failure");
+
+    // Recv reports that the call ended without a verdict.
+    let end = stream.recv().await;
+    assert!(
+        matches!(end, Err(Error::ClosedBeforeCompletion)),
+        "expected ClosedBeforeCompletion, got {end:?}"
+    );
+    server.await.unwrap();
+}

package/echo/integration_server.rs

@@ -1,3 +1,5 @@
+//! Echo transport endpoint used by cross-language integration checks.
+
 #[allow(dead_code)]
 mod gen;
 
@@ -9,21 +11,16 @@ use tokio::net::TcpListener;
 
 use gen::{EchoMsg, EchoerHandler, EchoerServer, Empty};
 
+/// Implements each echo call shape through the supplied transport stream.
 struct EchoServerImpl;
 
 #[async_trait]
 impl EchoerServer for EchoServerImpl {
-    async fn echo(&self, request: EchoMsg) -> Result<EchoMsg> {
-        Ok(EchoMsg {
-            body: request.body,
-        })
+    async fn echo(&self, _context: &starpc::Context, request: EchoMsg) -> Result<EchoMsg> {
+        Ok(EchoMsg { body: request.body })
     }
 
-    async fn echo_server_stream(
-        &self,
-        request: EchoMsg,
-        stream: Box<dyn Stream>,
-    ) -> Result<()> {
+    async fn echo_server_stream(&self, request: EchoMsg, stream: Box<dyn Stream>) -> Result<()> {
         for _ in 0..5 {
             let response = EchoMsg {
                 body: request.body.clone(),
@@ -34,10 +31,7 @@ impl EchoerServer for EchoServerImpl {
     }
 
     async fn echo_client_stream(&self, stream: &dyn Stream) -> Result<EchoMsg> {
-        match stream.msg_recv::<EchoMsg>().await {
-            Ok(msg) => Ok(msg),
-            Err(e) => Err(e),
-        }
+        stream.msg_recv().await
     }
 
     async fn echo_bidi_stream(&self, stream: Box<dyn Stream>) -> Result<()> {
@@ -63,11 +57,12 @@ impl EchoerServer for EchoServerImpl {
         Err(Error::Unimplemented)
     }
 
-    async fn do_nothing(&self, _request: Empty) -> Result<Empty> {
+    async fn do_nothing(&self, _context: &starpc::Context, _request: Empty) -> Result<Empty> {
         Ok(Empty {})
     }
 }
 
+/// Publishes a bound address and keeps accepted connections under the server owner.
 #[tokio::main]
 async fn main() -> Result<()> {
     let listener = TcpListener::bind("127.0.0.1:0").await?;
@@ -77,11 +72,9 @@ async fn main() -> Result<()> {
     let mux = Arc::new(Mux::new());
     mux.register(Arc::new(EchoerHandler::new(EchoServerImpl)))?;
 
-    loop {
-        let (stream, _) = listener.accept().await?;
-        let server = Server::with_arc(mux.clone());
-        tokio::spawn(async move {
-            let _ = server.handle_stream(stream).await;
-        });
-    }
+    // The server owns concurrent calls and their cancellation for this listener's lifetime.
+    let incoming = futures::stream::unfold(listener, |listener| async move {
+        Some((listener.accept().await.map(|(stream, _)| stream), listener))
+    });
+    Server::with_arc(mux).serve(Box::pin(incoming)).await
 }