1use mold_core::Config;
2use mold_inference::InferenceEngine;
3use std::path::PathBuf;
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::sync::{Arc, RwLock};
6use std::time::Instant;
7use tokio::sync::Mutex;
8
9use mold_inference::shared_pool::SharedPool;
10
11use crate::downloads::DownloadQueue;
12use crate::events::EventBroadcaster;
13use crate::gpu_pool::GpuPool;
14use crate::job_registry::{JobRegistry, SharedJobRegistry};
15use crate::model_cache::ModelCache;
16use crate::resources::ResourceBroadcaster;
17
18#[derive(Debug, Clone, Default)]
19pub struct EngineSnapshot {
20 pub model_name: Option<String>,
22 pub is_loaded: bool,
23 pub cached_models: Vec<String>,
25}
26
27#[derive(Debug, Clone)]
28pub struct ActiveGenerationSnapshot {
29 pub model: String,
30 pub prompt_sha256: String,
31 pub started_at_unix_ms: u64,
32 pub started_at: Instant,
33}
34
35pub enum SseMessage {
39 Progress(mold_core::SseProgressEvent),
40 Complete(mold_core::SseCompleteEvent),
41 UpscaleComplete(mold_core::SseUpscaleCompleteEvent),
42 Error(mold_core::SseErrorEvent),
43}
44
45pub struct GenerationJob {
47 pub id: String,
55 pub request: mold_core::GenerateRequest,
56 pub progress_tx: Option<tokio::sync::mpsc::UnboundedSender<SseMessage>>,
58 pub result_tx: tokio::sync::oneshot::Sender<Result<GenerationJobResult, String>>,
60 pub output_dir: Option<PathBuf>,
62}
63
64pub struct GenerationJobResult {
65 pub response: mold_core::GenerateResponse,
66 pub image: mold_core::ImageData,
67}
68
69#[derive(Clone)]
71pub struct QueueHandle {
72 job_tx: tokio::sync::mpsc::Sender<GenerationJob>,
73 pending_count: Arc<AtomicUsize>,
74}
75
76#[derive(Debug)]
78pub enum SubmitError {
79 Full { pending: usize, capacity: usize },
81 Shutdown,
83}
84
85impl QueueHandle {
86 pub fn new(job_tx: tokio::sync::mpsc::Sender<GenerationJob>) -> Self {
87 Self {
88 job_tx,
89 pending_count: Arc::new(AtomicUsize::new(0)),
90 }
91 }
92
93 pub async fn submit(&self, job: GenerationJob, capacity: usize) -> Result<usize, SubmitError> {
99 let prev = self.pending_count.fetch_add(1, Ordering::SeqCst);
100 if prev >= capacity {
101 self.pending_count.fetch_sub(1, Ordering::SeqCst);
102 return Err(SubmitError::Full {
103 pending: prev,
104 capacity,
105 });
106 }
107 if self.job_tx.send(job).await.is_err() {
108 self.pending_count.fetch_sub(1, Ordering::SeqCst);
109 return Err(SubmitError::Shutdown);
110 }
111 #[cfg(feature = "metrics")]
112 {
113 crate::metrics::record_queue_submit();
114 crate::metrics::record_queue_depth(self.pending_count.load(Ordering::SeqCst));
115 }
116 Ok(prev)
117 }
118
119 pub fn decrement(&self) {
120 self.pending_count.fetch_sub(1, Ordering::SeqCst);
121 }
122
123 pub fn pending(&self) -> usize {
124 self.pending_count.load(Ordering::SeqCst)
125 }
126}
127
128#[derive(Clone)]
131pub struct AppState {
132 pub gpu_pool: Arc<GpuPool>,
135 pub queue_capacity: usize,
137
138 pub model_cache: Arc<Mutex<ModelCache>>,
140 pub active_generation: Arc<RwLock<Option<ActiveGenerationSnapshot>>>,
144 pub config: Arc<tokio::sync::RwLock<Config>>,
145 pub start_time: Instant,
146 pub model_load_lock: Arc<Mutex<()>>,
148 pub pull_lock: Arc<Mutex<()>>,
150 pub queue: QueueHandle,
152 pub job_registry: SharedJobRegistry,
157 pub queue_pause: Arc<crate::queue::QueuePause>,
161 pub shared_pool: Arc<std::sync::Mutex<SharedPool>>,
163 pub shutdown_tx: Arc<tokio::sync::Mutex<Option<tokio::sync::oneshot::Sender<()>>>>,
165 pub upscaler_cache: Arc<std::sync::Mutex<Option<Box<dyn mold_inference::UpscaleEngine>>>>,
167 pub metadata_db: Arc<Option<mold_db::MetadataDb>>,
171 pub chain_jobs: Option<Arc<crate::chain_job_runner::ChainJobRunnerHandle>>,
174 pub downloads: Arc<DownloadQueue>,
177 pub resources: Arc<ResourceBroadcaster>,
179 pub events: Arc<EventBroadcaster>,
183 pub catalog_live_cache: mold_catalog::live::LiveCache,
186 pub catalog_live_civitai_base: Arc<String>,
189 pub catalog_intents: Arc<
194 tokio::sync::RwLock<
195 std::collections::HashMap<String, mold_catalog::synthesis::CatalogModelIntent>,
196 >,
197 >,
198}
199
200const CATALOG_LIVE_CACHE_TTL_SECS: u64 = 300;
204const CATALOG_LIVE_CACHE_MAX_KEYS: usize = 256;
207pub const CATALOG_LIVE_CIVITAI_BASE: &str = "https://civitai.com";
210
211fn default_live_cache() -> mold_catalog::live::LiveCache {
212 mold_catalog::live::LiveCache::new(
213 std::time::Duration::from_secs(CATALOG_LIVE_CACHE_TTL_SECS),
214 CATALOG_LIVE_CACHE_MAX_KEYS,
215 )
216}
217
218pub const DEFAULT_MAX_CACHED_MODELS: usize = 3;
220const MAX_CACHED_MODELS_LOWER: usize = 1;
224const MAX_CACHED_MODELS_UPPER: usize = 16;
225pub const MAX_CACHED_MODELS_ENV: &str = "MOLD_MAX_CACHED_MODELS";
227
228pub const DEFAULT_CACHE_IDLE_TTL_SECS: u64 = 1800;
233const CACHE_IDLE_TTL_LOWER_SECS: u64 = 60;
234const CACHE_IDLE_TTL_UPPER_SECS: u64 = 86_400;
235pub const CACHE_IDLE_TTL_ENV: &str = "MOLD_CACHE_IDLE_TTL_SECS";
237
238pub fn resolve_cache_idle_ttl_secs() -> u64 {
241 match std::env::var(CACHE_IDLE_TTL_ENV) {
242 Ok(raw) => match raw.trim().parse::<u64>() {
243 Ok(n) if (CACHE_IDLE_TTL_LOWER_SECS..=CACHE_IDLE_TTL_UPPER_SECS).contains(&n) => n,
244 Ok(n) => {
245 tracing::warn!(
246 env = CACHE_IDLE_TTL_ENV,
247 value = n,
248 lower = CACHE_IDLE_TTL_LOWER_SECS,
249 upper = CACHE_IDLE_TTL_UPPER_SECS,
250 "ignoring out-of-range cache idle-TTL; using default"
251 );
252 DEFAULT_CACHE_IDLE_TTL_SECS
253 }
254 Err(e) => {
255 tracing::warn!(
256 env = CACHE_IDLE_TTL_ENV,
257 raw = %raw,
258 error = %e,
259 "ignoring unparseable cache idle-TTL; using default"
260 );
261 DEFAULT_CACHE_IDLE_TTL_SECS
262 }
263 },
264 Err(_) => DEFAULT_CACHE_IDLE_TTL_SECS,
265 }
266}
267
268pub fn resolve_max_cached_models() -> usize {
272 match std::env::var(MAX_CACHED_MODELS_ENV) {
273 Ok(raw) => match raw.trim().parse::<usize>() {
274 Ok(n) if (MAX_CACHED_MODELS_LOWER..=MAX_CACHED_MODELS_UPPER).contains(&n) => n,
275 Ok(n) => {
276 tracing::warn!(
277 env = MAX_CACHED_MODELS_ENV,
278 value = n,
279 lower = MAX_CACHED_MODELS_LOWER,
280 upper = MAX_CACHED_MODELS_UPPER,
281 "ignoring out-of-range cache cap; using default"
282 );
283 DEFAULT_MAX_CACHED_MODELS
284 }
285 Err(e) => {
286 tracing::warn!(
287 env = MAX_CACHED_MODELS_ENV,
288 raw = %raw,
289 error = %e,
290 "ignoring unparseable cache cap; using default"
291 );
292 DEFAULT_MAX_CACHED_MODELS
293 }
294 },
295 Err(_) => DEFAULT_MAX_CACHED_MODELS,
296 }
297}
298
299impl AppState {
300 pub fn new(
302 engine: Box<dyn InferenceEngine>,
303 config: Config,
304 queue: QueueHandle,
305 gpu_pool: Arc<GpuPool>,
306 queue_capacity: usize,
307 ) -> Self {
308 let mut cache = ModelCache::new(resolve_max_cached_models());
309 cache.insert(engine, 0);
310 let events = EventBroadcaster::new();
311 Self {
312 gpu_pool,
313 queue_capacity,
314 model_cache: Arc::new(Mutex::new(cache)),
315 active_generation: Arc::new(RwLock::new(None)),
316 config: Arc::new(tokio::sync::RwLock::new(config)),
317 start_time: Instant::now(),
318 model_load_lock: Arc::new(Mutex::new(())),
319 pull_lock: Arc::new(Mutex::new(())),
320 queue,
321 job_registry: JobRegistry::with_events(events.clone()),
322 queue_pause: crate::queue::QueuePause::new(),
323 shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
324 shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
325 upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
326 metadata_db: Arc::new(None),
327 chain_jobs: None,
328 downloads: DownloadQueue::new(),
329 resources: ResourceBroadcaster::new(),
330 events,
331 catalog_live_cache: default_live_cache(),
332 catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
333 catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
334 }
335 }
336
337 pub fn empty(
339 config: Config,
340 queue: QueueHandle,
341 gpu_pool: Arc<GpuPool>,
342 queue_capacity: usize,
343 ) -> Self {
344 let events = EventBroadcaster::new();
345 Self {
346 gpu_pool,
347 queue_capacity,
348 model_cache: Arc::new(Mutex::new(ModelCache::new(resolve_max_cached_models()))),
349 active_generation: Arc::new(RwLock::new(None)),
350 config: Arc::new(tokio::sync::RwLock::new(config)),
351 start_time: Instant::now(),
352 model_load_lock: Arc::new(Mutex::new(())),
353 pull_lock: Arc::new(Mutex::new(())),
354 queue,
355 job_registry: JobRegistry::with_events(events.clone()),
356 queue_pause: crate::queue::QueuePause::new(),
357 shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
358 shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
359 upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
360 metadata_db: Arc::new(None),
361 chain_jobs: None,
362 downloads: DownloadQueue::new(),
363 resources: ResourceBroadcaster::new(),
364 events,
365 catalog_live_cache: default_live_cache(),
366 catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
367 catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
368 }
369 }
370
371 #[cfg(test)]
373 pub(crate) fn empty_gpu_pool() -> Arc<GpuPool> {
374 Arc::new(GpuPool {
375 workers: Vec::new(),
376 })
377 }
378
379 #[cfg(test)]
383 pub(crate) fn empty_gpu_pool_for_test() -> Arc<GpuPool> {
384 Self::empty_gpu_pool()
385 }
386
387 #[cfg(test)]
388 pub fn with_engine(engine: impl InferenceEngine + 'static) -> Self {
389 let (tx, _rx) = tokio::sync::mpsc::channel(16);
390 let queue = QueueHandle::new(tx);
391 let mut cache = ModelCache::new(resolve_max_cached_models());
392 cache.insert(Box::new(engine), 0);
393 let events = EventBroadcaster::new();
394 Self {
395 gpu_pool: Self::empty_gpu_pool(),
396 queue_capacity: 200,
397 model_cache: Arc::new(Mutex::new(cache)),
398 active_generation: Arc::new(RwLock::new(None)),
399 config: Arc::new(tokio::sync::RwLock::new(Config::default())),
400 start_time: Instant::now(),
401 model_load_lock: Arc::new(Mutex::new(())),
402 pull_lock: Arc::new(Mutex::new(())),
403 queue,
404 job_registry: JobRegistry::with_events(events.clone()),
405 queue_pause: crate::queue::QueuePause::new(),
406 shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
407 shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
408 upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
409 metadata_db: Arc::new(None),
410 chain_jobs: None,
411 downloads: DownloadQueue::new(),
412 resources: ResourceBroadcaster::new(),
413 events,
414 catalog_live_cache: default_live_cache(),
415 catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
416 catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
417 }
418 }
419
420 #[cfg(test)]
422 pub fn with_engine_and_queue(
423 engine: impl InferenceEngine + 'static,
424 ) -> (Self, tokio::sync::mpsc::Receiver<GenerationJob>) {
425 let (tx, rx) = tokio::sync::mpsc::channel(16);
426 let queue = QueueHandle::new(tx);
427 let mut cache = ModelCache::new(resolve_max_cached_models());
428 cache.insert(Box::new(engine), 0);
429 let events = EventBroadcaster::new();
430 let state = Self {
431 gpu_pool: Self::empty_gpu_pool(),
432 queue_capacity: 200,
433 model_cache: Arc::new(Mutex::new(cache)),
434 active_generation: Arc::new(RwLock::new(None)),
435 config: Arc::new(tokio::sync::RwLock::new(Config::default())),
436 start_time: Instant::now(),
437 model_load_lock: Arc::new(Mutex::new(())),
438 pull_lock: Arc::new(Mutex::new(())),
439 queue,
440 job_registry: JobRegistry::with_events(events.clone()),
441 queue_pause: crate::queue::QueuePause::new(),
442 shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
443 shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
444 upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
445 metadata_db: Arc::new(None),
446 chain_jobs: None,
447 downloads: DownloadQueue::new(),
448 resources: ResourceBroadcaster::new(),
449 events,
450 catalog_live_cache: default_live_cache(),
451 catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
452 catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
453 };
454 (state, rx)
455 }
456
457 pub fn for_tests() -> Self {
461 let (tx, _rx) = tokio::sync::mpsc::channel(16);
462 let queue = QueueHandle::new(tx);
463 let events = EventBroadcaster::new();
464 Self {
465 gpu_pool: Arc::new(GpuPool {
466 workers: Vec::new(),
467 }),
468 queue_capacity: 200,
469 model_cache: Arc::new(Mutex::new(ModelCache::new(resolve_max_cached_models()))),
470 active_generation: Arc::new(RwLock::new(None)),
471 config: Arc::new(tokio::sync::RwLock::new(Config::default())),
472 start_time: Instant::now(),
473 model_load_lock: Arc::new(Mutex::new(())),
474 pull_lock: Arc::new(Mutex::new(())),
475 queue,
476 job_registry: JobRegistry::with_events(events.clone()),
477 queue_pause: crate::queue::QueuePause::new(),
478 shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
479 shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
480 upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
481 metadata_db: Arc::new(None),
482 chain_jobs: None,
483 downloads: DownloadQueue::new(),
484 resources: ResourceBroadcaster::new(),
485 events,
486 catalog_live_cache: default_live_cache(),
487 catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
488 catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
489 }
490 }
491
492 pub fn with_civitai_base(mut self, base: impl Into<String>) -> Self {
496 self.catalog_live_civitai_base = Arc::new(base.into());
497 self
498 }
499}
500
501#[cfg(test)]
502mod tests {
503 use super::*;
504
505 #[test]
506 fn engine_snapshot_default_is_unloaded() {
507 let snap = EngineSnapshot::default();
508 assert!(snap.model_name.is_none());
509 assert!(!snap.is_loaded);
510 }
511
512 #[test]
513 fn active_generation_snapshot_stores_fields() {
514 let snap = ActiveGenerationSnapshot {
515 model: "flux-dev:q8".to_string(),
516 prompt_sha256: "abc123".to_string(),
517 started_at_unix_ms: 1700000000000,
518 started_at: std::time::Instant::now(),
519 };
520 assert_eq!(snap.model, "flux-dev:q8");
521 assert_eq!(snap.prompt_sha256, "abc123");
522 assert_eq!(snap.started_at_unix_ms, 1700000000000);
523 }
524
525 #[test]
526 fn queue_handle_pending_starts_at_zero() {
527 let (tx, _rx) = tokio::sync::mpsc::channel::<GenerationJob>(16);
528 let handle = QueueHandle::new(tx);
529 assert_eq!(handle.pending(), 0);
530 }
531
532 #[test]
533 fn upscaler_cache_starts_empty() {
534 let config = mold_core::Config::default();
535 let state = AppState::empty(
536 config,
537 QueueHandle::new(tokio::sync::mpsc::channel(1).0),
538 AppState::empty_gpu_pool(),
539 200,
540 );
541 let cache = state.upscaler_cache.lock().unwrap();
542 assert!(cache.is_none());
543 }
544
545 #[test]
546 fn upscaler_cache_cleared_by_setting_none() {
547 let config = mold_core::Config::default();
548 let state = AppState::empty(
549 config,
550 QueueHandle::new(tokio::sync::mpsc::channel(1).0),
551 AppState::empty_gpu_pool(),
552 200,
553 );
554 {
555 let mut cache = state.upscaler_cache.lock().unwrap();
556 *cache = None;
557 }
558 let cache = state.upscaler_cache.lock().unwrap();
559 assert!(cache.is_none());
560 }
561
562 #[test]
563 fn app_state_exposes_resources_broadcaster() {
564 let config = mold_core::Config::default();
565 let state = AppState::empty(
566 config,
567 QueueHandle::new(tokio::sync::mpsc::channel(1).0),
568 AppState::empty_gpu_pool(),
569 200,
570 );
571 assert!(state.resources.latest().is_none());
573 let _rx = state.resources.subscribe();
575 }
576
577 static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
581
582 fn with_env<R>(name: &str, value: Option<&str>, f: impl FnOnce() -> R) -> R {
586 let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
587 let prev = std::env::var(name).ok();
588 match value {
589 Some(v) => std::env::set_var(name, v),
590 None => std::env::remove_var(name),
591 }
592 let out = f();
593 match prev {
594 Some(v) => std::env::set_var(name, v),
595 None => std::env::remove_var(name),
596 }
597 out
598 }
599
600 #[test]
601 fn resolve_max_cached_uses_default_when_env_missing() {
602 let n = with_env(MAX_CACHED_MODELS_ENV, None, resolve_max_cached_models);
603 assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
604 }
605
606 #[test]
607 fn resolve_max_cached_honors_env_within_range() {
608 let n = with_env(MAX_CACHED_MODELS_ENV, Some("8"), resolve_max_cached_models);
609 assert_eq!(n, 8);
610 }
611
612 #[test]
613 fn resolve_max_cached_clamps_zero_back_to_default() {
614 let n = with_env(MAX_CACHED_MODELS_ENV, Some("0"), resolve_max_cached_models);
615 assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
616 }
617
618 #[test]
619 fn resolve_max_cached_clamps_overflow_back_to_default() {
620 let n = with_env(
621 MAX_CACHED_MODELS_ENV,
622 Some("999"),
623 resolve_max_cached_models,
624 );
625 assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
626 }
627
628 #[test]
629 fn resolve_max_cached_falls_back_when_env_unparseable() {
630 let n = with_env(
631 MAX_CACHED_MODELS_ENV,
632 Some("not-a-number"),
633 resolve_max_cached_models,
634 );
635 assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
636 }
637
638 #[test]
639 fn resolve_cache_idle_ttl_uses_default_when_env_missing() {
640 let n = with_env(CACHE_IDLE_TTL_ENV, None, resolve_cache_idle_ttl_secs);
641 assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
642 }
643
644 #[test]
645 fn resolve_cache_idle_ttl_honors_env_within_range() {
646 let n = with_env(CACHE_IDLE_TTL_ENV, Some("600"), resolve_cache_idle_ttl_secs);
648 assert_eq!(n, 600);
649 }
650
651 #[test]
652 fn resolve_cache_idle_ttl_clamps_zero_back_to_default() {
653 let n = with_env(CACHE_IDLE_TTL_ENV, Some("0"), resolve_cache_idle_ttl_secs);
655 assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
656 }
657
658 #[test]
659 fn resolve_cache_idle_ttl_clamps_overflow_back_to_default() {
660 let n = with_env(
662 CACHE_IDLE_TTL_ENV,
663 Some("100000"),
664 resolve_cache_idle_ttl_secs,
665 );
666 assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
667 }
668
669 #[test]
670 fn resolve_cache_idle_ttl_falls_back_when_env_unparseable() {
671 let n = with_env(
672 CACHE_IDLE_TTL_ENV,
673 Some("not-a-number"),
674 resolve_cache_idle_ttl_secs,
675 );
676 assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
677 }
678}