refactor: extract JobStore to services/job_store.rs
Moves JobEntry, JobStore, ProgressEvent, JOB_TTL, and emit_progress to a dedicated module. Updates imports in synthesis.rs, generation.rs, scheduler.rs, and app_state.rs. synthesis.rs re-exports for backward compatibility. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>master
parent
1ab9c817e4
commit
7f647bc656
@ -0,0 +1,335 @@
|
||||
//! In-memory job store for generation progress tracking.
|
||||
//!
|
||||
//! Manages active generation jobs keyed by UUID, with per-user
|
||||
//! deduplication and TTL-based cleanup. Progress is communicated
|
||||
//! via `tokio::sync::watch` channels.
|
||||
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use dashmap::DashMap;
|
||||
use dashmap::DashSet;
|
||||
use serde::Serialize;
|
||||
use tokio::sync::watch;
|
||||
use uuid::Uuid;
|
||||
|
||||
// ───────────────────────────────────────────────────────────────────
|
||||
// Progress Events
|
||||
// ───────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Progress event sent to SSE clients during generation.
|
||||
///
|
||||
/// The `watch` channel always holds the latest event, and new subscribers
|
||||
/// immediately receive the current state.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum ProgressEvent {
|
||||
/// Generation is in progress.
|
||||
#[serde(rename = "progress")]
|
||||
Progress {
|
||||
step: String,
|
||||
message: String,
|
||||
percent: u8,
|
||||
},
|
||||
/// Generation completed successfully.
|
||||
#[serde(rename = "complete")]
|
||||
Complete { synthesis_id: Uuid },
|
||||
/// Generation failed with an error.
|
||||
#[serde(rename = "error")]
|
||||
Error { message: String },
|
||||
}
|
||||
|
||||
// ───────────────────────────────────────────────────────────────────
|
||||
// Job Store
|
||||
// ───────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Entry in the job store, holding the progress channel and metadata.
|
||||
struct JobEntry {
|
||||
/// Sender side of the watch channel for progress updates.
|
||||
/// Wrapped in Arc so it can be shared with the background task
|
||||
/// without cloning the Sender itself.
|
||||
tx: Arc<watch::Sender<ProgressEvent>>,
|
||||
/// A receiver kept alive to prevent the channel from closing.
|
||||
/// Without at least one receiver, `Sender::send()` returns an error
|
||||
/// and does NOT update the stored value.
|
||||
_rx: watch::Receiver<ProgressEvent>,
|
||||
/// User who owns this job.
|
||||
user_id: Uuid,
|
||||
/// When the job was created (for TTL cleanup).
|
||||
created_at: Instant,
|
||||
/// Flag set to true when the user requests cancellation.
|
||||
cancelled: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
/// In-memory store for active generation jobs.
|
||||
///
|
||||
/// Uses `DashMap` for lock-free concurrent access. Jobs are keyed by
|
||||
/// a random UUID and automatically cleaned up after a TTL.
|
||||
#[derive(Clone)]
|
||||
pub struct JobStore {
|
||||
inner: Arc<DashMap<Uuid, JobEntry>>,
|
||||
generating_users: Arc<DashSet<Uuid>>,
|
||||
}
|
||||
|
||||
/// Jobs expire after 1 hour (allows SSE reconnection).
|
||||
const JOB_TTL: Duration = Duration::from_secs(3600);
|
||||
|
||||
impl Default for JobStore {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl JobStore {
|
||||
/// Create a new empty job store.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
inner: Arc::new(DashMap::new()),
|
||||
generating_users: Arc::new(DashSet::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new job for a user, returning the job ID, the watch Sender, and a cancellation flag.
|
||||
///
|
||||
/// Returns `None` if the user already has an active job.
|
||||
/// Uses an atomic DashSet insert to prevent race conditions on double-click.
|
||||
pub fn create_job(&self, user_id: Uuid) -> Option<(Uuid, Arc<watch::Sender<ProgressEvent>>, Arc<AtomicBool>)> {
|
||||
if !self.generating_users.insert(user_id) {
|
||||
return None;
|
||||
}
|
||||
let job_id = Uuid::new_v4();
|
||||
let (tx, rx) = watch::channel(ProgressEvent::Progress {
|
||||
step: "init".into(),
|
||||
message: "Initialisation...".into(),
|
||||
percent: 0,
|
||||
});
|
||||
let tx = Arc::new(tx);
|
||||
let cancelled = Arc::new(AtomicBool::new(false));
|
||||
self.inner.insert(job_id, JobEntry {
|
||||
tx: Arc::clone(&tx), _rx: rx, user_id, created_at: Instant::now(),
|
||||
cancelled: Arc::clone(&cancelled),
|
||||
});
|
||||
Some((job_id, tx, cancelled))
|
||||
}
|
||||
|
||||
/// Get a watch receiver for a job, if it exists and belongs to the given user.
|
||||
pub fn subscribe(&self, job_id: Uuid, user_id: Uuid) -> Option<watch::Receiver<ProgressEvent>> {
|
||||
self.inner.get(&job_id).and_then(|entry| {
|
||||
if entry.value().user_id == user_id {
|
||||
Some(entry.value().tx.subscribe())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Check if a user has an active (in-progress) job.
|
||||
pub fn has_active_job(&self, user_id: Uuid) -> Option<Uuid> {
|
||||
if !self.generating_users.contains(&user_id) { return None; }
|
||||
for entry in self.inner.iter() {
|
||||
if entry.value().user_id == user_id { return Some(*entry.key()); }
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Signal a job to stop. Returns true if the job was found and belongs to the user.
|
||||
pub fn cancel_job(&self, job_id: Uuid, user_id: Uuid) -> bool {
|
||||
if let Some(entry) = self.inner.get(&job_id) {
|
||||
if entry.value().user_id == user_id {
|
||||
entry.value().cancelled.store(true, Ordering::Relaxed);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// Release the generating lock for a user (called when job completes, errors, or times out).
|
||||
pub fn release_user(&self, user_id: Uuid) {
|
||||
self.generating_users.remove(&user_id);
|
||||
}
|
||||
|
||||
/// Remove expired jobs (older than TTL).
|
||||
pub fn cleanup_expired(&self) {
|
||||
let now = Instant::now();
|
||||
self.inner.retain(|_, entry| {
|
||||
let keep = now.duration_since(entry.created_at) < JOB_TTL;
|
||||
if !keep { self.generating_users.remove(&entry.user_id); }
|
||||
keep
|
||||
});
|
||||
}
|
||||
|
||||
/// Remove a specific job.
|
||||
pub fn remove(&self, job_id: &Uuid) {
|
||||
self.inner.remove(job_id);
|
||||
}
|
||||
|
||||
/// Get the number of active jobs (for testing/monitoring).
|
||||
pub fn len(&self) -> usize {
|
||||
self.inner.len()
|
||||
}
|
||||
|
||||
/// Check if the store is empty (for testing).
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.inner.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
/// Emit a progress event via the watch channel.
|
||||
pub fn emit_progress(tx: &watch::Sender<ProgressEvent>, step: &str, message: &str, percent: u8) {
|
||||
tx.send(ProgressEvent::Progress {
|
||||
step: step.into(),
|
||||
message: message.into(),
|
||||
percent,
|
||||
})
|
||||
.ok();
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn job_store_create_and_subscribe() {
|
||||
let store = JobStore::new();
|
||||
let user_id = Uuid::new_v4();
|
||||
|
||||
let (job_id, tx, _cancelled) = store.create_job(user_id).unwrap();
|
||||
assert_eq!(store.len(), 1);
|
||||
|
||||
// Subscribe
|
||||
let rx = store.subscribe(job_id, user_id);
|
||||
assert!(rx.is_some());
|
||||
|
||||
// Wrong user cannot subscribe
|
||||
let other_user = Uuid::new_v4();
|
||||
assert!(store.subscribe(job_id, other_user).is_none());
|
||||
|
||||
// Check active job
|
||||
assert_eq!(store.has_active_job(user_id), Some(job_id));
|
||||
assert_eq!(store.has_active_job(other_user), None);
|
||||
|
||||
drop(tx);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn job_store_prevents_duplicate_active_jobs() {
|
||||
let store = JobStore::new();
|
||||
let user_id = Uuid::new_v4();
|
||||
|
||||
let _result1 = store.create_job(user_id);
|
||||
assert!(_result1.is_some());
|
||||
|
||||
// Second job for same user should fail
|
||||
let result2 = store.create_job(user_id);
|
||||
assert!(result2.is_none());
|
||||
|
||||
// Different user should succeed
|
||||
let other_user = Uuid::new_v4();
|
||||
let result3 = store.create_job(other_user);
|
||||
assert!(result3.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn job_store_allows_new_job_after_completion() {
|
||||
let store = JobStore::new();
|
||||
let user_id = Uuid::new_v4();
|
||||
|
||||
let (_job_id, tx, _cancelled) = store.create_job(user_id).unwrap();
|
||||
|
||||
// Complete the job and release the user lock (as the pipeline does)
|
||||
tx.send(ProgressEvent::Complete {
|
||||
synthesis_id: Uuid::new_v4(),
|
||||
})
|
||||
.ok();
|
||||
store.release_user(user_id);
|
||||
|
||||
// Should now allow a new job
|
||||
let result2 = store.create_job(user_id);
|
||||
assert!(result2.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn job_store_allows_new_job_after_error() {
|
||||
let store = JobStore::new();
|
||||
let user_id = Uuid::new_v4();
|
||||
|
||||
let (_job_id, tx, _cancelled) = store.create_job(user_id).unwrap();
|
||||
|
||||
// Fail the job and release the user lock (as the pipeline does)
|
||||
tx.send(ProgressEvent::Error {
|
||||
message: "test error".into(),
|
||||
})
|
||||
.ok();
|
||||
store.release_user(user_id);
|
||||
|
||||
// Should now allow a new job
|
||||
let result2 = store.create_job(user_id);
|
||||
assert!(result2.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn job_store_cleanup_expired() {
|
||||
let store = JobStore::new();
|
||||
let user_id = Uuid::new_v4();
|
||||
|
||||
// Create a job and manually set its created_at to the past
|
||||
let (_job_id, _tx, _cancelled) = store.create_job(user_id).unwrap();
|
||||
assert_eq!(store.len(), 1);
|
||||
|
||||
// Cleanup should not remove recent jobs
|
||||
store.cleanup_expired();
|
||||
assert_eq!(store.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn job_store_remove() {
|
||||
let store = JobStore::new();
|
||||
let user_id = Uuid::new_v4();
|
||||
|
||||
let (job_id, _tx, _cancelled) = store.create_job(user_id).unwrap();
|
||||
assert_eq!(store.len(), 1);
|
||||
|
||||
store.remove(&job_id);
|
||||
assert!(store.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn progress_event_serialization_progress() {
|
||||
let event = ProgressEvent::Progress {
|
||||
step: "search".into(),
|
||||
message: "Searching...".into(),
|
||||
percent: 30,
|
||||
};
|
||||
|
||||
let json = serde_json::to_value(&event).unwrap();
|
||||
assert_eq!(json["type"], "progress");
|
||||
assert_eq!(json["step"], "search");
|
||||
assert_eq!(json["message"], "Searching...");
|
||||
assert_eq!(json["percent"], 30);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn progress_event_serialization_complete() {
|
||||
let synthesis_id = Uuid::nil();
|
||||
let event = ProgressEvent::Complete { synthesis_id };
|
||||
|
||||
let json = serde_json::to_value(&event).unwrap();
|
||||
assert_eq!(json["type"], "complete");
|
||||
assert_eq!(
|
||||
json["synthesis_id"],
|
||||
"00000000-0000-0000-0000-000000000000"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn progress_event_serialization_error() {
|
||||
let event = ProgressEvent::Error {
|
||||
message: "Something went wrong".into(),
|
||||
};
|
||||
|
||||
let json = serde_json::to_value(&event).unwrap();
|
||||
assert_eq!(json["type"], "error");
|
||||
assert_eq!(json["message"], "Something went wrong");
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue