@ -13,6 +13,7 @@ use std::time::{Duration, Instant};
use chrono ::Utc ;
use dashmap ::DashMap ;
use dashmap ::DashSet ;
use serde ::Serialize ;
use tokio ::sync ::watch ;
use uuid ::Uuid ;
@ -82,6 +83,7 @@ struct JobEntry {
#[ derive(Clone) ]
pub struct JobStore {
inner : Arc < DashMap < Uuid , JobEntry > > ,
generating_users : Arc < DashSet < Uuid > > ,
}
/// Jobs expire after 1 hour (allows SSE reconnection).
@ -98,49 +100,28 @@ impl JobStore {
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 and the watch Sender.
///
/// 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 > > ) > {
// Check if user already has an active job
for entry in self . inner . iter ( ) {
if entry . value ( ) . user_id = = user_id {
// Check if the job is still active (not completed/failed)
let current = entry . value ( ) . tx . borrow ( ) . clone ( ) ;
match current {
ProgressEvent ::Complete { .. } | ProgressEvent ::Error { .. } = > {
// Job finished, allow creating a new one
continue ;
}
ProgressEvent ::Progress { .. } = > {
return None ; // Active job exists
}
}
}
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 ) ;
self . inner . insert (
job_id ,
JobEntry {
tx : Arc ::clone ( & tx ) ,
_rx : rx ,
user_id ,
created_at : Instant ::now ( ) ,
} ,
) ;
self . inner . insert ( job_id , JobEntry {
tx : Arc ::clone ( & tx ) , _rx : rx , user_id , created_at : Instant ::now ( ) ,
} ) ;
Some ( ( job_id , tx ) )
}
@ -157,22 +138,25 @@ impl JobStore {
/// 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 {
let current = entry . value ( ) . tx . borrow ( ) . clone ( ) ;
if matches! ( current , ProgressEvent ::Progress { .. } ) {
return Some ( * entry . key ( ) ) ;
}
}
if entry . value ( ) . user_id = = user_id { return Some ( * entry . key ( ) ) ; }
}
None
}
/// 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 | {
now . duration_since ( entry . created_at ) < JOB_TTL
let keep = now . duration_since ( entry . created_at ) < JOB_TTL ;
if ! keep { self . generating_users . remove ( & entry . user_id ) ; }
keep
} ) ;
}
@ -1371,11 +1355,12 @@ mod tests {
let ( _job_id , tx ) = store . create_job ( user_id ) . unwrap ( ) ;
// Complete the job
// 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 ) ;
@ -1389,11 +1374,12 @@ mod tests {
let ( _job_id , tx ) = store . create_job ( user_id ) . unwrap ( ) ;
// Fail the job
// 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 ) ;