s3ueeze 0.1.0

use anyhow::Result;
use clap::{ArgAction, Parser};
use s3ueeze::S3ueeze;
use std::{fmt::Display, num::NonZeroUsize, process, time::Instant};
use tokio::{
    select,
    signal::unix::{SignalKind, signal},
};
use tokio_util::sync::CancellationToken;
use tracing::{debug, instrument, warn};
use uom::fmt::DisplayStyle;
use uom::si::f64::{Information, InformationRate, Time};
use uom::si::information::{byte, gibibyte, kibibyte, mebibyte};
use uom::si::information_rate::megabit_per_second;
use uom::si::time::second;
use url::Url;

/// Collect JSONline data from multiple objects in S3 and merge them together
/// into a single local file.
#[derive(Debug, Parser)]
#[command(
    author,
    version = format!(
        "{} {}{} {:?}",
        env!("CARGO_PKG_VERSION"),
        env!("VERGEN_GIT_SHA"),
        if env!("VERGEN_GIT_DIRTY") == "true" {
            "-dirty"
        } else {
            ""
        },
        env!("VERGEN_GIT_COMMIT_MESSAGE")
            .lines()
            .next()
            .expect("missing commit message"),
    ),
    max_term_width = 80,
    long_about = termimad::text(include_str!("../README.md")).to_string(),
    styles = clap_themes::CARGO,
)]
struct Cli {
    /// Increase logging verbosity.
    #[clap(short, long, action = ArgAction::Count, global = true)]
    verbose: u8,

    /// Number of workers to spawn. Defaults to number of cores.
    #[clap(short, long)]
    workers: Option<NonZeroUsize>,

    /// Destination where to write the file. May be local or remote; prefix with
    /// s3:// to use a remote destination.
    #[clap(short, long)]
    destination: Url,

    /// Sources of s3 files to merge. If a "directory", will recurse into the
    /// directory indefinitely. Don't pass a trailing slash.
    sources: Vec<Url>,
}

#[tokio::main]
#[instrument]
async fn main() -> Result<()> {
    let cli = Cli::parse();

    jacklog::from_level!(2 + cli.verbose)?;
    debug!(?cli);

    // Check the nofile (number of open file descriptors, `ulimit -n`, that we
    // can open. We need to ensure we'll have enough for all our network
    // connections.
    let (cur, max) = unsafe {
        let mut limit = libc::rlimit {
            rlim_cur: 0,
            rlim_max: 0,
        };
        libc::getrlimit(libc::RLIMIT_NOFILE, &mut limit);

        (limit.rlim_cur, limit.rlim_max)
    };
    debug!(?cur, ?max, max_safe_workers = cur / 2 - 1, "found fd limit");

    let workers = if matches!(cli, Cli {
        workers: Some(workers),
        ..
    } if workers.get() as u64 > cur / 2 - 1)
    {
        // We can use up to half the max fd -1 for the output file we'll hold
        // open.
        let n = cur / 2 - 1;
        warn!(workers = ?cli.workers, nofile = ?cur, "worker count > nofile/2-1; reducing to {n}");
        NonZeroUsize::new(n.try_into()?)
    } else {
        cli.workers
    };

    // Construct the config for the squeezer.
    let squeezer = S3ueeze::builder()
        .sources(cli.sources)
        .destination(cli.destination.clone())
        .and_worker_count(workers)
        .build()
        .await?;

    let start = Instant::now();
    let cancel = CancellationToken::new();
    let mut handler = signal(SignalKind::interrupt())?;

    // Launch the workers.
    eprintln!("spawning {} workers", squeezer.worker_count());
    let run = squeezer.run(cancel.clone());
    tokio::pin!(run);

    // Wait for the job to complete, or for sigint.
    let mut quit = false;
    let (objects, bytes, lines) = loop {
        select! {
            biased;

            _ = handler.recv() => {
                if quit {
                    eprintln!("received repeated SIGINT; aborting");
                    process::exit(1);
                }

                eprintln!("received SIGINT; shutting down cleanly...");
                cancel.cancel();
                quit = true;
                continue;
            }
            res = &mut run => break res?,
        }
    };
    let elapsed = Instant::now().checked_duration_since(start);

    // Get the bytes into uom so we can do more intelligent output.
    let uom_bytes = Information::new::<byte>(bytes as f64);

    const MIB: u64 = 2 ^ 20;
    const GIB: u64 = MIB * 1024;

    let human_bytes = match bytes {
        0..MIB => format!(
            "{:.2}",
            uom_bytes.into_format_args(kibibyte, DisplayStyle::Abbreviation)
        ),
        MIB..=GIB => format!(
            "{:.2}",
            uom_bytes.into_format_args(mebibyte, DisplayStyle::Abbreviation)
        ),
        _ => format!(
            "{:.2}",
            uom_bytes.into_format_args(gibibyte, DisplayStyle::Abbreviation)
        ),
    };
    eprint!(
        "read {human_bytes} in {} lines from {} objects",
        commatize(lines),
        commatize(objects),
    );

    if let Some(elapsed) = elapsed {
        let speed: InformationRate =
            (uom_bytes / Time::new::<second>(elapsed.as_secs() as f64)).into();
        eprint!(
            " in {} seconds ({:.2})",
            commatize(elapsed.as_secs()),
            speed.into_format_args(
                megabit_per_second,
                DisplayStyle::Abbreviation
            ),
        );
    }
    eprintln!("\ns3ueezed file now available at {}", cli.destination);

    if quit {
        eprintln!("forced clean shutdown");
        process::exit(9);
    }

    Ok(())
}

fn commatize<T: Display>(n: T) -> String {
    let mut s = n
        .to_string()
        .chars()
        .rev()
        .enumerate()
        .flat_map(|(n, c)| {
            if n > 0 && n % 3 == 0 {
                vec![',', c]
            } else {
                vec![c]
            }
        })
        .collect::<Vec<_>>();
    s.reverse();
    s.into_iter().collect()
}

#[test]
fn test_commatize() {
    assert_eq!(commatize(123), "123");
    assert_eq!(commatize(1234), "1,234");
}