mailks/status.rs

use std::{
    collections::HashMap,
    sync::{Arc, RwLock},
    time::{Duration, Instant},
};

use tracing;

use crate::pb::email::v1::SendStatus;

const STATUS_TTL: Duration = Duration::from_secs(24 * 60 * 60);
const MAX_STATUS_ENTRIES: usize = 10_000;

#[derive(Debug, Clone, Default)]
pub struct JobStatusStore {
    inner: Arc<RwLock<HashMap<u64, JobStatusEntry>>>,
}

#[derive(Debug, Clone)]
pub struct JobStatusEntry {
    pub status: SendStatus,
    pub error: Option<String>,
    updated_at: Instant,
}

impl JobStatusStore {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub fn set_queued(&self, id: u64) {
        self.write(id, SendStatus::Queued, None);
    }

    pub fn set_sending(&self, id: u64) {
        self.write(id, SendStatus::Sending, None);
    }

    pub fn set_sent(&self, id: u64) {
        self.write(id, SendStatus::Sent, None);
    }

    pub fn set_failed(&self, id: u64, error: String) {
        self.write(id, SendStatus::Failed, Some(error));
    }

    pub fn get(&self, id: u64) -> Option<JobStatusEntry> {
        let guard = match self.inner.read() {
            Ok(g) => g,
            Err(poisoned) => {
                tracing::error!("JobStatusStore read lock poisoned, recovering");
                poisoned.into_inner()
            }
        };
        guard.get(&id).cloned()
    }

    pub fn remove(&self, id: u64) {
        let mut guard = match self.inner.write() {
            Ok(g) => g,
            Err(poisoned) => {
                tracing::error!("JobStatusStore write lock poisoned, recovering");
                poisoned.into_inner()
            }
        };
        guard.remove(&id);
    }

    pub fn all_done(&self, ids: &[u64]) -> bool {
        let guard = match self.inner.read() {
            Ok(g) => g,
            Err(_) => return false,
        };
        ids.iter().all(|id| {
            matches!(
                guard.get(id).map(|e| e.status),
                Some(SendStatus::Sent | SendStatus::Failed)
            )
        })
    }

    fn write(&self, id: u64, status: SendStatus, error: Option<String>) {
        let mut guard = match self.inner.write() {
            Ok(g) => g,
            Err(poisoned) => {
                tracing::error!("JobStatusStore write lock poisoned, recovering");
                poisoned.into_inner()
            }
        };
        prune_statuses(&mut guard);
        guard.insert(
            id,
            JobStatusEntry {
                status,
                error,
                updated_at: Instant::now(),
            },
        );
    }
}

fn prune_statuses(entries: &mut HashMap<u64, JobStatusEntry>) {
    let now = Instant::now();
    entries.retain(|_, entry| now.duration_since(entry.updated_at) <= STATUS_TTL);

    while entries.len() >= MAX_STATUS_ENTRIES {
        let Some(oldest_id) = entries
            .iter()
            .min_by_key(|(_, entry)| entry.updated_at)
            .map(|(id, _)| *id)
        else {
            break;
        };
        entries.remove(&oldest_id);
    }
}