Fix indefinite hangs on non-EC2 network paths

coldsnap uploads can hang indefinitely when running from outside AWS.
A small fraction of PutSnapshotBlock requests complete the TCP handshake
but never receive an HTTP response. The AWS Rust SDK does not set
read_timeout, operation_attempt_timeout, or operation_timeout by
default, so these requests block the worker forever.

This has been reported in #362 (uploads stalling from GitHub Actions),
#374 (downloads hanging from non-EC2), #216 (excessive retries with no
visibility), and #95 (missing timeouts on remote calls).

Changes:

- Set SDK timeouts in build_client_config(): read_timeout 12s,
  operation_attempt_timeout 20s, operation_timeout 120s. These apply to
  both the upload and download CLI paths.

- Set SDK max_attempts to 1. coldsnap has its own per-block retry loop
  with backoff; layering SDK retries on top produced up to 36 attempts
  per block with no coordinated timeout.

- Reduce block retry count from 12 to 5. With bounded per-attempt
  timeouts, fewer retries are needed.

- Add --workers flag (default 64) to configure concurrent upload
  workers. Reject --workers 0 since for_each_concurrent treats 0 as
  unlimited.

- Add --client-shards flag (default 1) to create N independent
  EbsClient instances for uploads. Blocks are distributed by index
  (block_index % N). This is opt-in; the default preserves the
  existing single-client behavior. In our testing, --client-shards 8
  with 64 workers improved the per-block latency profile on
  high-latency paths.

- Log a latency histogram at INFO level after the upload completes
  (<250ms, 250-500ms, 500ms-1s, 1-2s, 2-5s, >5s, plus error count).

- Log per-block warnings on failure with block index, attempt number,
  elapsed time, and error. Previously failures were logged at DEBUG
  with no context.

Tested with 18 consecutive uploads of a 4.1 GiB image from GitHub
Actions runners across Virginia, Wyoming, and other Azure regions.
All 18 succeeded (20-51s depending on config and runner location).
Stock coldsnap on the same paths hung indefinitely or took 5-20+
minutes.

Author: nkk0

src/bin/coldsnap/main.rs

@@ -9,6 +9,8 @@ snapshots.
 use argh::FromArgs;
 use aws_config::default_provider::credentials::DefaultCredentialsChain;
 use aws_config::default_provider::region::DefaultRegionChain;
+use aws_config::retry::RetryConfig;
+use aws_config::timeout::TimeoutConfig;
 use aws_sdk_ebs::types::Tag;
 use aws_sdk_ebs::Client as EbsClient;
 use aws_sdk_ec2::Client as Ec2Client;
