vmm: fix stdout/stderr preemption

We introduced a preemptable receiver feature so that the new start containerd
can preempt the old stdio streaming process routine.

but it has a bug that when io copy thread finished, it will remove the
io channel, which has a preemption signal sender in it, the removal of
the sender may preempt the stream receiver to stop early, without all
message processed.

After this PR change, after copy of stdout and stderr finished, we don't
call remove of the iochannel, but just set a flag of sender_closed to
true. Before stream handle return, it just determin if the io channel
should be removed by this flag.

Signed-off-by: Abel Feng <fshb1988@gmail.com>

Author: abel-von

vmm/task/src/io.rs

@@ -53,7 +53,7 @@ use tokio_vsock::{VsockListener, VsockStream};
 
 use crate::{
     device::SYSTEM_DEV_PATH,
-    streaming::{get_output, get_stdin, remove_channel},
+    streaming::{close_stdout, get_output, get_stdin, remove_channel},
     vsock,
 };
 
@@ -281,7 +281,7 @@ where
         };
         copy(src, dst, exit_signal, on_close).await;
         if to.contains(STREAMING) {
-            remove_channel(&to).await.unwrap_or_default();
+            close_stdout(&to).await.unwrap_or_default();
         }
         debug!("finished copy io from container to {}", to);
     });

vmm/task/src/streaming.rs

@@ -80,6 +80,7 @@ pub struct IOChannel {
     remaining_data: Option<Any>,
     preemption_sender: Option<Sender<()>>,
     notifier: Arc<Notify>,
+    sender_closed: bool,
 }
 
 pub struct PreemptableReceiver {
@@ -116,6 +117,7 @@ impl IOChannel {
             remaining_data: None,
             preemption_sender: None,
             notifier: Arc::new(Notify::new()),
+            sender_closed: false,
         }
     }
 
@@ -135,10 +137,21 @@ impl IOChannel {
         None
     }
 
-    fn return_preempted_receiver(&mut self, r: PreemptableReceiver, remaining_data: Option<Any>) {
-        self.receiver = Some(r.receiver);
-        self.remaining_data = remaining_data;
-        self.notifier.notify_one();
+    fn return_preempted_receiver(
+        &mut self,
+        r: PreemptableReceiver,
+        remaining_data: Option<Any>,
+    ) -> bool {
+        // only return the receiver when sender is already moved to the io thread, and sender is not closed
+        if let None = self.sender {
+            if !self.sender_closed {
+                self.receiver = Some(r.receiver);
+                self.remaining_data = remaining_data;
+                self.notifier.notify_one();
+                return true;
+            }
+        }
+        return false;
     }
 }
 
@@ -214,11 +227,17 @@ impl Service {
         remaining_data: Option<Any>,
     ) {
         let mut ios = self.ios.lock().await;
+        let mut channel_used = false;
         if let Some(ch) = ios.get_mut(id) {
-            ch.return_preempted_receiver(r, remaining_data);
+            if ch.return_preempted_receiver(r, remaining_data) {
+                channel_used = true;
+            }
         } else {
             warn!("io channel removed when return the receiver");
         }
+        if !channel_used {
+            self.ios.lock().await.remove(id);
+        }
     }
 
     async fn get_remaining_data(&self, id: &str) -> Option<Any> {
@@ -283,18 +302,14 @@ impl Service {
                     Ok(d) => d,
                     Err(e) => {
                         debug!("failed to marshal update of stream {}, {}", stream_id, e);
-                        if let Some(c) = self.ios.lock().await.get_mut(stream_id) {
-                            c.sender = Some(sender);
-                        }
+                        self.ios.lock().await.remove(stream_id);
                         return Err(ttrpc::Error::Others(format!("failed to write data {}", e)));
                     }
                 };
                 let a = new_any!(WindowUpdate, update_bytes);
                 if let Err(e) = stream.send(&a).await {
                     debug!("failed to send update of stream {}, {}", stream_id, e);
-                    if let Some(c) = self.ios.lock().await.get_mut(stream_id) {
-                        c.sender = Some(sender);
-                    }
+                    self.ios.lock().await.remove(stream_id);
                     return Err(e);
                 }
                 window += WINDOW_SIZE;
@@ -310,6 +325,7 @@ impl Service {
                     };
                     let len: i32 = data_bytes.len().try_into().unwrap_or_default();
                     if let Err(e) = sender.send(data_bytes).await {
+                        self.ios.lock().await.remove(stream_id);
                         return Err(ttrpc::Error::Others(format!("failed to send data {}", e)));
                     }
                     window -= len;
@@ -327,6 +343,7 @@ impl Service {
         stream_id: &String,
         stream: ServerStream<Any, Any>,
     ) -> ttrpc::Result<()> {
+        // TODO the stream needs to return if client send close.
         let mut receiver = self.preempt_receiver(stream_id).await?;
         if let Some(a) = self.get_remaining_data(stream_id).await {
             if let Err(e) = stream.send(&a).await {
@@ -348,6 +365,7 @@ impl Service {
             match r {
                 Some(d) => {
                     if d.is_empty() {
+                        self.ios.lock().await.remove(stream_id);
                         return Ok(());
                     }
                     let mut data = Data::new();
@@ -376,6 +394,7 @@ impl Service {
                     };
                 }
                 None => {
+                    self.ios.lock().await.remove(stream_id);
                     return Ok(());
                 }
             }
@@ -394,6 +413,14 @@ pub async fn remove_channel(url: &str) -> containerd_shim::Result<()> {
     Ok(())
 }
 
+pub async fn close_stdout(url: &str) -> containerd_shim::Result<()> {
+    let id = get_id(url)?;
+    if let Some(ch) = STREAMING_SERVICE.ios.lock().await.get_mut(id) {
+        ch.sender_closed = true;
+    }
+    Ok(())
+}
+
 pub async fn get_output(url: &str) -> containerd_shim::Result<StreamingOutput> {
     let id = get_id(url)?;
     STREAMING_SERVICE.get_output(id).await