starpc 0.52.0 → 0.52.2

package/srpc/rpc.rs

@@ -1,14 +1,16 @@
 //! RPC state machines for client and server.
 //!
 //! This module provides the core RPC state machines that manage the lifecycle
-//! of RPC calls, matching the behavior of the Go and TypeScript implementations.
+//! of calls over the shared Go, TypeScript, and Rust wire protocol.
+
+use std::sync::{
+    atomic::{AtomicBool, Ordering},
+    Arc, Mutex,
+};
 
 use async_trait::async_trait;
 use bytes::Bytes;
-use std::collections::VecDeque;
-use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::Arc;
-use tokio::sync::{Mutex, Notify};
+use tokio::sync::{mpsc, Notify};
 
 use crate::error::{Error, Result};
 use crate::packet::{new_call_cancel, new_call_data_full, new_call_start, Validate};
@@ -44,7 +46,6 @@ pub struct CommonRpc {
     method: String,
 
     /// Whether we have completed locally (sent complete/cancel).
-    /// Note: not guarded by the mutex, uses atomic operations.
     local_completed: AtomicBool,
 
     /// Packet writer.
@@ -53,21 +54,43 @@ pub struct CommonRpc {
     /// Notification for state changes.
     notify: Notify,
 
+    /// Bounded incoming data, shared by the packet reader and RPC consumer.
+    data: mpsc::Sender<Bytes>,
+
     /// Internal state protected by mutex.
     state: Mutex<RpcState>,
 }
 
 /// Internal RPC state.
 struct RpcState {
-    /// Queue of incoming data messages.
-    /// Note: messages may be len() == 0 (empty data with data_is_zero).
-    data_queue: VecDeque<Bytes>,
+    /// Incoming messages, including empty payloads, retained until read.
+    data: mpsc::Receiver<Bytes>,
+
+    /// How the incoming side ended, once it has.
+    end: Option<RpcEnd>,
+}
+
+/// How the incoming side of an RPC ended.
+enum RpcEnd {
+    /// The remote completed the call, or the local side closed it.
+    Complete,
 
-    /// Whether the remote side has closed the stream.
-    data_closed: bool,
+    /// The remote failed or cancelled the call.
+    Remote(String),
 
-    /// Error from the remote side, if any.
-    remote_err: Option<String>,
+    /// The transport closed without a remote completion or error.
+    ClosedBeforeCompletion,
+}
+
+impl RpcEnd {
+    /// Returns the error a read reports after the queued messages.
+    fn error(&self) -> Error {
+        match self {
+            RpcEnd::Complete => Error::StreamClosed,
+            RpcEnd::Remote(err) => Error::Remote(err.clone()),
+            RpcEnd::ClosedBeforeCompletion => Error::ClosedBeforeCompletion,
+        }
+    }
 }
 
 impl CommonRpc {
@@ -78,6 +101,8 @@ impl CommonRpc {
         method: String,
         writer: Arc<dyn PacketWriter>,
     ) -> Self {
+        // Retain one unread message; the packet reader waits before accepting another.
+        let (data, receiver) = mpsc::channel(1);
         Self {
             ctx,
             service,
@@ -85,10 +110,10 @@ impl CommonRpc {
             local_completed: AtomicBool::new(false),
             writer,
             notify: Notify::new(),
+            data,
             state: Mutex::new(RpcState {
-                data_queue: VecDeque::new(),
-                data_closed: false,
-                remote_err: None,
+                data: receiver,
+                end: None,
             }),
         }
     }
@@ -113,32 +138,34 @@ impl CommonRpc {
         self.local_completed.load(Ordering::SeqCst)
     }
 
-    /// Waits for the RPC to finish (remote end closed the stream).
+    /// Waits for the remote verdict after preceding payloads have been accepted.
     ///
-    /// This matches the Go implementation's `Wait(ctx context.Context) error`.
+    /// Streaming consumers must receive data concurrently so the packet reader
+    /// can reach completion beyond a full receive channel.
     pub async fn wait(&self) -> Result<()> {
         loop {
+            // Capture the notification before reading state so a concurrent
+            // completion cannot fall between the check and the wait.
+            let changed = self.notify.notified();
             // Check current state
             {
-                let state = self.state.lock().await;
+                let state = self.state.lock().unwrap();
 
                 // Check cancellation first - local cancellation takes priority
                 if self.ctx.is_cancelled() {
                     return Err(Error::Cancelled);
                 }
 
-                if let Some(ref err) = state.remote_err {
-                    return Err(Error::Remote(err.clone()));
-                }
-
-                if state.data_closed {
-                    return Ok(());
+                match state.end {
+                    Some(RpcEnd::Complete) => return Ok(()),
+                    Some(ref end) => return Err(end.error()),
+                    None => {}
                 }
             }
 
             // Wait for notification or cancellation
             tokio::select! {
-                _ = self.notify.notified() => continue,
+                _ = changed => continue,
                 _ = self.ctx.cancelled() => return Err(Error::Cancelled),
             }
         }
@@ -146,36 +173,43 @@ impl CommonRpc {
 
     /// Reads one message from the data queue, blocking until available.
     ///
-    /// Returns `Err(Error::StreamClosed)` if the stream ended without a packet.
-    /// This matches the Go implementation which returns `io.EOF`.
+    /// After the queued messages, returns `Err(Error::StreamClosed)` if the
+    /// call completed, matching `io.EOF` in the Go implementation, or the
+    /// error that ended it.
     pub async fn read_one(&self) -> Result<Bytes> {
         loop {
+            // notify_waiters preserves notifications from this future's creation.
+            let changed = self.notify.notified();
             // Try to get a message from the queue first, before checking cancellation.
             // This ensures we drain any pending messages even if the context is cancelled.
             {
-                let mut state = self.state.lock().await;
+                let mut state = self.state.lock().unwrap();
 
-                if let Some(data) = state.data_queue.pop_front() {
+                if let Ok(data) = state.data.try_recv() {
                     return Ok(data);
                 }
 
-                // Check if the stream is closed (graceful close from remote)
-                if state.data_closed {
-                    if let Some(ref err) = state.remote_err {
-                        return Err(Error::Remote(err.clone()));
-                    }
-                    return Err(Error::StreamClosed);
+                if let Some(ref end) = state.end {
+                    return Err(end.error());
                 }
             }
 
             // Now check for cancellation - only if no data is available
             // and the stream isn't properly closed.
             if self.ctx.is_cancelled() {
-                // If context cancelled and data not closed, close it now
-                let mut state = self.state.lock().await;
-                if !state.data_closed {
-                    state.data_closed = true;
-                    drop(state);
+                // Settle cancellation before releasing the writer outside the lock.
+                let close = {
+                    let mut state = self.state.lock().unwrap();
+                    if state.end.is_none() {
+                        state.end = Some(RpcEnd::Complete);
+                        state.data.close();
+                        true
+                    } else {
+                        false
+                    }
+                };
+                if close {
+                    // Cancellation determines the result even if transport close fails.
                     let _ = self.writer.close().await;
                     self.ctx.cancel();
                     self.notify.notify_waiters();
@@ -185,7 +219,7 @@ impl CommonRpc {
 
             // Wait for notification or cancellation
             tokio::select! {
-                _ = self.notify.notified() => continue,
+                _ = changed => continue,
                 _ = self.ctx.cancelled() => {
                     // Loop will handle the cancellation
                     continue;
@@ -194,12 +228,9 @@ impl CommonRpc {
         }
     }
 
-    /// Writes a CallData packet.
-    ///
-    /// # Arguments
-    /// * `data` - Optional data to send
-    /// * `complete` - Whether this completes the RPC
-    /// * `error` - Optional error message
+    /// Writes a CallData packet, completing the call when either `complete`
+    /// or `error` is set. Returns `Error::Completed` once the call has
+    /// already completed locally.
     pub async fn write_call_data(
         &self,
         data: Option<Bytes>,
@@ -242,34 +273,51 @@ impl CommonRpc {
         self.writer.write_packet(new_call_cancel()).await
     }
 
-    /// Handles an incoming CallData packet.
+    /// Accepts incoming data, waiting for the consumer when one message is unread.
+    ///
+    /// The transport reader retains at most one further packet while waiting.
+    /// Cancellation or closing the RPC releases a blocked packet reader. Completion
+    /// remains ordered after its data, and control-only packets need no capacity.
     pub async fn handle_call_data(&self, call_data: CallData) -> Result<()> {
-        let mut state = self.state.lock().await;
-
-        // Check if already closed
-        if state.data_closed {
-            // If the packet is just indicating the call is complete, ignore it
-            // This matches Go behavior
-            if call_data.complete {
-                return Ok(());
-            }
-            return Err(Error::Completed);
+        // Reserve capacity without holding the state lock needed by the consumer.
+        let pending =
+            if let Some(data) = decode_optional_data(call_data.data, call_data.data_is_zero) {
+                let result = tokio::select! {
+                    biased;
+                    result = self.data.reserve() => result,
+                    () = self.ctx.cancelled() => return Err(Error::Cancelled),
+                };
+                match result {
+                    Ok(permit) => Some((data, permit)),
+                    Err(_) if call_data.complete => return Ok(()),
+                    Err(_) => return Err(Error::Completed),
+                }
+            } else {
+                None
+            };
+
+        // Closing and packet delivery share the lock, so no data arrives after an end.
+        let mut state = self.state.lock().unwrap();
+        if state.end.is_some() {
+            return if call_data.complete {
+                Ok(())
+            } else {
+                Err(Error::Completed)
+            };
         }
-
-        // Extract data if present
-        if let Some(data) = decode_optional_data(call_data.data, call_data.data_is_zero) {
-            state.data_queue.push_back(data);
+        if let Some((data, permit)) = pending {
+            permit.send(data);
         }
 
-        // Handle completion or error
+        // Publish the terminal verdict after its final payload and wake blocked readers.
         if !call_data.error.is_empty() {
-            state.remote_err = Some(call_data.error);
-            state.data_closed = true;
+            state.end = Some(RpcEnd::Remote(call_data.error));
         } else if call_data.complete {
-            state.data_closed = true;
+            state.end = Some(RpcEnd::Complete);
+        }
+        if state.end.is_some() {
+            state.data.close();
         }
-
-        // Notify waiters
         drop(state);
         self.notify.notify_waiters();
 
@@ -278,22 +326,29 @@ impl CommonRpc {
 
     /// Handles a CallCancel packet.
     pub async fn handle_call_cancel(&self) -> Result<()> {
-        self.handle_stream_close(Some("cancelled".to_string())).await
+        self.handle_stream_close(Some("cancelled".to_string()))
+            .await
     }
 
     /// Handles stream close from the transport.
+    ///
+    /// A clean close with no completion behind it leaves the call without a
+    /// verdict, so reads report `Error::ClosedBeforeCompletion` rather than the
+    /// clean end of the stream.
     pub async fn handle_stream_close(&self, err: Option<String>) -> Result<()> {
-        let mut state = self.state.lock().await;
-
-        if let Some(e) = err {
-            if state.remote_err.is_none() {
-                state.remote_err = Some(e);
+        // Preserve a verdict already received and release blocked packet producers.
+        {
+            let mut state = self.state.lock().unwrap();
+            if state.end.is_none() {
+                state.end = Some(match err {
+                    Some(err) => RpcEnd::Remote(err),
+                    None => RpcEnd::ClosedBeforeCompletion,
+                });
             }
+            state.data.close();
         }
-        state.data_closed = true;
-
-        drop(state);
 
+        // The incoming verdict determines reads even if closing the writer fails.
         let _ = self.writer.close().await;
         self.ctx.cancel();
         self.notify.notify_waiters();
@@ -301,22 +356,28 @@ impl CommonRpc {
         Ok(())
     }
 
-    /// Closes the RPC, releasing resources.
+    /// Ends the call locally and releases its resources.
     ///
-    /// This is called internally and handles cleanup.
-    /// Note: We don't set remote_err here because local cancellation is signaled
-    /// via ctx.is_cancelled(), which is checked first in wait()/read_one().
-    async fn close_locked(&self) {
-        let mut state = self.state.lock().await;
-        state.data_closed = true;
-        self.local_completed.store(true, Ordering::SeqCst);
-        // Don't set remote_err for local closure - rely on ctx.is_cancelled() instead
-        // to distinguish between local cancellation and remote errors.
-        drop(state);
+    /// A call with no verdict ends as complete, pending waiters and reads
+    /// wake, and the writer closes before the context cancels. Returns the
+    /// writer's close error; every caller still finds the call ended.
+    async fn close_local(&self) -> Result<()> {
+        // End data reception and release packet producers before transport shutdown.
+        {
+            let mut state = self.state.lock().unwrap();
+            if state.end.is_none() {
+                state.end = Some(RpcEnd::Complete);
+            }
+            self.local_completed.store(true, Ordering::SeqCst);
+            state.data.close();
+        }
 
-        let _ = self.writer.close().await;
+        // Wake observers even when the writer fails to close.
+        let closed = self.writer.close().await;
         self.notify.notify_waiters();
         self.ctx.cancel();
+
+        closed
     }
 }
 
@@ -356,7 +417,7 @@ impl ClientRpc {
         self.common.method()
     }
 
-    /// Waits for the RPC to finish.
+    /// Waits for the remote verdict; receive streaming payloads concurrently.
     pub async fn wait(&self) -> Result<()> {
         self.common.wait().await
     }
@@ -429,7 +490,7 @@ impl ClientRpc {
         let _ = self.common.write_call_cancel().await;
 
         // Close resources
-        self.common.close_locked().await;
+        let _ = self.common.close_local().await;
     }
 }
 
@@ -440,9 +501,7 @@ impl Stream for ClientRpc {
     }
 
     async fn send_bytes(&self, data: Bytes) -> Result<()> {
-        self.common
-            .write_call_data(Some(data), false, None)
-            .await
+        self.common.write_call_data(Some(data), false, None).await
     }
 
     async fn recv_bytes(&self) -> Result<Bytes> {
@@ -496,7 +555,7 @@ impl ServerRpc {
         self.common.method()
     }
 
-    /// Waits for the RPC to finish.
+    /// Waits for the remote verdict; receive streaming payloads concurrently.
     pub async fn wait(&self) -> Result<()> {
         self.common.wait().await
     }
@@ -536,15 +595,13 @@ impl Stream for ServerRpc {
     }
 
     async fn send_bytes(&self, data: Bytes) -> Result<()> {
-        self.common
-            .write_call_data(Some(data), false, None)
-            .await
+        self.common.write_call_data(Some(data), false, None).await
     }
 
     async fn recv_bytes(&self) -> Result<Bytes> {
         // First check for initial data
         {
-            let mut initial = self.initial_data.lock().await;
+            let mut initial = self.initial_data.lock().unwrap();
             if let Some(data) = initial.take() {
                 return Ok(data);
             }
@@ -559,31 +616,26 @@ impl Stream for ServerRpc {
     }
 
     async fn close(&self) -> Result<()> {
-        // Mark as locally completed
-        self.common.local_completed.store(true, Ordering::SeqCst);
-
-        // Close the writer
-        self.common.writer.close().await?;
-
-        // Cancel the context
-        self.common.ctx.cancel();
-
-        Ok(())
+        self.common.close_local().await
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use super::*;
     use std::sync::Mutex as StdMutex;
 
-    /// Mock packet writer for testing.
+    use super::*;
+
+    /// Records outgoing packets and whether the transport was closed.
     struct MockWriter {
+        /// Packets written in order.
         packets: StdMutex<Vec<Packet>>,
+        /// Whether close has run.
         closed: AtomicBool,
     }
 
     impl MockWriter {
+        /// Creates an open writer with no packets.
         fn new() -> Self {
             Self {
                 packets: StdMutex::new(Vec::new()),
@@ -591,10 +643,12 @@ mod tests {
             }
         }
 
+        /// Copies the recorded wire packets.
         fn packets(&self) -> Vec<Packet> {
             self.packets.lock().unwrap().clone()
         }
 
+        /// Reports transport closure.
         fn is_closed(&self) -> bool {
             self.closed.load(Ordering::SeqCst)
         }
@@ -613,11 +667,18 @@ mod tests {
         }
     }
 
+    /// Starting a call transmits its service, method, and initial data.
     #[tokio::test]
     async fn test_client_rpc_start() {
+        // Establish the call and its transport recorder.
         let writer = Arc::new(MockWriter::new());
         let ctx = Context::new();
-        let rpc = ClientRpc::new(ctx, "test.Service".into(), "TestMethod".into(), writer.clone());
+        let rpc = ClientRpc::new(
+            ctx,
+            "test.Service".into(),
+            "TestMethod".into(),
+            writer.clone(),
+        );
 
         rpc.start(Some(Bytes::from(vec![1, 2, 3]))).await.unwrap();
 
@@ -635,8 +696,10 @@ mod tests {
         }
     }
 
+    /// A client call sends its start packet once.
     #[tokio::test]
     async fn test_client_rpc_double_start_fails() {
+        // Establish the call and its transport recorder.
         let writer = Arc::new(MockWriter::new());
         let ctx = Context::new();
         let rpc = ClientRpc::new(ctx, "test.Service".into(), "TestMethod".into(), writer);
@@ -646,11 +709,18 @@ mod tests {
         assert!(matches!(result, Err(Error::Completed)));
     }
 
+    /// Closing a started call sends cancellation and closes the transport.
     #[tokio::test]
     async fn test_client_rpc_close_sends_cancel() {
+        // Establish the call and its transport recorder.
         let writer = Arc::new(MockWriter::new());
         let ctx = Context::new();
-        let rpc = ClientRpc::new(ctx, "test.Service".into(), "TestMethod".into(), writer.clone());
+        let rpc = ClientRpc::new(
+            ctx,
+            "test.Service".into(),
+            "TestMethod".into(),
+            writer.clone(),
+        );
 
         rpc.start(None).await.unwrap();
         rpc.close().await;
@@ -663,8 +733,10 @@ mod tests {
         assert!(writer.is_closed());
     }
 
+    /// A server delivers the initial payload before subsequent data.
     #[tokio::test]
     async fn test_server_rpc_from_call_start() {
+        // Establish the call and its transport recorder.
         let call_start = CallStart {
             rpc_service: "test.Service".into(),
             rpc_method: "TestMethod".into(),
@@ -684,8 +756,37 @@ mod tests {
         assert_eq!(&data[..], &[1, 2, 3]);
     }
 
+    /// Local closure rejects data arriving after the call ends.
+    #[tokio::test]
+    async fn test_server_rpc_close_ends_the_call() {
+        // Establish the call and its transport recorder.
+        let call_start = CallStart {
+            rpc_service: "test.Service".into(),
+            rpc_method: "TestMethod".into(),
+            data: Vec::new(),
+            data_is_zero: false,
+        };
+        let writer = Arc::new(MockWriter::new());
+        let rpc = ServerRpc::from_call_start(Context::new(), call_start, writer.clone());
+
+        // Closing releases the writer and cancels the call.
+        rpc.close().await.unwrap();
+        assert!(writer.is_closed());
+        assert!(rpc.context().is_cancelled());
+
+        // The call has an end state, so later data is refused instead of queued.
+        let late = CallData {
+            data: vec![1],
+            ..Default::default()
+        };
+        let result = rpc.common.handle_call_data(late).await;
+        assert!(matches!(result, Err(Error::Completed)));
+    }
+
+    /// Receiving preserves payload bytes.
     #[tokio::test]
     async fn test_common_rpc_read_one_with_data() {
+        // Establish the call and its transport recorder.
         let writer = Arc::new(MockWriter::new());
         let ctx = Context::new();
         let rpc = CommonRpc::new(ctx, "svc".into(), "method".into(), writer);
@@ -703,8 +804,10 @@ mod tests {
         assert_eq!(&data[..], &[1, 2, 3]);
     }
 
+    /// Completion without data ends reads cleanly.
     #[tokio::test]
     async fn test_common_rpc_read_one_stream_closed() {
+        // Establish the call and its transport recorder.
         let writer = Arc::new(MockWriter::new());
         let ctx = Context::new();
         let rpc = CommonRpc::new(ctx, "svc".into(), "method".into(), writer);
@@ -722,8 +825,10 @@ mod tests {
         assert!(matches!(result, Err(Error::StreamClosed)));
     }
 
+    /// A remote failure preserves its error text.
     #[tokio::test]
     async fn test_common_rpc_read_one_with_error() {
+        // Establish the call and its transport recorder.
         let writer = Arc::new(MockWriter::new());
         let ctx = Context::new();
         let rpc = CommonRpc::new(ctx, "svc".into(), "method".into(), writer);
@@ -744,8 +849,10 @@ mod tests {
         }
     }
 
+    /// Local completion rejects subsequent data writes.
     #[tokio::test]
     async fn test_write_call_data_after_complete() {
+        // Establish the call and its transport recorder.
         let writer = Arc::new(MockWriter::new());
         let ctx = Context::new();
         let rpc = CommonRpc::new(ctx, "svc".into(), "method".into(), writer);
@@ -754,12 +861,16 @@ mod tests {
         rpc.write_call_data(None, true, None).await.unwrap();
 
         // Trying to send more data should fail
-        let result = rpc.write_call_data(Some(Bytes::from(vec![1])), false, None).await;
+        let result = rpc
+            .write_call_data(Some(Bytes::from(vec![1])), false, None)
+            .await;
         assert!(matches!(result, Err(Error::Completed)));
     }
 
+    /// Cancellation sends exactly one wire packet.
     #[tokio::test]
     async fn test_write_call_cancel() {
+        // Establish the call and its transport recorder.
         let writer = Arc::new(MockWriter::new());
         let ctx = Context::new();
         let rpc = CommonRpc::new(ctx, "svc".into(), "method".into(), writer.clone());
@@ -774,4 +885,110 @@ mod tests {
         let result = rpc.write_call_cancel().await;
         assert!(matches!(result, Err(Error::Completed)));
     }
+
+    /// A full receive channel pauses packet processing until the consumer reads.
+    #[tokio::test]
+    async fn receive_capacity_preserves_data_and_terminal_order() {
+        // Fill the receive channel with an empty but present message.
+        let rpc = CommonRpc::new(
+            Context::new(),
+            "svc".into(),
+            "method".into(),
+            Arc::new(MockWriter::new()),
+        );
+        rpc.handle_call_data(CallData {
+            data_is_zero: true,
+            ..Default::default()
+        })
+        .await
+        .unwrap();
+
+        // The final payload waits for capacity without publishing completion early.
+        let final_packet = rpc.handle_call_data(CallData {
+            data: b"last".to_vec(),
+            complete: true,
+            error: "finished".into(),
+            ..Default::default()
+        });
+        tokio::pin!(final_packet);
+        assert!(futures::poll!(&mut final_packet).is_pending());
+        let wait = rpc.wait();
+        tokio::pin!(wait);
+        assert!(futures::poll!(&mut wait).is_pending());
+
+        // Draining the prior message admits the final payload before its error.
+        assert!(rpc.read_one().await.unwrap().is_empty());
+        final_packet.await.unwrap();
+        assert_eq!(rpc.read_one().await.unwrap().as_ref(), b"last");
+        assert!(matches!(rpc.read_one().await, Err(Error::Remote(err)) if err == "finished"));
+        assert!(matches!(wait.await, Err(Error::Remote(err)) if err == "finished"));
+    }
+
+    /// Cancellation and closure both wake a packet reader waiting for capacity.
+    #[tokio::test]
+    async fn full_receive_channel_releases_on_cancel_and_close() {
+        // Exercise context cancellation, local close, and transport loss separately.
+        for end in 0..3 {
+            let rpc = CommonRpc::new(
+                Context::new(),
+                "svc".into(),
+                "method".into(),
+                Arc::new(MockWriter::new()),
+            );
+            rpc.handle_call_data(CallData {
+                data: vec![1],
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+            let pending = rpc.handle_call_data(CallData {
+                data: vec![2],
+                ..Default::default()
+            });
+            tokio::pin!(pending);
+            assert!(futures::poll!(&mut pending).is_pending());
+
+            // End the RPC through its public lifetime operation while data remains unread.
+            match end {
+                0 => rpc.context().cancel(),
+                1 => rpc.close_local().await.unwrap(),
+                _ => rpc.handle_stream_close(None).await.unwrap(),
+            }
+            let result = tokio::time::timeout(std::time::Duration::from_secs(2), pending)
+                .await
+                .unwrap();
+            assert!(matches!(result, Err(Error::Cancelled | Error::Completed)));
+            assert_eq!(rpc.read_one().await.unwrap().as_ref(), &[1]);
+            assert!(rpc.read_one().await.is_err());
+        }
+    }
+
+    /// A completion with no payload bypasses a full data channel.
+    #[tokio::test]
+    async fn completion_does_not_need_receive_capacity() {
+        // Leave one unread payload when the peer completes the stream.
+        let rpc = CommonRpc::new(
+            Context::new(),
+            "svc".into(),
+            "method".into(),
+            Arc::new(MockWriter::new()),
+        );
+        rpc.handle_call_data(CallData {
+            data: vec![1],
+            ..Default::default()
+        })
+        .await
+        .unwrap();
+        rpc.handle_call_data(CallData {
+            complete: true,
+            ..Default::default()
+        })
+        .await
+        .unwrap();
+
+        // Wait observes completion, while receive preserves the payload preceding it.
+        rpc.wait().await.unwrap();
+        assert_eq!(rpc.read_one().await.unwrap().as_ref(), &[1]);
+        assert!(matches!(rpc.read_one().await, Err(Error::StreamClosed)));
+    }
 }