@@ -85,8 +87,25 @@ async fn run() -> Result<()> {
         }
 
         SubCommand::Upload(upload_args) => {
-            let client = EbsClient::new(&client_config);
-            let uploader = SnapshotUploader::new(client);
+            if upload_args.workers == Some(0) {
+                eprintln!("Error: --workers must be greater than zero");
+                std::process::exit(1);
+            }
+            if upload_args.client_shards == Some(0) {
+                eprintln!("Error: --client-shards must be greater than zero");
+                std::process::exit(1);
+            }
+
+            let num_shards = upload_args.client_shards.unwrap_or(1);
+            let uploader = if num_shards <= 1 {
+                SnapshotUploader::new(EbsClient::new(&client_config))
+            } else {
+                debug!("Creating {} EBS client shards", num_shards);
+                let clients = (0..num_shards)
+                    .map(|_| EbsClient::new(&client_config))
+                    .collect();
+                SnapshotUploader::with_client_shards(clients)
+            };
             ensure!(
                 upload_args.file.file_name().is_some(),
                 error::ValidateFilenameSnafu {
@@ -115,6 +134,7 @@ async fn run() -> Result<()> {
                     progress_bar?,
                     zero_blocks,
                     upload_args.kms_key_id,
+                    upload_args.workers,
                 )
                 .await
                 .context(error::UploadSnapshotSnafu)?;
@@ -222,6 +242,22 @@ async fn build_client_config(
         config = config.endpoint_url(endpoint);
     }
 
+    // The AWS SDK does not set response or per-attempt timeouts by default.
+    // Without these, a request that sends its body but never receives a response
+    // will block the worker indefinitely.
+    config = config
+        .timeout_config(
+            TimeoutConfig::builder()
+                .read_timeout(Duration::from_secs(12))
+                .operation_attempt_timeout(Duration::from_secs(20))
+                .operation_timeout(Duration::from_secs(120))
+                .build(),
+        )
+        // Disable SDK-level retries; coldsnap already has its own per-block retry
+        // loop with backoff.  Layering SDK retries on top of that leads to excessive
+        // total attempts and unpredictable wall-clock time.
+        .retry_config(RetryConfig::standard().with_max_attempts(1));
+
     config.load().await
 }
 
@@ -396,6 +432,14 @@ struct UploadArgs {
     #[argh(switch)]
     /// omit blocks of all zeros when uploading
     omit_zero_blocks: bool,
+
+    #[argh(option)]
+    /// number of concurrent upload workers (default: 64)
+    workers: Option<usize>,
+
+    #[argh(option)]
+    /// number of independent EBS clients for higher-concurrency uploads (default: 1)
+    client_shards: Option<usize>,
 }
 
 /// Turn a user-specified duration in seconds into a Duration object, for argh parsing.

src/lib.rs

@@ -34,7 +34,7 @@ let client = EbsClient::new(&aws_config::from_env().region("us-west-2").load().a
 let uploader = SnapshotUploader::new(client);
 let path = Path::new("./disk.img");
 
-let snapshot_id = uploader.upload_from_file(&path, None, None, None, None, None, None)
+let snapshot_id = uploader.upload_from_file(&path, None, None, None, None, None, None, None)
         .await
         .expect("failed to upload snapshot");
 # }

src/upload.rs

@@ -14,7 +14,7 @@ use base64::Engine as _;
 use bytes::BytesMut;
 use futures::stream::{self, StreamExt};
 use indicatif::ProgressBar;
-use log::debug;
+use log::{debug, info, warn};
 use sha2::{Digest, Sha256};
 use snafu::{ensure, OptionExt, ResultExt, Snafu};
 use std::cmp;
@@ -24,7 +24,7 @@ use std::ffi::OsStr;
 use std::io::SeekFrom;
 use std::os::unix::fs::FileTypeExt;
 use std::path::{Path, PathBuf};
-use std::sync::atomic::{AtomicI32, Ordering as AtomicOrdering};
+use std::sync::atomic::{AtomicI32, AtomicU64, Ordering as AtomicOrdering};
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 use tokio::fs::{self, File};
@@ -39,14 +39,58 @@ const GIBIBYTE: i64 = 1024 * 1024 * 1024;
 const SNAPSHOT_BLOCK_WORKERS: usize = 64;
 // How long to wait between attempts; this number * attempt number, in seconds.
 const SNAPSHOT_BLOCK_RETRY_SCALE: u64 = 2;
-// 12 retries with scale 2 gives us 132 seconds, chosen because it got past "snapshot does not
-// exist" errors in testing.
-const SNAPSHOT_BLOCK_ATTEMPTS: u64 = 12;
+// 5 retries with scale 2 gives us 20 seconds of backoff. With SDK-level timeouts
+// now in place, each attempt is bounded, so fewer retries are needed.
+const SNAPSHOT_BLOCK_ATTEMPTS: u64 = 5;
 const SNAPSHOT_TIMEOUT_MINUTES: i32 = 10;
 const SHA256_ALGORITHM: ChecksumAlgorithm = ChecksumAlgorithm::ChecksumAlgorithmSha256;
 const LINEAR_METHOD: ChecksumAggregationMethod =
     ChecksumAggregationMethod::ChecksumAggregationLinear;
 
+/// Collects per-block upload latencies for a summary logged after the upload.
+struct UploadStats {
+    buckets: [AtomicU64; 6],
+    errors: AtomicU64,
+}
+
+impl UploadStats {
+    fn new() -> Self {
+        Self {
+            buckets: Default::default(),
+            errors: AtomicU64::new(0),
+        }
+    }
+
+    fn record_success(&self, elapsed: Duration) {
+        let bucket = match elapsed.as_millis() {
+            0..250 => 0,
+            250..500 => 1,
+            500..1000 => 2,
+            1000..2000 => 3,
+            2000..5000 => 4,
+            _ => 5,
+        };
+        self.buckets[bucket].fetch_add(1, AtomicOrdering::Relaxed);
+    }
+
+    fn record_error(&self) {
+        self.errors.fetch_add(1, AtomicOrdering::Relaxed);
+    }
+
+    fn report(&self) {
+        let b: Vec<u64> = self
+            .buckets
+            .iter()
+            .map(|a| a.load(AtomicOrdering::Relaxed))
+            .collect();
+        let e = self.errors.load(AtomicOrdering::Relaxed);
+        info!(
+            "Upload complete: <250ms={} 250-500ms={} 500ms-1s={} 1-2s={} 2-5s={} >5s={} errors={}",
+            b[0], b[1], b[2], b[3], b[4], b[5], e
+        );
+    }
+}
+
 /// Specify how blocks of all zeroes should be handled.
 #[derive(Copy, Clone)]
 pub enum ZeroBlocks {
@@ -59,12 +103,27 @@ pub enum ZeroBlocks {
 }
 
 pub struct SnapshotUploader {
-    ebs_client: EbsClient,
+    ebs_clients: Vec<EbsClient>,
 }
 
 impl SnapshotUploader {
     pub fn new(ebs_client: EbsClient) -> Self {
-        SnapshotUploader { ebs_client }
+        SnapshotUploader {
+            ebs_clients: vec![ebs_client],
+        }
+    }
+
+    /// Create an uploader with multiple independent EBS clients.  Blocks are
+    /// distributed across clients by index, giving each a separate HTTP
+    /// connection pool.  This can reduce head-of-line blocking when many
+    /// workers share a single pool over high-latency paths.
+    pub fn with_client_shards(ebs_clients: Vec<EbsClient>) -> Self {
+        assert!(!ebs_clients.is_empty(), "need at least one EBS client");
+        SnapshotUploader { ebs_clients }
+    }
+
+    fn client_for_block(&self, block_index: i32) -> &EbsClient {
+        &self.ebs_clients[block_index as usize % self.ebs_clients.len()]
     }
 
     /// Upload a snapshot from the file at the specified path.
@@ -89,6 +148,7 @@ impl SnapshotUploader {
         progress_bar: Option<ProgressBar>,
         zero_blocks: Option<ZeroBlocks>,
         kms_key_id: Option<String>,
+        workers: Option<usize>,
     ) -> Result<String> {
         let path = path.as_ref();
         let description = description.map(|s| s.to_string()).unwrap_or_else(|| {
@@ -187,7 +247,7 @@ impl SnapshotUploader {
                 block_digests: Arc::clone(&block_digests),
                 block_errors: Arc::clone(&block_errors),
                 progress_bar: Arc::clone(&progress_bar),
-                ebs_client: self.ebs_client.clone(),
+                ebs_client: self.client_for_block(i).clone(),
                 zero_blocks,
             });
 
@@ -197,30 +257,56 @@ impl SnapshotUploader {
         // Distribute the work across a fixed number of concurrent workers.
         // New threads will be created by the runtime as needed, but we'll
         // only process this many blocks at once to limit resource usage.
-        let upload = stream::iter(block_contexts).for_each_concurrent(
-            SNAPSHOT_BLOCK_WORKERS,
-            |context| async move {
+        let worker_count = workers.unwrap_or(SNAPSHOT_BLOCK_WORKERS);
+        assert!(worker_count > 0, "--workers must be greater than zero");
+        debug!(
+            "Using {} concurrent upload workers across {} client shards",
+            worker_count,
+            self.ebs_clients.len()
+        );
+        let stats = Arc::new(UploadStats::new());
+        let upload = stream::iter(block_contexts).for_each_concurrent(worker_count, |context| {
+            let stats = Arc::clone(&stats);
+            async move {
                 for attempt in 0..SNAPSHOT_BLOCK_ATTEMPTS {
-                    // Increasing wait between attempts.  (No wait to start, on 0th attempt.)
-                    time::sleep(Duration::from_secs(attempt * SNAPSHOT_BLOCK_RETRY_SCALE)).await;
+                    if attempt > 0 {
+                        let backoff = Duration::from_secs(attempt * SNAPSHOT_BLOCK_RETRY_SCALE);
+                        debug!(
+                            "block {}: retry {}/{}, backoff {}s",
+                            context.block_index,
+                            attempt,
+                            SNAPSHOT_BLOCK_ATTEMPTS,
+                            backoff.as_secs()
+                        );
+                        time::sleep(backoff).await;
+                    }
 
+                    let start = std::time::Instant::now();
                     let block_result = self.upload_block(&context).await;
+                    let elapsed = start.elapsed();
+
                     let mut block_errors = context.block_errors.lock().expect("poisoned");
                     if let Err(e) = block_result {
-                        debug!(
-                            "Error uploading block, attempt {} of {}",
+                        stats.record_error();
+                        warn!(
+                            "block {}: attempt {}/{} failed after {:.1}s: {}",
+                            context.block_index,
                             attempt + 1,
-                            SNAPSHOT_BLOCK_ATTEMPTS
+                            SNAPSHOT_BLOCK_ATTEMPTS,
+                            elapsed.as_secs_f64(),
+                            e
                         );
                         block_errors.insert(context.block_index, e);
                         continue;
                     }
+                    stats.record_success(elapsed);
                     block_errors.remove(&context.block_index);
                     break;
                 }
-            },
-        );
+            }
+        });
         upload.await;
+        stats.report();
 
         // At this point, all the concurrent jobs have finished, so all of the Arcs we copied have
         // been dropped. Hence there's exactly one strong reference and it's safe to `try_unwrap`
@@ -283,8 +369,7 @@ impl SnapshotUploader {
         tags: Option<Vec<Tag>>,
         kms_key_id: Option<String>,
     ) -> Result<(String, i32)> {
-        let mut request = self
-            .ebs_client
+        let mut request = self.ebs_clients[0]
             .start_snapshot()
             .volume_size(volume_size)
             .set_description(Some(description))
@@ -316,7 +401,7 @@ impl SnapshotUploader {
         changed_blocks_count: i32,
         checksum: &str,
     ) -> Result<()> {
-        self.ebs_client
+        self.ebs_clients[0]
             .complete_snapshot()
             .snapshot_id(snapshot_id)
             .changed_blocks_count(changed_blocks_count)
@@ -579,3 +664,72 @@ mod error {
         },
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn histogram_bucket_boundaries() {
+        let stats = UploadStats::new();
+
+        stats.record_success(Duration::from_millis(0));
+        stats.record_success(Duration::from_millis(249));
+        stats.record_success(Duration::from_millis(250));
+        stats.record_success(Duration::from_millis(499));
+        stats.record_success(Duration::from_millis(500));
+        stats.record_success(Duration::from_millis(999));
+        stats.record_success(Duration::from_millis(1000));
+        stats.record_success(Duration::from_millis(1999));
+        stats.record_success(Duration::from_millis(2000));
+        stats.record_success(Duration::from_millis(4999));
+        stats.record_success(Duration::from_millis(5000));
+        stats.record_success(Duration::from_millis(60000));
+
+        let b: Vec<u64> = stats
+            .buckets
+            .iter()
+            .map(|a| a.load(AtomicOrdering::Relaxed))
+            .collect();
+
+        assert_eq!(b[0], 2); // <250ms: 0, 249
+        assert_eq!(b[1], 2); // 250-500ms: 250, 499
+        assert_eq!(b[2], 2); // 500ms-1s: 500, 999
+        assert_eq!(b[3], 2); // 1-2s: 1000, 1999
+        assert_eq!(b[4], 2); // 2-5s: 2000, 4999
+        assert_eq!(b[5], 2); // >5s: 5000, 60000
+    }
+
+    #[test]
+    fn error_counter() {
+        let stats = UploadStats::new();
+        stats.record_error();
+        stats.record_error();
+        stats.record_error();
+        assert_eq!(stats.errors.load(AtomicOrdering::Relaxed), 3);
+    }
+
+    #[test]
+    fn client_for_block_modulo_logic() {
+        // Verify the shard selection formula: block_index % num_shards.
+        let num_shards = 3usize;
+        let expected = [0, 1, 2, 0, 1, 2, 0, 1, 2];
+        for (i, &want) in expected.iter().enumerate() {
+            assert_eq!(i % num_shards, want);
+        }
+    }
+
+    #[test]
+    #[should_panic(expected = "need at least one EBS client")]
+    fn with_client_shards_rejects_empty() {
+        SnapshotUploader::with_client_shards(vec![]);
+    }
+
+    #[test]
+    #[should_panic(expected = "--workers must be greater than zero")]
+    fn worker_count_zero_panics() {
+        // Simulates what happens if workers=Some(0) gets past CLI validation.
+        let count: usize = 0;
+        assert!(count > 0, "--workers must be greater than zero");
+    }
+}