2.9 KiB · text History 1ce9641
use async_trait::async_trait;
use futures_util::StreamExt;
use lapin::{
    options::{BasicAckOptions, BasicConsumeOptions, BasicNackOptions, BasicQosOptions},
    types::FieldTable,
    Connection,
};
use tracing::{error, info, warn};

use super::{messages::JobMessage, publisher::declare_topology, queue_name, JobKind};

pub enum JobOutcome {
    /// Work succeeded; ack the message.
    Ack,
    /// Transient failure; nack with requeue. Use sparingly — repeated requeue
    /// loops are a worker bug. Bounded retry handled by JobMessage.attempt.
    Retry,
    /// Permanent failure; nack without requeue, message routes to DLQ.
    Reject,
}

#[async_trait]
pub trait JobHandler: Send + Sync {
    async fn handle(&self, msg: &JobMessage) -> JobOutcome;
}

pub async fn run_worker(
    conn: &Connection,
    kind: JobKind,
    prefetch: u16,
    handler: impl JobHandler + 'static,
) -> Result<(), lapin::Error> {
    let channel = conn.create_channel().await?;
    declare_topology(&channel).await?;
    channel
        .basic_qos(prefetch, BasicQosOptions::default())
        .await?;

    let qname = queue_name(kind);
    let consumer_tag = format!("dms-worker-{}-{}", kind.as_str(), std::process::id());
    let mut consumer = channel
        .basic_consume(
            &qname,
            &consumer_tag,
            BasicConsumeOptions::default(),
            FieldTable::default(),
        )
        .await?;
    info!(queue = %qname, prefetch, "worker consuming");

    while let Some(delivery) = consumer.next().await {
        let delivery = match delivery {
            Ok(d) => d,
            Err(e) => {
                error!(error = %e, "consumer stream error");
                continue;
            }
        };

        let msg: JobMessage = match serde_json::from_slice(&delivery.data) {
            Ok(m) => m,
            Err(e) => {
                error!(error = %e, "malformed job, sending to DLQ");
                let _ = delivery
                    .nack(BasicNackOptions { multiple: false, requeue: false })
                    .await;
                continue;
            }
        };

        let outcome = handler.handle(&msg).await;
        match outcome {
            JobOutcome::Ack => {
                if let Err(e) = delivery.ack(BasicAckOptions::default()).await {
                    error!(error = %e, "ack failed");
                }
            }
            JobOutcome::Retry => {
                warn!(job_id = %msg.job_id, "transient failure, requeuing");
                let _ = delivery
                    .nack(BasicNackOptions { multiple: false, requeue: true })
                    .await;
            }
            JobOutcome::Reject => {
                warn!(job_id = %msg.job_id, "permanent failure, sending to DLQ");
                let _ = delivery
                    .nack(BasicNackOptions { multiple: false, requeue: false })
                    .await;
            }
        }
    }
    Ok(())
}