refactor: use two stage accept (#87)

Author: rklaehn

examples/errors.rs

@@ -60,7 +60,7 @@ async fn main() -> anyhow::Result<()> {
     let server = RpcServer::new(server);
     let handle = tokio::task::spawn(async move {
         for _ in 0..1 {
-            let (req, chan) = server.accept().await?;
+            let (req, chan) = server.accept().await?.read_first().await?;
             match req {
                 IoRequest::Write(req) => chan.rpc_map_err(req, fs, Fs::write).await,
             }?

examples/modularize.rs

@@ -34,7 +34,10 @@ async fn main() -> Result<()> {
 async fn run_server<C: ServiceEndpoint<AppService>>(server_conn: C, handler: app::Handler) {
     let server = RpcServer::new(server_conn);
     loop {
-        match server.accept().await {
+        let Ok(accepting) = server.accept().await else {
+            continue;
+        };
+        match accepting.read_first().await {
             Err(err) => warn!(?err, "server accept failed"),
             Ok((req, chan)) => {
                 let handler = handler.clone();

examples/store.rs

@@ -168,7 +168,7 @@ async fn main() -> anyhow::Result<()> {
         let s = server;
         let store = Store;
         loop {
-            let (req, chan) = s.accept().await?;
+            let (req, chan) = s.accept().await?.read_first().await?;
             use StoreRequest::*;
             let store = store.clone();
             #[rustfmt::skip]

src/lib.rs

@@ -65,7 +65,7 @@
 //! let handler = Handler;
 //! loop {
 //!   // accept connections
-//!   let (msg, chan) = server.accept().await?;
+//!   let (msg, chan) = server.accept().await?.read_first().await?;
 //!   // dispatch the message to the appropriate handler
 //!   match msg {
 //!     PingRequest::Ping(ping) => chan.rpc(ping, handler, Handler::ping).await?,

src/server.rs

@@ -117,8 +117,14 @@ where
     }
 }
 
-impl<S: Service, C: ServiceEndpoint<S>> RpcServer<S, C> {
-    /// Accepts a new channel from a client, and reads the first request.
+/// The result of accepting a new connection.
+pub struct Accepting<S: Service, C: ServiceEndpoint<S>> {
+    send: C::SendSink,
+    recv: C::RecvStream,
+}
+
+impl<S: Service, C: ServiceEndpoint<S>> Accepting<S, C> {
+    /// Read the first message from the client.
     ///
     /// The return value is a tuple of `(request, channel)`.  Here `request` is the
     /// first request which is already read from the stream.  The `channel` is a
@@ -127,13 +133,8 @@ impl<S: Service, C: ServiceEndpoint<S>> RpcServer<S, C> {
     ///
     /// Often sink and stream will wrap an an underlying byte stream. In this case you can
     /// call into_inner() on them to get it back to perform byte level reads and writes.
-    pub async fn accept(&self) -> result::Result<(S::Req, RpcChannel<S, C>), RpcServerError<C>> {
-        let (send, mut recv) = self
-            .source
-            .accept_bi()
-            .await
-            .map_err(RpcServerError::Accept)?;
-
+    pub async fn read_first(self) -> result::Result<(S::Req, RpcChannel<S, C>), RpcServerError<C>> {
+        let Accepting { send, mut recv } = self;
         // get the first message from the client. This will tell us what it wants to do.
         let request: S::Req = recv
             .next()
@@ -144,6 +145,19 @@ impl<S: Service, C: ServiceEndpoint<S>> RpcServer<S, C> {
             .map_err(RpcServerError::RecvError)?;
         Ok((request, RpcChannel::new(send, recv)))
     }
+}
+
+impl<S: Service, C: ServiceEndpoint<S>> RpcServer<S, C> {
+    /// Accepts a new channel from a client. The result is an [Accepting] object that
+    /// can be used to read the first request.
+    pub async fn accept(&self) -> result::Result<Accepting<S, C>, RpcServerError<C>> {
+        let (send, recv) = self
+            .source
+            .accept_bi()
+            .await
+            .map_err(RpcServerError::Accept)?;
+        Ok(Accepting { send, recv })
+    }
 
     /// Get the underlying service endpoint
     pub fn into_inner(self) -> C {
@@ -309,7 +323,7 @@ where
 {
     let server = RpcServer::<S, C>::new(conn);
     loop {
-        let (req, chan) = server.accept().await?;
+        let (req, chan) = server.accept().await?.read_first().await?;
         let target = target.clone();
         handler(chan, req, target).await?;
     }

tests/flume.rs

@@ -67,7 +67,7 @@ async fn flume_channel_mapped_bench() -> anyhow::Result<()> {
         tokio::task::spawn(async move {
             let service = ComputeService;
             loop {
-                let (req, chan) = server.accept().await?;
+                let (req, chan) = server.accept().await?.read_first().await?;
                 let service = service.clone();
                 tokio::spawn(async move {
                     let req: OuterRequest = req;

tests/hyper.rs

@@ -176,8 +176,10 @@ async fn hyper_channel_errors() -> anyhow::Result<()> {
         let (res_tx, res_rx) = flume::unbounded();
         let handle = tokio::spawn(async move {
             loop {
-                let x = server.accept().await;
-                let res = match x {
+                let Ok(x) = server.accept().await else {
+                    continue;
+                };
+                let res = match x.read_first().await {
                     Ok((req, chan)) => match req {
                         TestRequest::BigRequest(req) => {
                             chan.rpc(req, TestService, TestService::big).await

tests/math.rs

@@ -167,7 +167,7 @@ impl ComputeService {
         let s = server;
         let service = ComputeService;
         loop {
-            let (req, chan) = s.accept().await?;
+            let (req, chan) = s.accept().await?.read_first().await?;
             let service = service.clone();
             tokio::spawn(async move { Self::handle_rpc_request(service, req, chan).await });
         }
@@ -206,7 +206,7 @@ impl ComputeService {
         let service = ComputeService;
         while received < count {
             received += 1;
-            let (req, chan) = s.accept().await?;
+            let (req, chan) = s.accept().await?.read_first().await?;
             let service = service.clone();
             tokio::spawn(async move {
                 use ComputeRequest::*;
@@ -236,7 +236,7 @@ impl ComputeService {
         let service = ComputeService;
         let request_stream = stream! {
             loop {
-                yield s2.accept().await;
+                yield s2.accept().await?.read_first().await;
             }
         };
         let process_stream = request_stream.map(move |r| {

tests/slow_math.rs

@@ -113,7 +113,7 @@ impl ComputeService {
         let s = server;
         let service = ComputeService;
         loop {
-            let (req, chan) = s.accept().await?;
+            let (req, chan) = s.accept().await?.read_first().await?;
             use ComputeRequest::*;
             let service = service.clone();
             #[rustfmt::skip]

tests/try.rs

@@ -77,7 +77,7 @@ async fn try_server_streaming() -> anyhow::Result<()> {
     let server = RpcServer::<TryService, _>::new(server);
     let server_handle = tokio::task::spawn(async move {
         loop {
-            let (req, chan) = server.accept().await?;
+            let (req, chan) = server.accept().await?.read_first().await?;
             let handler = Handler;
             match req {
                 TryRequest::StreamN(req) => {