pict-rs/src/generate.rs

mod ffmpeg;
mod magick;

use crate::{
    details::Details,
    error::{Error, UploadError},
    formats::{ImageFormat, InputProcessableFormat, InternalVideoFormat, ProcessableFormat},
    future::{WithMetrics, WithPollTimer, WithTimeout},
    repo::{Hash, NotificationEntry, VariantAlreadyExists},
    state::State,
    store::Store,
};

use std::{
    future::Future,
    sync::Arc,
    time::{Duration, Instant},
};

use tracing::Instrument;

struct MetricsGuard {
    start: Instant,
    armed: bool,
}

impl MetricsGuard {
    fn guard() -> Self {
        metrics::counter!(crate::init_metrics::GENERATE_START).increment(1);
        Self {
            start: Instant::now(),
            armed: true,
        }
    }

    fn disarm(mut self) {
        self.armed = false;
    }
}

impl Drop for MetricsGuard {
    fn drop(&mut self) {
        metrics::histogram!(crate::init_metrics::GENERATE_DURATION, "completed" => (!self.armed).to_string())
            .record(self.start.elapsed().as_secs_f64());
        metrics::counter!(crate::init_metrics::GENERATE_END, "completed" => (!self.armed).to_string())
            .increment(1);
    }
}

#[tracing::instrument(skip(state, original_details, hash))]
pub(crate) async fn generate<S: Store + 'static>(
    state: &State<S>,
    format: InputProcessableFormat,
    variant: String,
    variant_args: Vec<String>,
    original_details: &Details,
    hash: Hash,
) -> Result<(Details, Arc<str>), Error> {
    if state.config.server.danger_dummy_mode {
        let identifier = state
            .repo
            .identifier(hash)
            .await?
            .ok_or(UploadError::MissingIdentifier)?;

        Ok((original_details.clone(), identifier))
    } else {
        let mut attempts = 0;
        let tup = loop {
            if attempts > 2 {
                return Err(UploadError::ProcessTimeout.into());
            }

            match state
                .repo
                .claim_variant_processing_rights(hash.clone(), variant.clone())
                .await?
            {
                Ok(()) => {
                    // process
                    let process_future = process(
                        state,
                        format,
                        variant.clone(),
                        variant_args,
                        original_details,
                        hash.clone(),
                    )
                    .with_poll_timer("process-future");

                    let res = heartbeat(state, hash.clone(), variant.clone(), process_future)
                        .with_poll_timer("heartbeat-future")
                        .with_timeout(Duration::from_secs(state.config.media.process_timeout * 4))
                        .with_metrics(crate::init_metrics::GENERATE_PROCESS)
                        .await
                        .map_err(|_| Error::from(UploadError::ProcessTimeout));

                    state
                        .repo
                        .notify_variant(hash.clone(), variant.clone())
                        .await?;

                    break res???;
                }
                Err(entry) => {
                    if let Some(tuple) = wait_timeout(
                        hash.clone(),
                        variant.clone(),
                        entry,
                        state,
                        Duration::from_secs(20),
                    )
                    .await?
                    {
                        break tuple;
                    }

                    attempts += 1;
                }
            }
        };

        Ok(tup)
    }
}

pub(crate) async fn wait_timeout<S: Store + 'static>(
    hash: Hash,
    variant: String,
    mut entry: NotificationEntry,
    state: &State<S>,
    timeout: Duration,
) -> Result<Option<(Details, Arc<str>)>, Error> {
    let notified = entry.notified_timeout(timeout);

    if let Some(identifier) = state
        .repo
        .variant_identifier(hash.clone(), variant.clone())
        .await?
    {
        let details = crate::ensure_details_identifier(state, &identifier).await?;

        return Ok(Some((details, identifier)));
    }

    match notified.await {
        Ok(()) => tracing::debug!("notified"),
        Err(_) => tracing::debug!("timeout"),
    }

    Ok(None)
}

async fn heartbeat<S, O>(
    state: &State<S>,
    hash: Hash,
    variant: String,
    future: impl Future<Output = O>,
) -> Result<O, Error> {
    let repo = state.repo.clone();

    let handle = crate::sync::abort_on_drop(crate::sync::spawn("heartbeat-task", async move {
        let mut interval = tokio::time::interval(Duration::from_secs(5));

        loop {
            interval.tick().await;

            if let Err(e) = repo.variant_heartbeat(hash.clone(), variant.clone()).await {
                break Error::from(e);
            }
        }
    }));

    let future = std::pin::pin!(future);

    tokio::select! {
        biased;
        output = future => {
            Ok(output)
        }
        res = handle => {
            Err(res.map_err(|_| UploadError::Canceled)?)
        }
    }
}

