hyperfile 0.11.3

//! How long a flush takes when data has been streaming out while it was written.
//!
//! A producer writes `BENCH_MIB` MiB in 1 MiB writes at `BENCH_RATES` MB/s (0 is as
//! fast as it can), then flushes. Three containers:
//!
//! - `parted, all at flush`: `PartedSegment`, everything goes out at the flush;
//! - `stream N MiB`: `StreamSegment`, data sent while it is written with
//!   `stream_segment_bytes` N MiB, for each N in `BENCH_SIZES_MIB`.
//!
//! Prints the write phase, the flush, and their sum, and what was sent ahead and
//! of it died; checks the data cold. With `BENCH_HOT_MIB`, the writes go round a
//! working set of that many MiB, written over and over. `BENCH_LIMITS=default`
//! keeps the default dirty-data limits; `BENCH_REACTOR=1` goes through the
//! reactor's handler.
//!
//! ```bash
//! BENCH_MIB=512 BENCH_RATES=0,400,200,100 BENCH_SIZES_MIB=4,16,256,512,1024 \
//!   cargo test --release --test bench_stream_flush -- --ignored --nocapture --test-threads=1
//! ```

mod common;
use common::*;

use std::time::{Duration, Instant};

use hyperfile::config::{HyperFileConfigBuilder, HyperFileMetaConfig, HyperFileRuntimeConfig};
use hyperfile::file::flags::{FileFlags, HyperFileFlags};
use hyperfile::file::hyper::Hyper;
use hyperfile::file::mode::{FileMode, HyperFileMode};
use hyperfile::meta_format::BlockPtrFormat;
use hyperfile::staging::config::StagingConfig;

const MIB: usize = 1 << 20;

fn env_num(name: &str, default: usize) -> usize {
    std::env::var(name).ok().and_then(|v| v.parse().ok()).unwrap_or(default)
}

fn chunk(k: usize) -> Vec<u8> {
    (0..MIB).map(|i| (k as u8).wrapping_mul(31) ^ (i as u8) ^ ((i >> 12) as u8)).collect()
}

/// The sizes of a container's segment objects, in MiB, sorted.
async fn objects(client: &aws_sdk_s3::Client, uri: &str) -> Vec<f64> {
    let (bucket, prefix) = uri.trim_start_matches("s3://").split_once('/').unwrap();
    let mut v = Vec::new();
    let mut pages = client.list_objects_v2().bucket(bucket).prefix(format!("{}/", prefix)).into_paginator().send();
    while let Some(page) = pages.next().await {
        for o in page.expect("list").contents().iter().filter(|o| !o.key().unwrap_or("").ends_with("/inode")) {
            v.push(o.size().unwrap_or(0) as f64 / MIB as f64);
        }
    }
    v.sort_by(|a, b| a.partial_cmp(b).unwrap());
    v
}

