use crate::{ concurrent_processor::ProcessMap, config::Configuration, error::{Error, UploadError}, formats::InputProcessableFormat, future::LocalBoxFuture, magick::ArcPolicyDir, repo::{Alias, ArcRepo, DeleteToken, Hash, JobId, UploadId}, serde_str::Serde, store::Store, tmp_file::ArcTmpDir, }; use reqwest_middleware::ClientWithMiddleware; use std::{ path::PathBuf, sync::Arc, time::{Duration, Instant}, }; use tracing::Instrument; pub(crate) mod cleanup; mod process; const CLEANUP_QUEUE: &str = "cleanup"; const PROCESS_QUEUE: &str = "process"; const OUTDATED_PROXIES_UNIQUE_KEY: &str = "outdated-proxies"; const OUTDATED_VARIANTS_UNIQUE_KEY: &str = "outdated-variants"; const ALL_VARIANTS_UNIQUE_KEY: &str = "all-variants"; const PRUNE_MISSING_UNIQUE_KEY: &str = "prune-missing"; #[derive(Debug, serde::Deserialize, serde::Serialize)] enum Cleanup { Hash { hash: Hash, }, Identifier { identifier: String, }, Alias { alias: Serde, token: Serde, }, Variant { hash: Hash, #[serde(skip_serializing_if = "Option::is_none")] variant: Option, }, AllVariants, OutdatedVariants, OutdatedProxies, Prune, } #[derive(Debug, serde::Deserialize, serde::Serialize)] enum Process { Ingest { identifier: String, upload_id: Serde, declared_alias: Option>, }, Generate { target_format: InputProcessableFormat, source: Serde, process_path: PathBuf, process_args: Vec, }, } pub(crate) async fn cleanup_alias( repo: &ArcRepo, alias: Alias, token: DeleteToken, ) -> Result<(), Error> { let job = serde_json::to_value(Cleanup::Alias { alias: Serde::new(alias), token: Serde::new(token), }) .map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job, None).await?; Ok(()) } pub(crate) async fn cleanup_hash(repo: &ArcRepo, hash: Hash) -> Result<(), Error> { let job = serde_json::to_value(Cleanup::Hash { hash }).map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job, None).await?; Ok(()) } pub(crate) async fn cleanup_identifier(repo: &ArcRepo, identifier: &Arc) -> Result<(), Error> { let job = serde_json::to_value(Cleanup::Identifier { identifier: identifier.to_string(), }) .map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job, None).await?; Ok(()) } async fn cleanup_variants( repo: &ArcRepo, hash: Hash, variant: Option, ) -> Result<(), Error> { let job = serde_json::to_value(Cleanup::Variant { hash, variant }).map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job, None).await?; Ok(()) } pub(crate) async fn cleanup_outdated_proxies(repo: &ArcRepo) -> Result<(), Error> { let job = serde_json::to_value(Cleanup::OutdatedProxies).map_err(UploadError::PushJob)?; if repo .push(CLEANUP_QUEUE, job, Some(OUTDATED_PROXIES_UNIQUE_KEY)) .await? .is_none() { tracing::debug!("outdated proxies conflict"); } Ok(()) } pub(crate) async fn cleanup_outdated_variants(repo: &ArcRepo) -> Result<(), Error> { let job = serde_json::to_value(Cleanup::OutdatedVariants).map_err(UploadError::PushJob)?; if repo .push(CLEANUP_QUEUE, job, Some(OUTDATED_VARIANTS_UNIQUE_KEY)) .await? .is_none() { tracing::debug!("outdated variants conflict"); } Ok(()) } pub(crate) async fn cleanup_all_variants(repo: &ArcRepo) -> Result<(), Error> { let job = serde_json::to_value(Cleanup::AllVariants).map_err(UploadError::PushJob)?; if repo .push(CLEANUP_QUEUE, job, Some(ALL_VARIANTS_UNIQUE_KEY)) .await? .is_none() { tracing::debug!("all variants conflict"); } Ok(()) } pub(crate) async fn prune_missing(repo: &ArcRepo) -> Result<(), Error> { let job = serde_json::to_value(Cleanup::Prune).map_err(UploadError::PushJob)?; if repo .push(CLEANUP_QUEUE, job, Some(PRUNE_MISSING_UNIQUE_KEY)) .await? .is_none() { tracing::debug!("prune missing conflict"); } Ok(()) } pub(crate) async fn queue_ingest( repo: &ArcRepo, identifier: &Arc, upload_id: UploadId, declared_alias: Option, ) -> Result<(), Error> { let job = serde_json::to_value(Process::Ingest { identifier: identifier.to_string(), declared_alias: declared_alias.map(Serde::new), upload_id: Serde::new(upload_id), }) .map_err(UploadError::PushJob)?; repo.push(PROCESS_QUEUE, job, None).await?; Ok(()) } pub(crate) async fn queue_generate( repo: &ArcRepo, target_format: InputProcessableFormat, source: Alias, process_path: PathBuf, process_args: Vec, ) -> Result<(), Error> { let job = serde_json::to_value(Process::Generate { target_format, source: Serde::new(source), process_path, process_args, }) .map_err(UploadError::PushJob)?; repo.push(PROCESS_QUEUE, job, None).await?; Ok(()) } pub(crate) async fn process_cleanup( repo: ArcRepo, store: S, config: Configuration, ) { process_jobs(&repo, &store, &config, CLEANUP_QUEUE, cleanup::perform).await } pub(crate) async fn process_images( tmp_dir: ArcTmpDir, policy_dir: ArcPolicyDir, repo: ArcRepo, store: S, client: ClientWithMiddleware, process_map: ProcessMap, config: Configuration, ) { process_image_jobs( &tmp_dir, &policy_dir, &repo, &store, &client, &process_map, &config, PROCESS_QUEUE, process::perform, ) .await } struct MetricsGuard { worker_id: uuid::Uuid, queue: &'static str, start: Instant, armed: bool, } impl MetricsGuard { fn guard(worker_id: uuid::Uuid, queue: &'static str) -> Self { metrics::counter!("pict-rs.job.start", "queue" => queue, "worker-id" => worker_id.to_string()).increment(1); Self { worker_id, queue, start: Instant::now(), armed: true, } } fn disarm(mut self) { self.armed = false; } } impl Drop for MetricsGuard { fn drop(&mut self) { metrics::histogram!("pict-rs.job.duration", "queue" => self.queue, "worker-id" => self.worker_id.to_string(), "completed" => (!self.armed).to_string()).record(self.start.elapsed().as_secs_f64()); metrics::counter!("pict-rs.job.end", "queue" => self.queue, "worker-id" => self.worker_id.to_string(), "completed" => (!self.armed).to_string()).increment(1); } } async fn process_jobs( repo: &ArcRepo, store: &S, config: &Configuration, queue: &'static str, callback: F, ) where S: Store, for<'a> F: Fn( &'a ArcRepo, &'a S, &'a Configuration, serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { let worker_id = uuid::Uuid::new_v4(); loop { tracing::trace!("process_jobs: looping"); tokio::task::yield_now().await; let res = job_loop(repo, store, config, worker_id, queue, callback).await; if let Err(e) = res { tracing::warn!("Error processing jobs: {}", format!("{e}")); tracing::warn!("{}", format!("{e:?}")); if e.is_disconnected() { tokio::time::sleep(Duration::from_secs(10)).await; } continue; } break; } } async fn job_loop( repo: &ArcRepo, store: &S, config: &Configuration, worker_id: uuid::Uuid, queue: &'static str, callback: F, ) -> Result<(), Error> where S: Store, for<'a> F: Fn( &'a ArcRepo, &'a S, &'a Configuration, serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { loop { tracing::trace!("job_loop: looping"); tokio::task::yield_now().await; async { let (job_id, job) = repo.pop(queue, worker_id).await?; let guard = MetricsGuard::guard(worker_id, queue); let res = heartbeat( repo, queue, worker_id, job_id, (callback)(repo, store, config, job), ) .await; repo.complete_job(queue, worker_id, job_id).await?; res?; guard.disarm(); Ok(()) as Result<(), Error> } .instrument(tracing::info_span!("tick", %queue, %worker_id)) .await?; } } #[allow(clippy::too_many_arguments)] async fn process_image_jobs( tmp_dir: &ArcTmpDir, policy_dir: &ArcPolicyDir, repo: &ArcRepo, store: &S, client: &ClientWithMiddleware, process_map: &ProcessMap, config: &Configuration, queue: &'static str, callback: F, ) where S: Store, for<'a> F: Fn( &'a ArcTmpDir, &'a ArcPolicyDir, &'a ArcRepo, &'a S, &'a ClientWithMiddleware, &'a ProcessMap, &'a Configuration, serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { let worker_id = uuid::Uuid::new_v4(); loop { tracing::trace!("process_image_jobs: looping"); tokio::task::yield_now().await; let res = image_job_loop( tmp_dir, policy_dir, repo, store, client, process_map, config, worker_id, queue, callback, ) .await; if let Err(e) = res { tracing::warn!("Error processing jobs: {}", format!("{e}")); tracing::warn!("{}", format!("{e:?}")); if e.is_disconnected() { tokio::time::sleep(Duration::from_secs(10)).await; } continue; } break; } } #[allow(clippy::too_many_arguments)] async fn image_job_loop( tmp_dir: &ArcTmpDir, policy_dir: &ArcPolicyDir, repo: &ArcRepo, store: &S, client: &ClientWithMiddleware, process_map: &ProcessMap, config: &Configuration, worker_id: uuid::Uuid, queue: &'static str, callback: F, ) -> Result<(), Error> where S: Store, for<'a> F: Fn( &'a ArcTmpDir, &'a ArcPolicyDir, &'a ArcRepo, &'a S, &'a ClientWithMiddleware, &'a ProcessMap, &'a Configuration, serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { loop { tracing::trace!("image_job_loop: looping"); tokio::task::yield_now().await; async { let (job_id, job) = repo.pop(queue, worker_id).await?; let guard = MetricsGuard::guard(worker_id, queue); let res = heartbeat( repo, queue, worker_id, job_id, (callback)( tmp_dir, policy_dir, repo, store, client, process_map, config, job, ), ) .await; repo.complete_job(queue, worker_id, job_id).await?; res?; guard.disarm(); Ok(()) as Result<(), Error> } .instrument(tracing::info_span!("tick", %queue, %worker_id)) .await?; } } #[tracing::instrument("running-job", skip(repo, queue, worker_id, fut))] async fn heartbeat( repo: &ArcRepo, queue: &'static str, worker_id: uuid::Uuid, job_id: JobId, fut: Fut, ) -> Fut::Output where Fut: std::future::Future, { let mut fut = std::pin::pin!(fut.instrument(tracing::info_span!("job-future"))); let mut interval = tokio::time::interval(Duration::from_secs(5)); let mut hb = None; loop { tracing::trace!("heartbeat: looping"); tokio::task::yield_now().await; tokio::select! { biased; output = &mut fut => { return output; } _ = interval.tick() => { if hb.is_none() { hb = Some(repo.heartbeat(queue, worker_id, job_id)); } } opt = poll_opt(hb.as_mut()), if hb.is_some() => { hb.take(); if let Some(Err(e)) = opt { tracing::warn!("Failed heartbeat\n{}", format!("{e:?}")); } } } } } async fn poll_opt(opt: Option<&mut Fut>) -> Option where Fut: std::future::Future + Unpin, { match opt { None => None, Some(fut) => Some(fut.await), } }