Slightly refactor the low-level network client code.

src/network/client.rs

@@ -1,9 +1,9 @@
 use std::collections::BTreeMap;
-use std::fmt;
 use std::time::Duration;
+use std::{fmt, mem};
 
 use core::pin::Pin;
-use futures::future::{self, Either, Future};
+use futures::future::Future;
 use futures::sink::{Sink, SinkExt};
 use futures::stream::{Stream, StreamExt};
 use futures::task::{Context, Poll};
@@ -37,7 +37,8 @@ where
     let (cmd_send, cmd_receive) = mpsc::unbounded_channel();
     let (event_send, event_receive) = mpsc::channel(1);
     tokio::spawn(async move {
-        client_connection_task(stream, cmd_receive, event_send).await;
+        let task = ClientConnectionTask::new(stream, cmd_receive, event_send);
+        task.run().await;
     });
 
     (
@@ -53,131 +54,156 @@ enum ClientTaskCommand<PacketType, ErrorType> {
     SendPacket(CallData<PacketType, ErrorType>),
 }
 
-async fn client_connection_task<StreamType, ErrorType, PacketType>(
-    mut stream: StreamType,
-    mut cmd_receive: mpsc::UnboundedReceiver<ClientTaskCommand<PacketType, ErrorType>>,
-    mut event_send: mpsc::Sender<PacketType>,
-) where
+struct ClientConnectionTask<StreamType, ErrorType, PacketType> {
+    free_ids: IdAllocator,
+    calls_in_progress: BTreeMap<u32, oneshot::Sender<Result<PacketType, RPCError<ErrorType>>>>,
+    stream: StreamType,
+    cmd_receive: mpsc::UnboundedReceiver<ClientTaskCommand<PacketType, ErrorType>>,
+    event_send: mpsc::Sender<PacketType>,
+}
+
+impl<StreamType, ErrorType, PacketType> ClientConnectionTask<StreamType, ErrorType, PacketType>
+where
     PacketType: DeserializeOwned + Serialize,
     StreamType:
         Stream<Item = Result<Vec<u8>, ErrorType>> + Sink<Vec<u8>, Error = ErrorType> + Unpin,
     ErrorType: fmt::Debug,
 {
-    let mut free_ids = IdAllocator::new();
-    let mut calls_in_progress =
-        BTreeMap::<u32, oneshot::Sender<Result<PacketType, RPCError<ErrorType>>>>::new();
-
-    // TODO: Refactor this function.
+    fn new(
+        stream: StreamType,
+        cmd_receive: mpsc::UnboundedReceiver<ClientTaskCommand<PacketType, ErrorType>>,
+        event_send: mpsc::Sender<PacketType>,
+    ) -> Self {
+        Self {
+            free_ids: IdAllocator::new(),
+            calls_in_progress: BTreeMap::new(),
+            stream: stream,
+            cmd_receive: cmd_receive,
+            event_send,
+        }
+    }
 
-    let mut connection_closed = false;
-    loop {
-        match future::select(Box::pin(stream.next()), Box::pin(cmd_receive.recv())).await {
-            Either::Left((packet, _cmd_future)) => {
-                match packet {
-                    Some(Ok(data)) => {
-                        match rmps::decode::from_slice(&data) as Result<Packet<PacketType>, _> {
-                            Ok(deserialized) => {
-                                if deserialized.id == SERVER_EVENT_ID {
-                                    // We received an asynchronous event from the server.
-                                    event_send.send(deserialized.payload).await.ok();
-                                } else {
-                                    // We received a response to a call - find the call and forward
-                                    // the reply.
-                                    if let Some(call) = calls_in_progress.remove(&deserialized.id) {
-                                        free_ids.free(deserialized.id);
-                                        call.send(Ok(deserialized.payload)).ok();
-                                    } else {
-                                        error!(
-                                            "Received a reply for an unknown call ({:?})!",
-                                            deserialized.id
-                                        );
-                                    }
-                                }
-                            }
-                            Err(e) => {
-                                // We received an invalid packet. We cannot determine for which
-                                // call the packet was, so we just ignore it. The corresponding
-                                // call will freeze until the connection is dropped.
-                                error!("Invalid packet received: {:?}", e);
-                            }
-                        };
-                    }
-                    Some(Err(e)) => {
-                        error!("Network error: {:?}", e);
-                    }
-                    None => {
-                        // Stream closed. We still need to return errors for any
-                        // packet sent by the client, so do not exit the loop
-                        // immediately but rather return an error for any
-                        // incoming packet.
-                        connection_closed = true;
-                        break;
+    async fn run(mut self) {
+        let mut connection_closed = false;
+        loop {
+            tokio::select! {
+                packet = &mut self.stream.next() => {
+                    match packet {
+                        Some(Ok(data)) => {
+                            self.received_packet(&data).await;
+                        }
+                        Some(Err(e)) => {
+                            error!("Network error: {:?}", e);
+                        }
+                        None => {
+                            // Stream closed. We still need to return errors for any
+                            // packet sent by the client, so do not exit the loop
+                            // immediately but rather return an error for any
+                            // incoming packet.
+                            connection_closed = true;
+                            break;
+                        }
                     }
                 }
-            }
-            Either::Right((cmd, _stream_future)) => {
-                match cmd {
-                    Some(ClientTaskCommand::Close) => {
-                        // The RPC interface for this connection was dropped, so
-                        // we just exit the loop and close the connection.
-                        break;
-                    }
-                    Some(ClientTaskCommand::SendPacket(call)) => {
-                        // Add an ID for the outgoing packet.
-                        let id = free_ids.alloc();
-                        let packet = Packet {
-                            id: id,
-                            payload: call.data,
-                        };
-                        let mut serialized = Vec::new();
-                        packet
-                            .serialize(
-                                &mut Serializer::new(&mut serialized)
-                                    .with_struct_map()
-                                    .with_string_variants(),
-                            )
-                            .unwrap();
-
-                        // Send the serialized packet.
-                        // TODO: Send without flushing?
-                        match stream.send(serialized).await {
-                            Ok(()) => {
-                                calls_in_progress.insert(id, call.result);
-                            }
-                            Err(e) => {
-                                // We potentially leak the ID here, but that is
-                                // okay as the state of the call is not well
-                                // defined anymore.
-                                call.result.send(Err(RPCError::Stream(e))).ok();
-                            }
+                cmd = &mut self.cmd_receive.next() => {
+                    match cmd {
+                        Some(ClientTaskCommand::Close) => {
+                            // The RPC interface for this connection was dropped, so
+                            // we just exit the loop and close the connection.
+                            break;
+                        }
+                        Some(ClientTaskCommand::SendPacket(call)) => {
+                            self.send_packet(call).await;
+                        }
+                        None => {
+                            // This should never happen.
+                            break;
                         }
-                    }
-                    None => {
-                        // This should never happen.
-                        break;
                     }
                 }
             }
-        };
+        }
+
+        drop(self.event_send);
+
+        // Return errors for all pending packets as there will be no responses.
+        let mut calls_in_progress = BTreeMap::new();
+        mem::swap(&mut calls_in_progress, &mut self.calls_in_progress);
+        for (_, result) in calls_in_progress {
+            result.send(Err(RPCError::Closed)).ok();
+        }
+
+        if connection_closed {
+            // The stream was closed, but the command channel was not. We still need
+            // to return errors for any packet sent by the client, so do not exit
+            // the loop immediately but rather return an error for any incoming
+            // packet.
+            client_connection_closed(self.cmd_receive).await;
+        } else {
+            // The command channel was closed, but the stream was not. Gracefully
+            // close the stream.
+            // TODO
+        }
     }
 
-    drop(event_send);
+    async fn send_packet(&mut self, call: CallData<PacketType, ErrorType>) {
+        // Add an ID for the outgoing packet.
+        let id = self.free_ids.alloc();
+        let packet = Packet {
+            id: id,
+            payload: call.data,
+        };
+        let mut serialized = Vec::new();
+        packet
+            .serialize(
+                &mut Serializer::new(&mut serialized)
+                    .with_struct_map()
+                    .with_string_variants(),
+            )
+            .unwrap();
 
-    // Return errors for all pending packets as there will be no responses.
-    for (_, result) in calls_in_progress {
-        result.send(Err(RPCError::Closed)).ok();
+        // Send the serialized packet.
+        // TODO: Send without flushing?
+        match self.stream.send(serialized).await {
+            Ok(()) => {
+                self.calls_in_progress.insert(id, call.result);
+            }
+            Err(e) => {
+                // We potentially leak the ID here, but that is
+                // okay as the state of the call is not well
+                // defined anymore.
+                call.result.send(Err(RPCError::Stream(e))).ok();
+            }
+        }
     }
 
-    if connection_closed {
-        // The stream was closed, but the command channel was not. We still need
-        // to return errors for any packet sent by the client, so do not exit
-        // the loop immediately but rather return an error for any incoming
-        // packet.
-        client_connection_closed(cmd_receive).await;
-    } else {
-        // The command channel was closed, but the stream was not. Gracefully
-        // close the stream.
-        // TODO
+    async fn received_packet(&mut self, data: &[u8]) {
+        match rmps::decode::from_slice(data) as Result<Packet<PacketType>, _> {
+            Ok(deserialized) => {
+                if deserialized.id == SERVER_EVENT_ID {
+                    // We received an asynchronous event from the server.
+                    self.event_send.send(deserialized.payload).await.ok();
+                } else {
+                    // We received a response to a call - find the call and forward
+                    // the reply.
+                    if let Some(call) = self.calls_in_progress.remove(&deserialized.id) {
+                        self.free_ids.free(deserialized.id);
+                        call.send(Ok(deserialized.payload)).ok();
+                    } else {
+                        error!(
+                            "Received a reply for an unknown call ({:?})!",
+                            deserialized.id
+                        );
+                    }
+                }
+            }
+            Err(e) => {
+                // We received an invalid packet. We cannot determine for which
+                // call the packet was, so we just ignore it. The corresponding
+                // call will freeze until the connection is dropped.
+                error!("Invalid packet received: {:?}", e);
+            }
+        };
     }
 }