From ec636b6002174ebc102c3c537f0efbc068c995df Mon Sep 17 00:00:00 2001
From: Gary Guo <gary.guo@lowrisc.org>
Date: Sun, 3 Mar 2024 23:27:49 +0000
Subject: [PATCH] Improve wait handling when docker is killed

---
 src/docker/container.rs | 23 ++++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/src/docker/container.rs b/src/docker/container.rs
index a352c49..3bd6e9c 100644
--- a/src/docker/container.rs
+++ b/src/docker/container.rs
@@ -2,7 +2,7 @@ use crate::cli::Timeout;
 
 use super::{IoStream, IoStreamSource};
 
-use anyhow::{anyhow, ensure, Context, Error, Result};
+use anyhow::{anyhow, Context, Error, Result};
 use bollard::service::EventMessage;
 use futures::future::{BoxFuture, Shared};
 use tokio::io::AsyncWriteExt;
@@ -139,16 +139,21 @@ impl Container {
         let options = bollard::container::WaitContainerOptions {
             condition: "not-running",
         };
-        let mut response = self.docker.wait_container(self.id.as_str(), Some(options));
 
-        let mut last = None;
-        while let Some(wait_response) = response.next().await {
-            last = Some(wait_response?);
+        let response = self
+            .docker
+            .wait_container(self.id.as_str(), Some(options))
+            .next()
+            .await
+            .context("No response received for wait")?;
+
+        match response {
+            Ok(response) => Ok(response.status_code),
+            // If the container does not complete, e.g. it's killed, then we will receive
+            // an error code through docker.
+            Err(bollard::errors::Error::DockerContainerWaitError { error: _, code }) => Ok(code),
+            Err(err) => Err(err)?,
         }
-
-        ensure!(last.is_some(), "Unexpected exit status");
-
-        Ok(last.unwrap().status_code)
     }
 
     pub async fn mkdir<T: AsRef<std::path::Path>>(&self, path: T) -> Result<()> {