dist: make installation asynchronous

Author: djc

src/diskio/mod.rs

@@ -290,7 +290,7 @@ impl IncrementalFileState {
 /// Trait object for performing IO. At this point the overhead
 /// of trait invocation is not a bottleneck, but if it becomes
 /// one we could consider an enum variant based approach instead.
-pub(crate) trait Executor {
+pub(crate) trait Executor: Send {
     /// Perform a single operation.
     /// During overload situations previously queued items may
     /// need to be completed before the item is accepted:

src/dist/manifestation.rs

@@ -4,11 +4,17 @@
 #[cfg(test)]
 mod tests;
 
+use std::collections::VecDeque;
 use std::path::Path;
+use std::pin::Pin;
 use std::sync::Arc;
+use std::task::{Context, Poll};
+use std::vec;
 
-use anyhow::{Context, Result, anyhow, bail};
+use anyhow::{Context as _, Result, anyhow, bail};
+use futures_util::Stream;
 use futures_util::stream::{FuturesUnordered, StreamExt};
+use tokio::task::{JoinHandle, spawn_blocking};
 use tracing::{info, warn};
 
 use crate::diskio::{Executor, IO_CHUNK_SIZE, get_executor, unpack_ram};
@@ -211,34 +217,44 @@ impl Manifestation {
         }
 
         info!("downloading component(s)");
-        let mut downloads = FuturesUnordered::new();
-        let mut component_iter = components.into_iter();
-        let mut cleanup_downloads = vec![];
-        let manifestation = Arc::new(self);
-        loop {
-            if downloads.is_empty() && component_iter.len() == 0 {
-                break;
-            }
+        let mut tx = if !components.is_empty() {
+            let mut stream = InstallEvents::new(components.into_iter(), Arc::new(self));
+            let mut transaction = Some(tx);
+            let tx = loop {
+                // Refill downloads when there's capacity
+                // Must live outside of `InstallEvents` because we can't write the type of future
+                while stream.components.len() > 0 && stream.downloads.len() < concurrent_downloads {
+                    if let Some(bin) = stream.components.next() {
+                        stream.downloads.push(bin.download(max_retries));
+                    }
+                }
 
-            let installable = downloads.next().await.transpose()?;
-            while component_iter.len() > 0 && downloads.len() < concurrent_downloads {
-                if let Some(bin) = component_iter.next() {
-                    downloads.push(bin.download(max_retries));
+                // Trigger another installation if no other installation is in progress, as evidenced
+                // by whether `transaction` is `Some` (not held by another installation task).
+                stream.try_install(&mut transaction);
+                match stream.next().await {
+                    // Completed an installation, yielding the transaction back
+                    Some(Ok(tx)) => match stream.is_done() {
+                        true => break tx,
+                        false => transaction = Some(tx),
+                    },
+                    Some(Err(e)) => return Err(e),
+                    // A download completed, so we can trigger another one
+                    None => {}
                 }
-            }
+            };
 
-            if let Some((installable, hash)) = installable {
-                cleanup_downloads.push(hash);
-                tx = installable.install(tx, manifestation.clone())?;
-            }
-        }
+            download_cfg.clean(&stream.cleanup_downloads)?;
+            drop(stream);
+            tx
+        } else {
+            tx
+        };
 
         // Install new distribution manifest
         let new_manifest_str = new_manifest.clone().stringify()?;
         tx.modify_file(rel_installed_manifest_path)?;
         utils::write_file("manifest", &installed_manifest_path, &new_manifest_str)?;
-        download_cfg.clean(&cleanup_downloads)?;
-        drop(downloads);
 
         // Write configuration.
         //
@@ -445,6 +461,94 @@ impl Manifestation {
     }
 }
 
+struct InstallEvents<'a, F> {
+    manifestation: Arc<Manifestation>,
+    components: vec::IntoIter<ComponentBinary<'a>>,
+    cleanup_downloads: Vec<&'a str>,
+    install_queue: VecDeque<ComponentInstall>,
+    installing: Option<JoinHandle<Result<Transaction>>>,
+    downloads: FuturesUnordered<F>,
+}
+
+impl<'a, F> InstallEvents<'a, F> {
+    fn new(
+        components: vec::IntoIter<ComponentBinary<'a>>,
+        manifestation: Arc<Manifestation>,
+    ) -> Self {
+        Self {
+            manifestation,
+            cleanup_downloads: Vec::with_capacity(components.len()),
+            components,
+            install_queue: VecDeque::new(),
+            installing: None,
+            downloads: FuturesUnordered::new(),
+        }
+    }
+
+    fn try_install(&mut self, tx: &mut Option<Transaction>) {
+        let Some(installable) = self.install_queue.pop_front() else {
+            return;
+        };
+
+        if let Some(tx) = tx.take() {
+            let manifestation = self.manifestation.clone();
+            self.installing = Some(spawn_blocking(|| installable.install(tx, manifestation)));
+        } else {
+            self.install_queue.push_front(installable);
+        }
+    }
+
+    fn is_done(&self) -> bool {
+        self.components.len() == 0 && self.downloads.is_empty() && self.install_queue.is_empty()
+    }
+}
+
+impl<'a, F: Future<Output = Result<(ComponentInstall, &'a str)>>> Stream for InstallEvents<'a, F> {
+    type Item = Result<Transaction>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        loop {
+            let mut this = self.as_mut();
+
+            // First, see if any of the downloads is complete; if so, yield `None`
+            // to the caller so it can trigger another download.
+            match Pin::new(&mut this.downloads).poll_next(cx) {
+                Poll::Ready(Some(Ok((installable, hash)))) => {
+                    this.cleanup_downloads.push(hash);
+                    this.install_queue.push_back(installable);
+                    return Poll::Ready(None);
+                }
+                Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))),
+                Poll::Ready(None) | Poll::Pending => {}
+            }
+
+            let Some(handle) = &mut this.installing else {
+                return match self.install_queue.is_empty() {
+                    // Nothing to do, yield control to the runtime
+                    true => Poll::Pending,
+                    // Can try to start the next installation
+                    false => Poll::Ready(None),
+                };
+            };
+
+            match Pin::new(handle).poll(cx) {
+                Poll::Ready(Ok(Ok(tx))) => {
+                    // Current `handle` must not be polled again
+                    this.installing = None;
+                    return Poll::Ready(Some(Ok(tx)));
+                }
+                Poll::Ready(Ok(Err(e))) => return Poll::Ready(Some(Err(e))),
+                Poll::Ready(Err(e)) => {
+                    return Poll::Ready(Some(Err(anyhow!(
+                        "internal error during installation: {e}"
+                    ))));
+                }
+                Poll::Pending => return Poll::Pending,
+            }
+        }
+    }
+}
+
 #[derive(Debug, Default)]
 struct Update {
     components_to_uninstall: Vec<Component>,