s3ueeze 0.0.3

mod metric;
mod worker;

use anyhow::{bail, Result};
use aws_config::BehaviorVersion;
use aws_sdk_s3::Client;
use buildstructor::buildstructor;
use metric::Metric;
use std::{fmt::Debug, num::NonZeroUsize, thread, time::Duration};
use tokio::{
    fs::OpenOptions,
    io::{AsyncWriteExt, BufWriter},
    select, spawn,
    sync::mpsc::{self},
    time::interval,
};
use tokio_util::sync::CancellationToken;
#[cfg(feature = "tracing")]
use tracing::{debug, info, instrument, trace};
use uom::{
    fmt::DisplayStyle,
    si::{
        f64::{Information, InformationRate, Time},
        information::{byte, mebibyte},
        information_rate::megabit_per_second,
        time::second,
    },
};
use url::Url;
use worker::Worker;

#[derive(Clone, Debug)]
pub struct S3ueeze {
    client: Client,
    destination: Url,
    sources: Vec<Url>,
    worker_count: usize,
}

#[buildstructor]
impl S3ueeze {
    #[builder]
    #[must_use]
    #[cfg_attr(feature = "tracing", instrument)]
    pub async fn new(
        client: Option<Client>,
        sources: Vec<Url>,
        destination: Url,
        worker_count: Option<NonZeroUsize>,
    ) -> Result<Self> {
        let client = if let Some(client) = client {
            client
        } else {
            let config =
                aws_config::load_defaults(BehaviorVersion::latest()).await;
            Client::new(&config)
        };

        for url in &sources {
            if url.scheme() != "s3" {
                bail!("source schemes must all be s3://: {:?}", url);
            }

            if url.host().is_none() {
                bail!(
                    "source hosts must all be valid s3 bucket names: {:?}",
                    url
                );
            }
        }

        if destination.scheme() != "file" {
            bail!("destination scheme must be path://: {:?}", destination);
        }

        let worker_count = if let Some(workers) = worker_count {
            workers.get()
        } else {
            thread::available_parallelism()?
                .get()
                .checked_mul(5)
                .unwrap()
        };

        Ok(Self {
            client,
            destination,
            sources,
            worker_count,
        })
    }

    #[must_use]
    pub fn worker_count(&self) -> usize {
        self.worker_count
    }