/// The handle being benchmarked: direct, or behind the reactor's handler.
enum Handle {
    Direct(Hyper<'static>),
    #[cfg(feature = "reactor")]
    Reactor(hyperfile::file::fh::HyperFileHandler<'static>),
}

impl Handle {
    async fn write(&mut self, off: usize, buf: &[u8]) {
        match self {
            Self::Direct(h) => { h.fs_write(off, buf).await.expect("write"); },
            #[cfg(feature = "reactor")]
            Self::Reactor(fh) => { fh.fh_write(off, buf).await.expect("write"); },
        }
    }
    async fn flush(&mut self) {
        match self {
            Self::Direct(h) => { h.fs_flush().await.expect("flush"); },
            #[cfg(feature = "reactor")]
            Self::Reactor(fh) => { fh.fh_flush().await.expect("flush"); },
        }
    }
    async fn finish(self) -> hyperfile::file::FlushTimingSnapshot {
        match self {
            Self::Direct(mut h) => { let t = h.flush_timing().snapshot(); let _ = h.fs_release().await; t },
            #[cfg(feature = "reactor")]
            Self::Reactor(mut fh) => { let t = fh.fh_flush_timing().await.expect("timing"); let _ = fh.fh_release().await; t },
        }
    }
}

async fn run(client: &aws_sdk_s3::Client, fmt: BlockPtrFormat, ahead: usize, partials: bool, mib: usize, rate: usize)
    -> (Duration, Duration, Duration, u64, Vec<f64>, hyperfile::file::FlushTimingSnapshot)
{
    // Writes wrap within `BENCH_HOT_MIB` when set: a working set written over and over.
    let hot = env_num("BENCH_HOT_MIB", 0);
    let span = if hot > 0 { hot.min(mib) } else { mib };
    let tf = TestFile::new(client).await;
    let mut rt = HyperFileRuntimeConfig::default();
    // `BENCH_LIMITS=default` keeps the default dirty-data limits, which write a
    // partial — and close the object being sent ahead — when crossed.
    let default_limits = std::env::var("BENCH_LIMITS").is_ok_and(|v| v == "default");
    if !default_limits {
        rt.data_cache_dirty_max_bytes_threshold = 1 << 40;
        rt.data_cache_dirty_max_blocks_threshold = (1 << 40) / 4096;
        // With `partials`, memory pressure writes partials at the default point,
        // which the writer waits for; without, nothing goes out before the flush.
        if !partials {
            rt.segment_buffer_size = 1 << 40;
        }
    }
    rt.data_cache_dirty_max_flush_interval = u64::MAX / 2;
    rt.data_cache_blocks = 0;
    rt.stream_segment_bytes = ahead;
    let mut meta = HyperFileMetaConfig::default();
    meta.block_ptr_format = fmt;
    let config = HyperFileConfigBuilder::new()
        .with_staging_config(&StagingConfig::new_s3_uri(tf.uri(), None))
        .with_meta_config(&meta)
        .with_runtime_config(&rt)
        .build();
    let hyper = Hyper::create(client.clone(), config.clone(), HyperFileFlags::from_flags(FileFlags::rdwr()),
        HyperFileMode::from_mode(FileMode::default_file())).await.expect("create");
    #[cfg(feature = "reactor")]
    let reactor = hyperfile_reactor::Reactor::new_current().expect("reactor");
    let mut h = if std::env::var("BENCH_REACTOR").is_ok_and(|v| v == "1") {
        #[cfg(feature = "reactor")]
        { Handle::Reactor(hyperfile::file::fh::HyperFileHandler::fh_from_hyper(&reactor, hyper).await.expect("handler")) }
        #[cfg(not(feature = "reactor"))]
        { panic!("BENCH_REACTOR needs the reactor feature") }
    } else {
        Handle::Direct(hyper)
    };
    let data: Vec<Vec<u8>> = (0..8).map(chunk).collect();
    let mut longest = Duration::ZERO;
    let start = Instant::now();
    for k in 0..mib {
        if rate > 0 {
            // The producer's own pace: chunk k is ready at k MiB / rate.
            let due = Duration::from_secs_f64(k as f64 / rate as f64 * (MIB as f64 / 1e6));
            if let Some(wait) = due.checked_sub(start.elapsed()) {
                tokio::time::sleep(wait).await;
            }
        }
        let t = Instant::now();
        h.write((k % span) * MIB, &data[k % 8]).await;
        longest = longest.max(t.elapsed());
    }
    let wrote = start.elapsed();
    let t = Instant::now();
    h.flush().await;
    let flushed = t.elapsed();
    let timing = h.finish().await;

    let mut c = Hyper::open(client.clone(), config, HyperFileFlags::from_flags(FileFlags::rdonly())).await.expect("open");
    let mut buf = vec![0u8; MIB];
    let mut bad = 0;
    for k in (0..span).step_by(37).chain([span - 1]) {
        // The last write to each MiB of the span.
        let last = (mib - 1 - k) / span * span + k;
        let n = c.fs_read(k * MIB, &mut buf).await.expect("read");
        if n != MIB || buf != data[last % 8] {
            bad += 1;
        }
    }
    let _ = c.fs_release().await;
    let n = objects(client, tf.uri()).await;
    tf.cleanup(client).await;
    (wrote, flushed, longest, bad, n, timing)
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
#[ignore]
async fn bench_stream_flush() {
    let client = make_client().await;
    let mib = env_num("BENCH_MIB", 512);
    let rates: Vec<usize> = std::env::var("BENCH_RATES").unwrap_or_else(|_| "0,400,200".into())
        .split(',').filter_map(|r| r.parse().ok()).collect();
    let sizes: Vec<usize> = std::env::var("BENCH_SIZES_MIB").unwrap_or_else(|_| "4,16,256,512,1024".into())
        .split(',').filter_map(|r| r.parse().ok()).collect();
    println!("BENCH {} MiB; stream_segment_bytes {:?} MiB", mib, sizes);
    for rate in rates {
        let mut cases: Vec<(String, BlockPtrFormat, usize)> = vec![("parted, all at flush".into(), BlockPtrFormat::PartedSegment, 0)];
        cases.extend(sizes.iter().map(|m| (format!("stream {} MiB", m), BlockPtrFormat::StreamSegment, m * MIB)));
        for (name, fmt, a) in cases {
            let (w, f, longest, bad, n, t) = run(&client, fmt, a, false, mib, rate).await;
            // Data uploaded against the data the checkpoint keeps: the span, once.
            let hot = env_num("BENCH_HOT_MIB", 0);
            let span = if hot > 0 { hot.min(mib) } else { mib } as u64;
            let uploaded = if a == 0 && fmt == BlockPtrFormat::PartedSegment { span << 20 }
                else { t.stream_sent_bytes + t.stream_flush_bytes };
            let sizes = if n.is_empty() { String::new() } else {
                format!("{:.0}/{:.0}/{:.0}", n[0], n[n.len() / 2], n[n.len() - 1]) };
            println!("BENCH rate {:>4} MB/s  {:<22} flush {:>5.2} s  total {:>6.2} s  longest write {:>5.3} s  objects {:>4} (MiB min/med/max {:>11})  sent ahead {:>4} MiB (died in flight {:>4})  at flush {:>4} MiB  data uploaded {:>4.1}x  meta {:>4.1} MiB  {}",
                if rate == 0 { "max".to_string() } else { rate.to_string() }, name,
                f.as_secs_f64(), (w + f).as_secs_f64(), longest.as_secs_f64(), n.len(), sizes,
                t.stream_sent_bytes >> 20, t.stream_dead_bytes >> 20, t.stream_flush_bytes >> 20,
                uploaded as f64 / (span << 20) as f64, t.stream_meta_bytes as f64 / MIB as f64,
                if bad == 0 { "ok" } else { "DATA WRONG" });
            assert_eq!(bad, 0);
        }
    }
}