feat: add disk space check for downloads on unix

Author: motorailgun

src/dist/download.rs

@@ -15,7 +15,9 @@ use url::Url;
 use crate::config::Cfg;
 use crate::dist::manifest::Manifest;
 use crate::dist::{Channel, DEFAULT_DIST_SERVER, ToolchainDesc, temp};
-use crate::download::{download_file, download_file_with_resume, is_network_failure};
+use crate::download::{
+    content_length, download_file, download_file_with_resume, is_network_failure,
+};
 use crate::errors::RustupError;
 use crate::process::Process;
 use crate::utils;
@@ -129,6 +131,10 @@ impl<'a> DownloadCfg<'a> {
         }
     }
 
+    pub(crate) async fn content_length(&self, url: &Url) -> Result<Option<u64>> {
+        content_length(url, self.process).await
+    }
+
     pub(crate) fn clean(&self, hashes: &[impl AsRef<Path>]) -> Result<()> {
         for hash in hashes.iter() {
             let used_file = self.download_dir.join(hash);

src/dist/manifestation.rs

@@ -240,6 +240,50 @@ impl Manifestation {
                 );
             };
 
+            // TODO: enable this for windows
+            //       disk_free() is only available for cfg(unix) for now
+            #[cfg(unix)]
+            {
+                let estimated_download_size = futures_util::future::join_all(
+                    components
+                        .iter()
+                        .map(|component| component.required_size_via_head()),
+                )
+                .await
+                .iter()
+                .try_fold(0u64, |acc, elem| {
+                    match elem {
+                        Ok(Some(size)) => Some(acc + size),
+                        // This component size's unknown (None or Err)
+                        _ => {
+                            warn!("failed to fetch component size, continueing...");
+                            None
+                        }
+                    }
+                });
+
+                let path = &components
+                    .first()
+                    .unwrap()
+                    .download_cfg
+                    .tmp_cx
+                    .root_directory;
+                let disk_free = match utils::disk_free(path) {
+                    Ok(size) => Some(size),
+                    Err(e) => {
+                        warn!("failed to acquire disk free space: {e}\ncontinueing... ");
+                        None
+                    }
+                };
+
+                if let (Some(df), Some(est)) = (disk_free, estimated_download_size)
+                    && df < est
+                    && !force_update
+                {
+                    bail!("insufficient storage: {est} bytes required, {df} bytes available")
+                }
+            };
+
             let mut stream = InstallEvents::new(components.into_iter(), Arc::new(self));
             let mut transaction = Some(tx);
             tx = loop {
@@ -767,6 +811,11 @@ impl<'a> ComponentBinary<'a> {
         }))
     }
 
+    async fn required_size_via_head(&self) -> Result<Option<u64>> {
+        let url = self.download_cfg.url(&self.binary.url)?;
+        self.download_cfg.content_length(&url).await
+    }
+
     async fn download(self, max_retries: usize) -> Result<(ComponentInstall, &'a str)> {
         use tokio_retry::{RetryIf, strategy::FixedInterval};
 

src/dist/temp.rs

@@ -73,7 +73,7 @@ impl Drop for File {
 }
 
 pub struct Context {
-    root_directory: PathBuf,
+    pub root_directory: PathBuf,
     pub dist_server: String,
 }