s3ueeze 0.1.0

use anyhow::{Context, Result, bail};
use aws_sdk_s3::{Client, types::Object};
use buildstructor::buildstructor;
use tokio::{
    io::AsyncBufReadExt,
    select,
    sync::mpsc::{UnboundedReceiver, UnboundedSender},
};
use tokio_util::sync::CancellationToken;
use tracing::{error, info, info_span, instrument, trace, warn};
use url::Url;

#[derive(Debug, Clone)]
pub struct Worker {
    client: Client,
    objects_tx: UnboundedSender<(String, Vec<Object>)>,
    bytes_tx: UnboundedSender<String>,
    cancel: CancellationToken,
}

#[buildstructor]
impl Worker {
    #[builder]
    pub fn new(
        client: Client,
        objects_tx: UnboundedSender<(String, Vec<Object>)>,
        bytes_tx: UnboundedSender<String>,
        cancel: CancellationToken,
    ) -> Self {
        Self {
            client,
            objects_tx,
            bytes_tx,
            cancel,
        }
    }

    #[cfg_attr(feature = "tracing", instrument(level = "debug", err, skip_all))]
    pub async fn work(
        self,
        mut sources_rx: UnboundedReceiver<Url>,
        mut contents_rx: UnboundedReceiver<(String, Object)>,
    ) -> Result<()> {
        // Just run an endless loop, waiting for events to come in on a channel.
        loop {
            select! {
                // Bias for responding to cancellation, and reading objects.
                // Otherwise we wast time and memory queueing up a ton of
                // objects to process. Usually there are far fewer sources than
                // objects.
                biased;

                () = self.cancel.cancelled() => {
                    #[cfg(feature = "tracing")]
                    info!("cancellation received; shutting down");

                    return Ok(());
                }
                Some((bucket, object)) = contents_rx.recv() => {
                    // Go download an object from a bucket.
                    get_object(&self.client, bucket, object, self.bytes_tx.clone()).await?;
                },
                Some(url) = sources_rx.recv() => {
                    list_objects(&self.client, url, self.objects_tx.clone()).await?;
                },
            }
        }
    }
}

#[cfg_attr(feature = "tracing",
    instrument(level = "TRACE", err, skip(client, tx), fields(url = %url))
)]
async fn list_objects(
    client: &Client,
    url: Url,
    tx: UnboundedSender<(String, Vec<Object>)>,
) -> Result<()> {
    let bucket = url.host().unwrap().to_string();
    let prefix = url.path().trim_start_matches('/');
    #[cfg(feature = "tracing")]
    trace!(%bucket, %prefix);

    // Construct a paginated stream of object list results.
    let mut stream = client
        .list_objects_v2()
        .bucket(&bucket)
        .prefix(prefix)
        .into_paginator()
        .send();

    // Process each page.
    while let Some(page) = stream.next().await {
        let page = page.context("list_objects_v2")?;

        // Get the objects; we should never have an empty page.
        let Some(objects) = page.contents else {
            #[cfg(feature = "tracing")]
            warn!(%bucket, %prefix, "no objects");

            return Ok(());
        };

        // Send the objects out to be downloaded.
        tx.send((bucket.clone(), objects))?;
    }

    Ok(())
}

/// Download an object from a bucket, and send the contents back to the main
/// thread line-wise.
#[cfg_attr(
    feature = "tracing",
    instrument(level = "TRACE", ret, err, skip_all, fields(
        bucket = %bucket,
        key = object.key,
    ))
)]
async fn get_object(
    client: &Client,
    bucket: String,
    object: Object,
    tx: UnboundedSender<String>,
) -> Result<()> {
    // Fetch the object.
    let object = match client
        .get_object()
        .bucket(bucket)
        .set_key(object.key().map(Into::into))
        .send()
        .await
    {
        Ok(o) => o,
        Err(e) => {
            #[cfg(feature = "tracing")]
            error!(?e, "get_object");
            bail!(e);
        },
    };

    // Get an AsyncBufReader.
    let reader = object.body.into_async_read();

    // Read line-wise from the reader.
    let mut lines = reader.lines();
    let span = info_span!("send.lines");

    // Send each line back over the data channel for writing to the file.
    while let Some(line) = lines.next_line().await? {
        let _guard = span.enter();
        tx.send(line)?;
    }

    Ok(())
}