    #[allow(clippy::too_many_lines)]
    #[cfg_attr(feature = "tracing", instrument(err, skip_all))]
    pub async fn run(
        &self,
        cancel: CancellationToken,
    ) -> Result<(u64, u64, u64)> {
        // Create channels for all the kinds of responses we need from the workers.
        let (objects_tx, mut objects_rx) = mpsc::unbounded_channel();
        let (bytes_tx, mut bytes_rx) = mpsc::unbounded_channel();

        // Collection of worker channels.
        let mut workers = vec![];

        // Go through and spawn workers connected to all the channels.
        for n in 0..self.worker_count {
            #[cfg(feature = "tracing")]
            trace!("spawning worker {n}");

            // Create channels for all the requests the workers need to receive.
            // Make a new channel on which the worker will receive messages.
            let (sources_tx, sources_rx) = mpsc::unbounded_channel();
            let (contents_tx, contents_rx) = mpsc::unbounded_channel();

            // Construct a worker.
            let worker = Worker::builder()
                .client(self.client.clone())
                .objects_tx(objects_tx.clone())
                .bytes_tx(bytes_tx.clone())
                .cancel(cancel.clone())
                .build();

            // Spawn the async worker, storing its join handle and channels.
            workers.push((
                spawn(
                    async move { worker.work(sources_rx, contents_rx).await },
                ),
                sources_tx,
                contents_tx,
            ));
        }
        #[cfg(feature = "tracing")]
        debug!("done spawning {} workers", workers.len());

        // Open a file for writing the output to.
        let f = OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(self.destination.path())
            .await?;
        let mut f = BufWriter::new(f);

        for (source, (n, (_, tx, _))) in
            self.sources.iter().zip(workers.iter().enumerate().cycle())
        {
            #[cfg(feature = "tracing")]
            debug!("sending source {source} to worker {n}");
            tx.send(source.clone())?;
        }

        // Set up a ticker for status updates.
        let mut ticker = interval(Duration::from_secs(5));

        let mut processed_sources = 0;
        let mut processed_bytes: Metric<u64> = Metric::default();
        let mut total_objects: u64 = 0;
        let mut total_bytes: u64 = 0;
        let mut processed_lines = 0;
        loop {
            select! {
                biased;

                () = cancel.cancelled() => {
                    break;
                }
                _ = ticker.tick() => {
                    let new = processed_bytes.flush();
                    if new < 1 && processed_bytes.total()? > 0 {
                        cancel.cancel();

                        break;
                    }

                    let bytes = Information::new::<byte>(new as f64);

                    // Calculate Mbps.
                    let speed: InformationRate = (
                        bytes / Time::new::<second>(5.0)
                    ).into();
                    let speed = speed
                        .into_format_args(
                            megabit_per_second,
                            DisplayStyle::Abbreviation,
                        );
                    let bytes = bytes
                        .into_format_args(
                            mebibyte,
                            DisplayStyle::Abbreviation,
                        );

                        let percent = if total_bytes > 0 {
                            processed_bytes.total()? as f64 / total_bytes as f64 * 100.0
                        } else {
                            0.0
                        };

                    eprintln!("{percent:.1}%, {speed}");
                    #[cfg(feature = "tracing")]
                    info!(
                        %speed,
                        %processed_lines,
                        %processed_sources,
                        %total_objects,
                        processed_bytes = processed_bytes.total()?,
                        %bytes,
                    );

                    if self.sources.len() == processed_sources
                        && total_bytes == processed_bytes.total()?
                    {
                        #[cfg(feature = "tracing")]
                        info!("done; shutting down workers");

                        // Shut down the workers cleanly.
                        cancel.cancel();

                        break;
                    }
                }
                // Receive a line of JSON (these are JSON-object-per-line
                // files) and write it out to disk, collecting metrics along
                // the way.
                Some(line) = bytes_rx.recv() => {
                    f.write_all(line.as_bytes()).await?;
                    f.write_u8(b'\n').await?;

                    processed_bytes.add(line.len() as u64 +1);
                    processed_lines += 1;
                }
                // Receive a list of objects from a bucket and fan them out to
                // workers to fetch.
                Some((bucket, objects)) = objects_rx.recv() => {
                    #[cfg(feature = "tracing")]
                    debug!(?bucket, objects = ?objects.len());

                    for (
                        object,
                        (n, (_, _, contents_tx))
                    ) in objects.iter().zip(workers.iter().enumerate().cycle())
                    {
                        #[cfg(feature = "tracing")]
                        trace!("sending {:?} to worker {n}", object.key());
                        contents_tx.send((
                            bucket.clone(),
                            object.clone(),
                        ))?;

                        let Some(size) = object.size() else {
                            bail!("object missing size: object key: {:?}", object.key());
                        };

                        total_bytes += u64::try_from(size)?;
                    }

                    processed_sources += 1;
                    total_objects += objects.len() as u64;
                },
            }
        }

        // Go through and shut down each worker, checking the return value.
        for (handle, _, _) in workers {
            handle.await??;
        }

        processed_bytes.flush();
        Ok((total_objects, processed_bytes.total()?, processed_lines))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    macro_rules! url {
        ($str:expr) => {
            Url::try_from($str).unwrap()
        };
    }

    #[tokio::test]
    async fn test_config() {
        // Base case.
        S3ueeze::builder()
            .destination(url!("file:///tmp/foobar"))
            .build()
            .await
            .unwrap();

        // With one source.
        S3ueeze::builder()
            .source(url!("s3://bucket/tmp/foobar"))
            .destination(url!("file:///tmp/foobar"))
            .build()
            .await
            .unwrap();

        // With multiple sources.
        S3ueeze::builder()
            .sources(vec![
                url!("s3://bucket/tmp/foobar"),
                url!("s3://bucket/tmp/foobar"),
            ])
            .destination(url!("file:///tmp/foobar"))
            .build()
            .await
            .unwrap();
    }
}