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