post3/crates/post3-sdk/examples/multipart_upload.rs

//! Stress test: multipart upload and verify huge files (4–16 GiB).
//!
//! Uses the SDK's multipart_upload convenience method which splits data into
//! parts and uploads them sequentially via CreateMultipartUpload / UploadPart /
//! CompleteMultipartUpload.
//!
//! Prerequisites: post3-server running on localhost:9000
//!   mise run up && mise run dev
//!
//! Run:
//!   cargo run -p post3-sdk --example multipart_upload --release
//!
//! Or with custom sizes (in MB) and part size:
//!   POST3_SIZES=4096,8192,16384 POST3_PART_SIZE=64 cargo run -p post3-sdk --example multipart_upload --release

use post3_sdk::Post3Client;
use std::time::Instant;

fn generate_data(size_bytes: usize) -> Vec<u8> {
    let mut data = Vec::with_capacity(size_bytes);
    let mut state: u64 = 0xdeadbeef;
    while data.len() < size_bytes {
        state ^= state << 13;
        state ^= state >> 7;
        state ^= state << 17;
        data.extend_from_slice(&state.to_le_bytes());
    }
    data.truncate(size_bytes);
    data
}

fn format_size(bytes: usize) -> String {
    if bytes >= 1024 * 1024 * 1024 {
        format!("{:.1} GiB", bytes as f64 / (1024.0 * 1024.0 * 1024.0))
    } else if bytes >= 1024 * 1024 {
        format!("{:.1} MiB", bytes as f64 / (1024.0 * 1024.0))
    } else if bytes >= 1024 {
        format!("{:.1} KiB", bytes as f64 / 1024.0)
    } else {
        format!("{} B", bytes)
    }
}

fn format_throughput(bytes: usize, duration: std::time::Duration) -> String {
    let secs = duration.as_secs_f64();
    if secs == 0.0 {
        return "∞".to_string();
    }
    let mb_per_sec = bytes as f64 / (1024.0 * 1024.0) / secs;
    format!("{:.1} MiB/s", mb_per_sec)
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let endpoint = std::env::var("POST3_ENDPOINT")
        .unwrap_or_else(|_| "http://localhost:9000".to_string());
    let client = Post3Client::new(&endpoint);

    let sizes_mb: Vec<usize> = std::env::var("POST3_SIZES")
        .unwrap_or_else(|_| "100,1024,4096,8192,16384".to_string())
        .split(',')
        .filter_map(|s| s.trim().parse().ok())
        .collect();

    // Part size in MiB (default 64 MiB — good balance of part count vs memory)
    let part_size_mb: usize = std::env::var("POST3_PART_SIZE")
        .unwrap_or_else(|_| "64".to_string())
        .parse()
        .unwrap_or(64);
    let part_size = part_size_mb * 1024 * 1024;

    println!("=== post3 Multipart Upload Stress Test ===");
    println!("Endpoint: {}", endpoint);
    println!("Sizes: {:?} MB", sizes_mb);
    println!("Part size: {} MiB", part_size_mb);
    println!();

    client.create_bucket("mp-stress").await?;

    for size_mb in &sizes_mb {
        let size_bytes = size_mb * 1024 * 1024;
        let key = format!("mp-test-{}mb.bin", size_mb);
        let num_parts = (size_bytes + part_size - 1) / part_size;

        println!("--- {} ({} parts of {} each) ---",
            format_size(size_bytes),
            num_parts,
            format_size(part_size.min(size_bytes)),
        );

        // Generate data
        print!("  Generating data... ");
        let gen_start = Instant::now();
        let data = generate_data(size_bytes);
        println!("done ({:.1}s)", gen_start.elapsed().as_secs_f64());

        // Multipart upload
        print!("  Uploading (multipart)... ");
        let upload_start = Instant::now();
        match client.multipart_upload("mp-stress", &key, &data, part_size).await {
            Ok(()) => {
                let upload_dur = upload_start.elapsed();
                println!(
                    "done ({:.1}s, {})",
                    upload_dur.as_secs_f64(),
                    format_throughput(size_bytes, upload_dur)
                );
            }
            Err(e) => {
                println!("FAILED: {}", e);
                println!("  Skipping remaining sizes");
                break;
            }
        }

        // Head (verify metadata)
        let head = client.head_object("mp-stress", &key).await?;
        if let Some(info) = &head {
            println!(
                "  Head:              size={}, etag={:?}",
                format_size(info.size as usize),
                info.etag
            );
            // Verify the compound ETag format (md5-N)
            if let Some(ref etag) = info.etag {
                let stripped = etag.trim_matches('"');
                if stripped.contains('-') {
                    let parts_str = stripped.split('-').last().unwrap_or("?");
                    println!("  ETag format:       compound ({} parts)", parts_str);
                }
            }
        }

        // Download and verify
        print!("  Downloading...     ");
        let download_start = Instant::now();
        match client.get_object("mp-stress", &key).await {
            Ok(downloaded) => {
                let download_dur = download_start.elapsed();
                println!(
                    "done ({:.1}s, {})",
                    download_dur.as_secs_f64(),
                    format_throughput(size_bytes, download_dur)
                );

                print!("  Verifying...       ");
                if downloaded.len() != data.len() {
                    println!(
                        "FAILED: size mismatch (expected {}, got {})",
                        data.len(),
                        downloaded.len()
                    );
                } else if downloaded.as_ref() == data.as_slice() {
                    println!("OK (byte-for-byte match)");
                } else {
                    let pos = data
                        .iter()
                        .zip(downloaded.iter())
                        .position(|(a, b)| a != b)
                        .unwrap_or(0);
                    println!("FAILED: mismatch at byte {}", pos);
                }
            }
            Err(e) => {
                println!("FAILED: {}", e);
            }
        }

        // Cleanup
        client.delete_object("mp-stress", &key).await?;
        println!();
    }

    client.delete_bucket("mp-stress").await?;
    println!("=== Done ===");
    Ok(())
}