4.1 KiB · text History 376a05c
use serde::{Deserialize, Serialize};
use uuid::Uuid;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum JobKind {
    /// Extract a still frame from a video so the still-image categorizers
    /// have something to work on. Chains into thumbnail/ocr/describe.
    Poster,
    /// Short animated WebP for the grid, from a video or GIF. Independent of
    /// the poster: it reads the original, so runs alongside it.
    Preview,
    Thumbnail,
    Ocr,
    Describe,
    ExtractTags,
    SimilaritySearch,
}

/// Every kind, in pipeline order. Anything that enumerates kinds — queue
/// topology, depth probes, the retry UI — must iterate this rather than
/// keeping its own list: a kind missing from `declare_topology` has no queue,
/// and messages published to it are dropped by the broker without a word.
pub const ALL_KINDS: [JobKind; 7] = [
    JobKind::Poster,
    JobKind::Preview,
    JobKind::Thumbnail,
    JobKind::Ocr,
    JobKind::Describe,
    JobKind::ExtractTags,
    JobKind::SimilaritySearch,
];

impl JobKind {
    pub fn as_str(&self) -> &'static str {
        match self {
            JobKind::Poster => "poster",
            JobKind::Preview => "preview",
            JobKind::Thumbnail => "thumbnail",
            JobKind::Ocr => "ocr",
            JobKind::Describe => "describe",
            JobKind::ExtractTags => "extract_tags",
            JobKind::SimilaritySearch => "similarity_search",
        }
    }

    pub fn from_str(s: &str) -> Option<Self> {
        match s {
            "poster" | "frame" => Some(JobKind::Poster),
            "preview" => Some(JobKind::Preview),
            "thumbnail" | "thumbs" => Some(JobKind::Thumbnail),
            "ocr" => Some(JobKind::Ocr),
            "describe" => Some(JobKind::Describe),
            "extract_tags" | "extract-tags" => Some(JobKind::ExtractTags),
            "similarity_search" | "similarity" => Some(JobKind::SimilaritySearch),
            _ => None,
        }
    }
}

/// Envelope sent over RabbitMQ for every job. The kind is duplicated in the
/// routing key but kept here so consumers can re-publish without inspecting
/// AMQP metadata.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JobMessage {
    pub kind: JobKind,
    pub upload_id: Uuid,
    pub user_id: Uuid,
    /// UUID generated at publish time so workers can dedupe redeliveries
    /// across queue restarts.
    pub job_id: String,
    pub attempt: u32,
}

impl JobMessage {
    pub fn new(kind: JobKind, upload_id: Uuid, user_id: Uuid) -> Self {
        Self {
            kind,
            upload_id,
            user_id,
            job_id: Uuid::new_v4().to_string(),
            attempt: 1,
        }
    }
}

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

    #[test]
    fn all_kinds_covers_every_variant() {
        // The match below is deliberately exhaustive with no wildcard: adding a
        // JobKind variant stops this file compiling until the author has been
        // here, and the assert then forces it into ALL_KINDS. That list drives
        // `declare_topology`, and a kind with no declared queue has its
        // messages dropped by the broker without an error anywhere.
        fn index(kind: JobKind) -> usize {
            match kind {
                JobKind::Poster => 0,
                JobKind::Preview => 1,
                JobKind::Thumbnail => 2,
                JobKind::Ocr => 3,
                JobKind::Describe => 4,
                JobKind::ExtractTags => 5,
                JobKind::SimilaritySearch => 6,
            }
        }
        let mut seen = [false; ALL_KINDS.len()];
        for kind in ALL_KINDS {
            seen[index(kind)] = true;
        }
        assert!(
            seen.iter().all(|s| *s),
            "a JobKind variant is missing from ALL_KINDS"
        );
    }

    #[test]
    fn kind_strings_round_trip_and_are_unique() {
        let mut names = std::collections::HashSet::new();
        for kind in ALL_KINDS {
            let name = kind.as_str();
            assert!(names.insert(name), "duplicate kind string: {name}");
            assert_eq!(JobKind::from_str(name), Some(kind), "{name} did not round-trip");
        }
    }
}