Use range requests to download bundles, add progress bar (#1113)

- Uses range requests to ask for chunks of the support bundle more incrementally
- Adds a progress bar to show progress, which is particularly useful for larger downloads
- Tested manually by using:

```
# In omicron:
$ cargo xtask omicron-dev run-all

# From this repo
$ cargo run -- bundle create
$ cargo run -- bundle download --id a2ae42e6-e863-415e-b211-94d5757defef output.zip
```

This downloads the bundle quite quickly - it's only KiB in size - but tuning the chunk size up or down can slow progress enough to show the progress bar.

Author: smklein

cli/docs/cli.json

@@ -593,6 +593,11 @@
           "name": "download",
           "about": "Downloads a support bundle",
           "args": [
+            {
+              "long": "chunk-size",
+              "short": "c",
+              "help": "The size, in bytes, of each range request to use when downloading bundles"
+            },
             {
               "long": "id",
               "help": "ID of the bundle"

cli/src/cmd_support_bundle.rs

@@ -15,9 +15,13 @@ use camino::Utf8PathBuf;
 use clap::Parser;
 use futures::Stream;
 use futures::StreamExt;
+use futures::TryStreamExt;
+use indicatif::ProgressBar;
+use indicatif::ProgressStyle;
 use oxide::Client;
 use oxide::ClientHiddenExt;
 use std::io;
+use std::num::NonZeroU64;
 use std::pin::Pin;
 use std::task::Context;
 use std::task::Poll;
@@ -28,6 +32,7 @@ use support_bundle_viewer::SupportBundleIndex;
 use tokio::io::AsyncRead;
 use tokio::io::AsyncWriteExt;
 use tokio::io::ReadBuf;
+use tokio::sync::watch;
 use uuid::Uuid;
 
 /// Downloads a support bundle
@@ -42,6 +47,71 @@ pub struct CmdDownload {
     /// Path where the bundle should be downloaded
     #[clap(short, long)]
     output: Utf8PathBuf,
+
+    /// The size, in bytes, of each range request to use when downloading bundles
+    #[clap(short, long, default_value_t = NonZeroU64::new(100 * (1 << 20)).unwrap())]
+    chunk_size: NonZeroU64,
+}
+
+// Downloads a portion of a support bundle using range requests.
+//
+// "range" is in bytes, and is inclusive on both sides.
+//
+// Returns: The observed content length of the stream, and the stream.
+async fn support_bundle_download_range(
+    client: &Client,
+    id: Uuid,
+    range: (u64, u64),
+) -> anyhow::Result<(
+    u64,
+    impl futures::Stream<Item = anyhow::Result<bytes::Bytes>>,
+)> {
+    let range_str = format!("bytes={}-{}", range.0, range.1);
+    let response = client
+        .support_bundle_download()
+        .bundle_id(id)
+        .range(&range_str)
+        .send()
+        .await
+        .with_context(|| format!("downloading support bundle {}", id))?;
+
+    let Some(len) = response.content_length() else {
+        bail!("No content length observed when downloading bundle");
+    };
+
+    Ok((
+        len,
+        response
+            .into_inner_stream()
+            .map(|r| r.map_err(|err| anyhow::anyhow!(err))),
+    ))
+}
+
+// Downloads all ranges of a support bundle, and combines them into a single
+// stream.
+//
+// Starts the download at "start" bytes (inclusive) and continues up to "end"
+// bytes (exclusive).
+fn support_bundle_download_ranges(
+    client: &Client,
+    id: Uuid,
+    chunk_size: NonZeroU64,
+    start: u64,
+    end: u64,
+) -> impl futures::Stream<Item = anyhow::Result<bytes::Bytes>> + use<'_> {
+    futures::stream::try_unfold(
+        (start, start + chunk_size.get() - 1),
+        move |range| async move {
+            if end <= range.0 {
+                return Ok(None);
+            }
+
+            let (actual_len, stream) = support_bundle_download_range(client, id, range).await?;
+            let next_range = (range.0 + actual_len, range.1 + actual_len);
+            Ok::<_, anyhow::Error>(Some((stream, next_range)))
+        },
+    )
+    .try_flatten()
 }
 
 #[async_trait]
@@ -51,24 +121,67 @@ impl crate::AuthenticatedCmd for CmdDownload {
             .await
             .with_context(|| format!("Cannot create output file: {}", self.output))?;
 
-        // NOTE: It might be worth adding a progress bar here?
-        let mut stream = client
-            .support_bundle_download()
+        let (progress_tx, progress_rx) = watch::channel(0);
+
+        let total_length = client
+            .support_bundle_head()
             .bundle_id(self.id)
             .send()
             .await?
-            .into_inner_stream();
+            .content_length()
+            .ok_or_else(|| anyhow::anyhow!("No content length"))?;
+
+        start_progress_bar(progress_rx, total_length, self.id)?;
 
+        let mut stream =
+            support_bundle_download_ranges(client, self.id, self.chunk_size, 0, total_length);
+        let mut stream = std::pin::pin!(stream);
         while let Some(data) = stream.next().await {
             match data {
                 Err(err) => bail!(err),
-                Ok(data) => output.write_all(&data).await?,
+                Ok(data) => {
+                    output.write_all(&data).await?;
+                    progress_tx.send_modify(|l| *l += data.len() as u64);
+                }
             }
         }
         Ok(())
     }
 }
 
+fn start_progress_bar(
+    mut progress_rx: watch::Receiver<u64>,
+    file_size: u64,
+    bundle_id: Uuid,
+) -> Result<ProgressBar> {
+    let pb = ProgressBar::new(file_size);
+    pb.set_style(ProgressStyle::default_bar().template(
+        "[{elapsed_precise}] [{wide_bar:.green}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta}",
+    )?);
+    pb.set_position(0);
+    pb.println(format!("Downloading bundle \"{bundle_id}\""));
+    let pb_updater = pb.clone();
+
+    tokio::spawn(async move {
+        loop {
+            match progress_rx.changed().await {
+                Ok(_) => {
+                    let p = *progress_rx.borrow();
+                    pb_updater.set_position(p);
+
+                    if p >= file_size {
+                        pb_updater.finish();
+                        return;
+                    }
+                }
+                Err(_) => return, // Sender has dropped.
+            }
+        }
+    });
+
+    Ok(pb)
+}
+
 /// Inspects a support bundle
 ///
 /// Support bundles may be inspected before they are downloaded (via

cli/tests/test_bundles.rs

@@ -0,0 +1,142 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+// Copyright 2025 Oxide Computer Company
+
+use assert_cmd::Command;
+use httpmock::MockServer;
+use oxide_httpmock::MockServerExt;
+
+#[test]
+fn test_bundle_download_ranges() {
+    let server = MockServer::start();
+
+    let id = uuid::Uuid::try_from("0a27f7c9-a18c-4c6b-9ba9-7e461e60f556").unwrap();
+    let output_path = tempfile::NamedTempFile::new().expect("Failed to make tempfile");
+    let chunk_size: u64 = 1024;
+
+    let mock_head = server.support_bundle_head(|when, then| {
+        when.bundle_id(&id).into_inner().any_request();
+        then.into_inner()
+            .header("accept-ranges", "bytes")
+            .header("content-length", "2048")
+            .header("content-type", "application/octet-stream")
+            .status(200);
+    });
+
+    let mock_download1 = server.support_bundle_download(|when, then| {
+        when.bundle_id(&id)
+            .into_inner()
+            .header("range", "bytes=0-1023")
+            .any_request();
+
+        then.into_inner()
+            .header("accept-ranges", "bytes")
+            .header("content-length", "1024")
+            .header("content-range", "bytes 0-1023/2048")
+            .header("content-type", "application/octet-stream")
+            .body([0; 1024]);
+    });
+    let mock_download2 = server.support_bundle_download(|when, then| {
+        when.bundle_id(&id)
+            .into_inner()
+            .header("range", "bytes=1024-2047")
+            .any_request();
+
+        then.into_inner()
+            .header("accept-ranges", "bytes")
+            .header("content-length", "1024")
+            .header("content-range", "bytes 1024-2047/2048")
+            .header("content-type", "application/octet-stream")
+            .body([1; 1024]);
+    });
+
+    Command::cargo_bin("oxide")
+        .unwrap()
+        .env("RUST_BACKTRACE", "1")
+        .env("OXIDE_HOST", server.url(""))
+        .env("OXIDE_TOKEN", "fake-token")
+        .arg("bundle")
+        .arg("download")
+        .arg("--id")
+        .arg(id.to_string())
+        .arg("--output")
+        .arg(output_path.path())
+        .arg("--chunk-size")
+        .arg(chunk_size.to_string())
+        .assert()
+        .success();
+
+    assert_eq!(
+        std::fs::read(output_path.path()).expect("Failed to read output"),
+        [[0; 1024], [1; 1024]].concat()
+    );
+
+    mock_head.assert();
+    mock_download1.assert();
+    mock_download2.assert();
+}
+
+// This server claims to support range requests, but it returns the entire file
+// anyway.
+//
+// Our CLI should cope by streaming the file as it is returned.
+#[test]
+fn test_bundle_download_range_too_large_response() {
+    let server = MockServer::start();
+
+    let id = uuid::Uuid::try_from("0a27f7c9-a18c-4c6b-9ba9-7e461e60f556").unwrap();
+    let output_path = tempfile::NamedTempFile::new().expect("Failed to make tempfile");
+    let chunk_size: u64 = 1024;
+
+    let mock_head = server.support_bundle_head(|when, then| {
+        when.bundle_id(&id).into_inner().any_request();
+        then.into_inner()
+            .header("accept-ranges", "bytes")
+            .header("content-length", "4096")
+            .header("content-type", "application/octet-stream")
+            .status(200);
+    });
+
+    // Even though the server advertised that range requests were supported,
+    // it'll return the whole file at once. Confirm that we download it
+    // in a single call anyway.
+    let mock_download = server.support_bundle_download(|when, then| {
+        when.bundle_id(&id)
+            .into_inner()
+            .header("range", "bytes=0-1023")
+            .any_request();
+
+        then.into_inner()
+            .header("accept-ranges", "bytes")
+            .header("content-length", "4096")
+            .header("content-range", "bytes 0-4095/4096")
+            .header("content-type", "application/octet-stream")
+            .body([0; 4096]);
+    });
+
+    Command::cargo_bin("oxide")
+        .unwrap()
+        .env("RUST_BACKTRACE", "1")
+        .env("OXIDE_HOST", server.url(""))
+        .env("OXIDE_TOKEN", "fake-token")
+        .arg("bundle")
+        .arg("download")
+        .arg("--id")
+        .arg(id.to_string())
+        .arg("--output")
+        .arg(output_path.path())
+        .arg("--chunk-size")
+        .arg(chunk_size.to_string())
+        .assert()
+        .success();
+
+    assert_eq!(
+        std::fs::read(output_path.path()).expect("Failed to read output"),
+        [0; 4096]
+    );
+
+    mock_head.assert();
+    mock_download.assert();
+}