lakehouse/crates/vectord/src/jobs.rs

/// Background job system for long-running embedding tasks.
/// POST /vectors/index returns a job_id immediately.
/// GET /vectors/jobs/{id} returns progress.
/// Embedding runs in background via tokio::spawn.

use serde::Serialize;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum JobType {
    Embed,
    ProfileActivation,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum JobStatus {
    Running,
    Completed,
    Failed,
}

#[derive(Debug, Clone, Serialize)]
pub struct Job {
    pub id: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub job_type: Option<JobType>,
    pub status: JobStatus,
    pub index_name: Option<String>,
    pub profile_id: Option<String>,
    pub total_chunks: usize,
    #[serde(rename = "processed")]
    pub processed_alias: usize,
    pub progress_pct: f32,
    pub storage_key: Option<String>,
    pub error: Option<String>,
    pub started_at: String,
    pub completed_at: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<serde_json::Value>,
}

impl Job {
    pub fn new_embed(index_name: &str, total_chunks: usize) -> Self {
        Self {
            id: format!("job-{}", chrono::Utc::now().timestamp_millis()),
            job_type: Some(JobType::Embed),
            status: JobStatus::Running,
            index_name: Some(index_name.to_string()),
            profile_id: None,
            total_chunks,
            processed_alias: 0,
            progress_pct: 0.0,
            storage_key: None,
            error: None,
            started_at: chrono::Utc::now().to_rfc3339(),
            completed_at: None,
            result: None,
        }
    }

    pub fn new_profile_activation(profile_id: &str) -> Self {
        Self {
            id: format!("job-{}", chrono::Utc::now().timestamp_millis()),
            job_type: Some(JobType::ProfileActivation),
            status: JobStatus::Running,
            index_name: None,
            profile_id: Some(profile_id.to_string()),
            total_chunks: 0,
            processed_alias: 0,
            progress_pct: 0.0,
            storage_key: None,
            error: None,
            started_at: chrono::Utc::now().to_rfc3339(),
            completed_at: None,
            result: None,
        }
    }
}

#[derive(Clone)]
pub struct JobTracker {
    jobs: Arc<RwLock<HashMap<String, Job>>>,
}

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

    pub async fn create_embed(&self, index_name: &str, total_chunks: usize) -> String {
        let job = Job::new_embed(index_name, total_chunks);
        let id = job.id.clone();
        self.jobs.write().await.insert(id.clone(), job);
        id
    }

    pub async fn create_profile_activation(&self, profile_id: &str) -> String {
        let job = Job::new_profile_activation(profile_id);
        let id = job.id.clone();
        self.jobs.write().await.insert(id.clone(), job);
        id
    }

    pub async fn update_embed_progress(&self, id: &str, embedded: usize, _rate: f32) {
        let mut jobs = self.jobs.write().await;
        if let Some(job) = jobs.get_mut(id) {
            job.processed_alias = embedded;
            job.progress_pct = if job.total_chunks > 0 {
                (embedded as f32 / job.total_chunks as f32) * 100.0
            } else {
                0.0
            };
        }
    }

    pub async fn complete(&self, id: &str, result: Option<serde_json::Value>) {
        let mut jobs = self.jobs.write().await;
        if let Some(job) = jobs.get_mut(id) {
            job.status = JobStatus::Completed;
            job.progress_pct = 100.0;
            job.completed_at = Some(chrono::Utc::now().to_rfc3339());
            job.result = result;
        }
    }

    pub async fn fail(&self, id: &str, error: String) {
        let mut jobs = self.jobs.write().await;
        if let Some(job) = jobs.get_mut(id) {
            job.status = JobStatus::Failed;
            job.error = Some(error);
            job.completed_at = Some(chrono::Utc::now().to_rfc3339());
        }
    }

    pub async fn get(&self, id: &str) -> Option<Job> {
        self.jobs.read().await.get(id).cloned()
    }

    pub async fn list(&self) -> Vec<Job> {
        self.jobs.read().await.values().cloned().collect()
    }
}