#[tracing::instrument(skip(state, hash))]
async fn process<S: Store + 'static>(
    state: &State<S>,
    output_format: InputProcessableFormat,
    variant: String,
    variant_args: Vec<String>,
    original_details: &Details,
    hash: Hash,
) -> Result<(Details, Arc<str>), Error> {
    let guard = MetricsGuard::guard();
    let permit = crate::process_semaphore().acquire().await?;

    let identifier = input_identifier(state, output_format, hash.clone(), original_details).await?;

    let input_details = crate::ensure_details_identifier(state, &identifier).await?;

    let input_format = input_details
        .internal_format()
        .processable_format()
        .expect("Already verified format is processable");

    let format = input_format.process_to(output_format);

    let quality = match format {
        ProcessableFormat::Image(format) => state.config.media.image.quality_for(format),
        ProcessableFormat::Animation(format) => state.config.media.animation.quality_for(format),
    };

    let stream = state.store.to_stream(&identifier, None, None).await?;

    let bytes =
        crate::magick::process_image_command(state, variant_args, input_format, format, quality)
            .await?
            .drive_with_stream(stream)
            .into_bytes_stream()
            .instrument(tracing::info_span!(
                "Reading processed image to BytesStream"
            ))
            .await?;

    drop(permit);

    let details = Details::from_bytes_stream(state, bytes.clone()).await?;

    let identifier = state
        .store
        .save_stream(
            bytes.into_io_stream(),
            details.media_type(),
            Some(details.file_extension()),
        )
        .await?;

    let identifier = if let Err(VariantAlreadyExists) = state
        .repo
        .relate_variant_identifier(hash.clone(), variant.clone(), &identifier)
        .await?
    {
        state.store.remove(&identifier).await?;
        state
            .repo
            .variant_identifier(hash, variant)
            .await?
            .ok_or(UploadError::MissingIdentifier)?
    } else {
        state.repo.relate_details(&identifier, &details).await?;
        identifier
    };

    guard.disarm();

    Ok((details, identifier)) as Result<(Details, Arc<str>), Error>
}

pub(crate) async fn ensure_motion_identifier<S>(
    state: &State<S>,
    hash: Hash,
    original_details: &Details,
) -> Result<Arc<str>, Error>
where
    S: Store + 'static,
{
    if let Some(identifier) = state.repo.motion_identifier(hash.clone()).await? {
        return Ok(identifier);
    };

    let identifier = state
        .repo
        .identifier(hash.clone())
        .await?
        .ok_or(UploadError::MissingIdentifier)?;

    let (reader, media_type, file_extension) =
        if let Some(processable_format) = original_details.internal_format().processable_format() {
            let thumbnail_format = state.config.media.image.format.unwrap_or(ImageFormat::Webp);

            let stream = state.store.to_stream(&identifier, None, None).await?;

            let process =
                magick::thumbnail_command(state, processable_format, thumbnail_format).await?;

            (
                process.drive_with_stream(stream),
                thumbnail_format.media_type(),
                thumbnail_format.file_extension(),
            )
        } else {
            let thumbnail_format = match state.config.media.image.format {
                Some(ImageFormat::Webp | ImageFormat::Avif | ImageFormat::Jxl) => {
                    ffmpeg::ThumbnailFormat::Webp
                }
                Some(ImageFormat::Png) => ffmpeg::ThumbnailFormat::Png,
                Some(ImageFormat::Jpeg) | None => ffmpeg::ThumbnailFormat::Jpeg,
            };

            let reader = ffmpeg::thumbnail(
                state,
                identifier,
                original_details
                    .video_format()
                    .unwrap_or(InternalVideoFormat::Mp4),
                thumbnail_format,
            )
            .await?;

            (
                reader,
                thumbnail_format.media_type(),
                thumbnail_format.file_extension(),
            )
        };

    let motion_identifier = reader
        .with_stdout(|stdout| async {
            state
                .store
                .save_stream(
                    tokio_util::io::ReaderStream::with_capacity(stdout, 1024 * 64),
                    media_type,
                    Some(file_extension),
                )
                .await
        })
        .await??;

    state
        .repo
        .relate_motion_identifier(hash, &motion_identifier)
        .await?;

    Ok(motion_identifier)
}

#[tracing::instrument(skip_all)]
async fn input_identifier<S>(
    state: &State<S>,
    output_format: InputProcessableFormat,
    hash: Hash,
    original_details: &Details,
) -> Result<Arc<str>, Error>
where
    S: Store + 'static,
{
    let should_thumbnail =
        if let Some(input_format) = original_details.internal_format().processable_format() {
            let output_format = input_format.process_to(output_format);

            input_format.should_thumbnail(output_format)
        } else {
            // video case
            true
        };

    if should_thumbnail {
        return ensure_motion_identifier(state, hash.clone(), original_details).await;
    }

    state
        .repo
        .identifier(hash)
        .await?
        .ok_or(UploadError::MissingIdentifier)
        .map_err(From::from)
}