Extract non-generic code from base::Sender::send() and base::Receiver::recv()

- send_chunks_from_stream: streaming chunk send task extracted
- connect_ports_and_spawn: port connection + callback spawning extracted
- feed_recv_chunks: chunk feeding loop extracted from recv()
- receive_and_connect_ports: port acceptance + callback spawning extracted

Author: surban

remoc/src/rch/base/receiver.rs

@@ -194,6 +194,87 @@ impl<T, Codec> fmt::Debug for Receiver<T, Codec> {
     }
 }
 
+/// Result of feeding received chunks to the deserialization thread.
+enum FeedChunksResult {
+    /// All chunks have been fed successfully.
+    Done,
+    /// The current send operation was cancelled; restart.
+    Cancelled,
+    /// The chmux connection failed.
+    ChMux,
+    /// Maximum item size was exceeded.
+    MaxItemSizeExceeded,
+}
+
+/// Non-generic: feed received data chunks from chmux to the deserialization thread.
+///
+/// Separated from the generic `Receiver::recv` to avoid monomorphization.
+#[inline(never)]
+async fn feed_recv_chunks(
+    receiver: &mut chmux::Receiver,
+    tx: &tokio::sync::mpsc::Sender<Result<Bytes, ()>>,
+    total: &mut usize,
+    max_item_size: usize,
+) -> FeedChunksResult {
+    loop {
+        let tx_permit = match tx.reserve().await {
+            Ok(tx_permit) => tx_permit,
+            _ => return FeedChunksResult::Done,
+        };
+
+        match receiver.recv_chunk().await {
+            Ok(Some(chunk)) => {
+                *total += chunk.remaining();
+                if *total > max_item_size {
+                    return FeedChunksResult::MaxItemSizeExceeded;
+                }
+                tx_permit.send(Ok(chunk));
+            }
+            Ok(None) => return FeedChunksResult::Done,
+            Err(RecvChunkError::Cancelled) => return FeedChunksResult::Cancelled,
+            Err(RecvChunkError::ChMux) => return FeedChunksResult::ChMux,
+        }
+    }
+}
+
+/// Non-generic: receive port requests and spawn callbacks from deserialized objects.
+///
+/// Separated from the generic `Receiver::recv` to avoid monomorphization.
+///
+/// Returns `Ok(None)` on success, `Ok(Some(received))` if the received message
+/// was not port requests (requires restart), or `Err` on failure.
+#[inline(never)]
+async fn receive_and_connect_ports(
+    receiver: &mut chmux::Receiver,
+    pds: &mut PortDeserializer,
+    default_max_ports: usize,
+) -> Result<Option<Option<Received>>, RecvError> {
+    if !pds.expected.is_empty() {
+        receiver.set_max_ports(pds.expected.len() + default_max_ports);
+
+        let requests = match receiver.recv_any().await? {
+            Some(chmux::Received::Requests(requests)) => requests,
+            other => return Ok(Some(other)),
+        };
+
+        for request in requests {
+            if let Some((local_port, callback)) = pds.expected.remove(&request.id()) {
+                exec::spawn(callback(local_port, request).in_current_span());
+            }
+        }
+
+        if !pds.expected.is_empty() {
+            return Err(RecvError::MissingPorts(pds.expected.keys().copied().collect()));
+        }
+    }
+
+    for task in pds.tasks.drain(..) {
+        exec::spawn(task.in_current_span());
+    }
+
+    Ok(None)
+}
+
 enum DataSource<T> {
     None,
     Buffered(Option<chmux::DataBuf>),
@@ -293,46 +374,20 @@ where
 
                     // Observe deserialization of streamed data.
                     DataSource::Streamed { tx, task, total } => {
-                        enum FeedError {
-                            RecvChunkError(RecvChunkError),
-                            MaxItemSizeExceeded,
-                        }
-
-                        // Feed received data chunks to deserialization thread.
-                        if let Some(tx) = &tx {
-                            let res = loop {
-                                let tx_permit = match tx.reserve().await {
-                                    Ok(tx_permit) => tx_permit,
-                                    _ => {
-                                        break Ok(());
-                                    }
-                                };
-
-                                match self.receiver.recv_chunk().await {
-                                    Ok(Some(chunk)) => {
-                                        *total += chunk.remaining();
-                                        if *total > self.max_item_size {
-                                            break Err(FeedError::MaxItemSizeExceeded);
-                                        }
-
-                                        tx_permit.send(Ok(chunk));
-                                    }
-                                    Ok(None) => break Ok(()),
-                                    Err(err) => break Err(FeedError::RecvChunkError(err)),
-                                }
-                            };
-
-                            match res {
-                                Ok(()) => (),
-                                Err(FeedError::RecvChunkError(RecvChunkError::Cancelled)) => {
+                        // Feed received data chunks to deserialization thread (non-generic).
+                        if let Some(tx_ref) = &tx {
+                            match feed_recv_chunks(&mut self.receiver, tx_ref, total, self.max_item_size).await
+                            {
+                                FeedChunksResult::Done => (),
+                                FeedChunksResult::Cancelled => {
                                     self.data = DataSource::None;
                                     continue 'restart;
                                 }
-                                Err(FeedError::RecvChunkError(RecvChunkError::ChMux)) => {
+                                FeedChunksResult::ChMux => {
                                     self.data = DataSource::None;
                                     return Err(RecvError::Receive(chmux::RecvError::ChMux));
                                 }
-                                Err(FeedError::MaxItemSizeExceeded) => {
+                                FeedChunksResult::MaxItemSizeExceeded => {
                                     self.data = DataSource::None;
                                     return Err(RecvError::MaxItemSizeExceeded);
                                 }
@@ -365,47 +420,26 @@ where
                 }
             }
 
-            // Connect received ports.
+            // Connect received ports (non-generic).
             let pds = self.port_deser.as_mut().unwrap();
-            if !pds.expected.is_empty() {
-                // Set port limit.
-                //
-                // Allow the reception of additional ports for forward compatibility,
-                // i.e. our deserializer may use an older version of the struct
-                // which is missing some ports that the remote endpoint sent.
-                self.receiver.set_max_ports(pds.expected.len() + self.default_max_ports.unwrap());
-
-                // Receive port requests from chmux.
-                let requests = match self.receiver.recv_any().await? {
-                    Some(chmux::Received::Requests(requests)) => requests,
-                    other => {
-                        // Current send operation has been aborted and this is data from
-                        // next send operation, so we restart.
-                        self.recved = Some(other);
-                        self.data = DataSource::None;
-                        self.item = None;
-                        self.port_deser = None;
-                        continue 'restart;
-                    }
-                };
-
-                // Call port callbacks from received objects, ignoring superfluous requests for
-                // forward compatibility.
-                for request in requests {
-                    if let Some((local_port, callback)) = pds.expected.remove(&request.id()) {
-                        exec::spawn(callback(local_port, request).in_current_span());
-                    }
-                }
-
-                // But error on ports that we expect but that are missing.
-                if !pds.expected.is_empty() {
-                    return Err(RecvError::MissingPorts(pds.expected.keys().copied().collect()));
+            match receive_and_connect_ports(
+                &mut self.receiver,
+                pds,
+                self.default_max_ports.unwrap(),
+            )
+            .await
+            {
+                Ok(None) => (),
+                Ok(Some(other)) => {
+                    // Current send operation has been aborted and this is data from
+                    // next send operation, so we restart.
+                    self.recved = Some(other);
+                    self.data = DataSource::None;
+                    self.item = None;
+                    self.port_deser = None;
+                    continue 'restart;
                 }
-            }
-
-            // Spawn registered tasks.
-            for task in pds.tasks.drain(..) {
-                exec::spawn(task.in_current_span());
+                Err(err) => return Err(err),
             }
 
             return Ok(Some(self.item.take().unwrap()));

remoc/src/rch/base/sender.rs

@@ -226,6 +226,61 @@ impl PortSerializer {
     }
 }
 
+/// Non-generic: stream serialized chunks over chmux.
+///
+/// Separated from the generic `Sender::send` to avoid monomorphization.
+#[inline(never)]
+async fn send_chunks_from_stream(
+    sender: &mut chmux::Sender,
+    mut rx: tokio::sync::mpsc::Receiver<BytesMut>,
+    max_item_size: usize,
+) -> Result<(), SendErrorKind> {
+    let mut sc = sender.send_chunks();
+    let mut total = 0;
+    while let Some(chunk) = rx.recv().await {
+        total += chunk.len();
+        if total > max_item_size {
+            return Err(SendErrorKind::MaxItemSizeExceeded);
+        }
+        sc = sc.send(chunk.freeze()).await.map_err(SendErrorKind::Send)?;
+    }
+    sc.finish().await.map_err(SendErrorKind::Send)
+}
+
+/// Non-generic: connect ports obtained during serialization and spawn callbacks.
+///
+/// Separated from the generic `Sender::send` to avoid monomorphization.
+#[inline(never)]
+async fn connect_ports_and_spawn(
+    sender: &mut chmux::Sender,
+    ps: PortSerializer,
+) -> Result<(), chmux::SendError> {
+    let PortSerializer { requests, tasks, .. } = ps;
+
+    let mut ports = Vec::new();
+    let mut callbacks = Vec::new();
+    for (port, callback) in requests {
+        ports.push(PortReq::new(port));
+        callbacks.push(callback);
+    }
+
+    let connects = if ports.is_empty() {
+        Vec::new()
+    } else {
+        sender.connect(ports, true).await?
+    };
+
+    for (callback, connect) in callbacks.into_iter().zip(connects.into_iter()) {
+        exec::spawn(callback(connect).in_current_span());
+    }
+
+    for task in tasks {
+        exec::spawn(task.in_current_span());
+    }
+
+    Ok(())
+}
+
 /// Sends arbitrary values to a remote endpoint.
 ///
 /// Values may be or contain any channel from this crate.
@@ -365,54 +420,24 @@ where
 
             None => {
                 // Stream data while serializing.
-                let (tx, mut rx) = tokio::sync::mpsc::channel(BIG_DATA_CHUNK_QUEUE);
-                let ser_task = Self::serialize_streaming(
-                    self.sender.port_allocator(),
-                    self.sender.storage(),
-                    item,
-                    tx,
-                    self.sender.chunk_size(),
-                );
-
-                enum SendTaskError {
-                    SendError(chmux::SendError),
-                    MaxItemSizeExceeded,
-                }
-
-                let mut sc = self.sender.send_chunks();
-                let max_item_size = self.max_item_size;
-                let send_task = async move {
-                    let mut total = 0;
-                    while let Some(chunk) = rx.recv().await {
-                        total += chunk.len();
-                        if total > max_item_size {
-                            return Err(SendTaskError::MaxItemSizeExceeded);
-                        }
-
-                        sc = sc.send(chunk.freeze()).await.map_err(SendTaskError::SendError)?;
-                    }
-                    Ok(sc)
-                };
+                let (tx, rx) = tokio::sync::mpsc::channel(BIG_DATA_CHUNK_QUEUE);
+                let allocator = self.sender.port_allocator();
+                let storage = self.sender.storage();
+                let chunk_size = self.sender.chunk_size();
+                let ser_task = Self::serialize_streaming(allocator, storage, item, tx, chunk_size);
+                let send_task = send_chunks_from_stream(&mut self.sender, rx, self.max_item_size);
 
                 match tokio::join!(ser_task, send_task) {
-                    (Ok((item, ps, size)), Ok(sc)) => {
-                        if let Err(err) = sc.finish().await {
-                            return Err(SendError::new(SendErrorKind::Send(err), item));
-                        }
-
+                    (Ok((item, ps, size)), Ok(())) => {
                         if size <= self.sender.max_data_size() {
                             self.big_data = (self.big_data - 1).max(-BIG_DATA_LIMIT);
                         }
 
                         (item, ps)
                     }
-                    (Ok((item, _, _)), Err(err)) | (Err((_, item)), Err(err)) => {
+                    (Ok((item, _, _)), Err(kind)) | (Err((_, item)), Err(kind)) => {
                         // When sending fails, the serialization task will either finish
                         // or fail due to rx being dropped.
-                        let kind = match err {
-                            SendTaskError::SendError(err) => SendErrorKind::Send(err),
-                            SendTaskError::MaxItemSizeExceeded => SendErrorKind::MaxItemSizeExceeded,
-                        };
                         return Err(SendError::new(kind, item));
                     }
                     (Err((err, item)), _) => {
@@ -424,42 +449,14 @@ where
             }
         };
 
-        let PortSerializer { requests, tasks, .. } = ps;
-
-        // Extract ports and connect callbacks.
-        let mut ports = Vec::new();
-        let mut callbacks = Vec::new();
-        for (port, callback) in requests {
-            ports.push(PortReq::new(port));
-            callbacks.push(callback);
+        // Connect ports obtained during serialization (non-generic).
+        if let Err(err) = connect_ports_and_spawn(&mut self.sender, ps).await {
+            return Err(SendError::new(SendErrorKind::Send(err), item));
         }
 
-        // Request connecting chmux ports.
-        let connects = if ports.is_empty() {
-            Vec::new()
-        } else {
-            match self.sender.connect(ports, true).await {
-                Ok(connects) => connects,
-                Err(err) => return Err(SendError::new(SendErrorKind::Send(err), item)),
-            }
-        };
-
-        // Ensure that item is dropped before calling connection callbacks.
+        // Ensure that item is dropped before connection callbacks run.
         drop(item);
 
-        // Call callbacks of BaseSenders and BaseReceivers with obtained
-        // chmux connect requests.
-        //
-        // We have to spawn a task for this to ensure cancellation safety.
-        for (callback, connect) in callbacks.into_iter().zip(connects.into_iter()) {
-            exec::spawn(callback(connect).in_current_span());
-        }
-
-        // Spawn registered tasks.
-        for task in tasks {
-            exec::spawn(task.in_current_span());
-        }
-
         Ok(())
     }