Ignore client messages after stopping the IO task (#1126)

This protects us from a subtle race condition:

- The main task decides to stop the client task (e.g. because we've hit `IO_OUTSTANDING_MAX`)
  - It sends a stop request down the `tokio::sync::oneshot` channel
  - Then, it marks all jobs as skipped, which may result in discarding some of them
- The IO task receives a message from the downstairs
  - It passes this message to the main task
  - This message is a reply to one of the jobs that we skipped
  - The main task gets mad

Author: mkeeter

upstairs/src/client.rs

@@ -285,13 +285,21 @@ impl DownstairsClient {
     /// must the select expressions be cancel safe, but the **bodies** must also
     /// be cancel-safe.  This is why we simply return a single value in the body
     /// of each statement.
+    ///
+    /// This function will wait forever if we have asked for the client task to
+    /// stop, so it should only be called in a higher-level `select!`.
     pub(crate) async fn select(&mut self) -> ClientAction {
-        tokio::select! {
-            d = self.client_task.client_response_rx.recv() => {
-                match d {
-                    Some(c) => c.into(),
-                    None => ClientAction::ChannelClosed,
-                }
+        loop {
+            let out = match self.client_task.client_response_rx.recv().await {
+                Some(c) => c.into(),
+                None => break ClientAction::ChannelClosed,
+            };
+            // Ignore client responses if we have told the client to exit (we
+            // still propagate other ClientAction variants, e.g. TaskStopped).
+            if self.client_task.client_stop_tx.is_some()
+                || !matches!(out, ClientAction::Response(..))
+            {
+                break out;
             }
         }
     }