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 queue: QueueHandle,
151 pub job_registry: SharedJobRegistry,
156 pub shared_pool: Arc<std::sync::Mutex<SharedPool>>,
158 pub shutdown_tx: Arc<tokio::sync::Mutex<Option<tokio::sync::oneshot::Sender<()>>>>,
160 pub upscaler_cache: Arc<std::sync::Mutex<Option<Box<dyn mold_inference::UpscaleEngine>>>>,
162 pub metadata_db: Arc<Option<mold_db::MetadataDb>>,
166 pub chain_jobs: Option<Arc<crate::chain_job_runner::ChainJobRunnerHandle>>,
169 pub downloads: Arc<DownloadQueue>,
172 pub resources: Arc<ResourceBroadcaster>,
174 pub catalog_live_cache: mold_catalog::live::LiveCache,
177 pub catalog_live_civitai_base: Arc<String>,
180 pub catalog_intents: Arc<
185 tokio::sync::RwLock<
186 std::collections::HashMap<String, mold_catalog::synthesis::CatalogModelIntent>,
187 >,
188 >,
189}
190
191const CATALOG_LIVE_CACHE_TTL_SECS: u64 = 300;
195const CATALOG_LIVE_CACHE_MAX_KEYS: usize = 256;
198pub const CATALOG_LIVE_CIVITAI_BASE: &str = "https://civitai.com";
201
202fn default_live_cache() -> mold_catalog::live::LiveCache {
203 mold_catalog::live::LiveCache::new(
204 std::time::Duration::from_secs(CATALOG_LIVE_CACHE_TTL_SECS),
205 CATALOG_LIVE_CACHE_MAX_KEYS,
206 )
207}
208
209pub const DEFAULT_MAX_CACHED_MODELS: usize = 3;
211const MAX_CACHED_MODELS_LOWER: usize = 1;
215const MAX_CACHED_MODELS_UPPER: usize = 16;
216pub const MAX_CACHED_MODELS_ENV: &str = "MOLD_MAX_CACHED_MODELS";
218
219pub const DEFAULT_CACHE_IDLE_TTL_SECS: u64 = 1800;
224const CACHE_IDLE_TTL_LOWER_SECS: u64 = 60;
225const CACHE_IDLE_TTL_UPPER_SECS: u64 = 86_400;
226pub const CACHE_IDLE_TTL_ENV: &str = "MOLD_CACHE_IDLE_TTL_SECS";
228
229pub fn resolve_cache_idle_ttl_secs() -> u64 {
232 match std::env::var(CACHE_IDLE_TTL_ENV) {
233 Ok(raw) => match raw.trim().parse::<u64>() {
234 Ok(n) if (CACHE_IDLE_TTL_LOWER_SECS..=CACHE_IDLE_TTL_UPPER_SECS).contains(&n) => n,
235 Ok(n) => {
236 tracing::warn!(
237 env = CACHE_IDLE_TTL_ENV,
238 value = n,
239 lower = CACHE_IDLE_TTL_LOWER_SECS,
240 upper = CACHE_IDLE_TTL_UPPER_SECS,
241 "ignoring out-of-range cache idle-TTL; using default"
242 );
243 DEFAULT_CACHE_IDLE_TTL_SECS
244 }
245 Err(e) => {
246 tracing::warn!(
247 env = CACHE_IDLE_TTL_ENV,
248 raw = %raw,
249 error = %e,
250 "ignoring unparseable cache idle-TTL; using default"
251 );
252 DEFAULT_CACHE_IDLE_TTL_SECS
253 }
254 },
255 Err(_) => DEFAULT_CACHE_IDLE_TTL_SECS,
256 }
257}
258
259pub fn resolve_max_cached_models() -> usize {
263 match std::env::var(MAX_CACHED_MODELS_ENV) {
264 Ok(raw) => match raw.trim().parse::<usize>() {
265 Ok(n) if (MAX_CACHED_MODELS_LOWER..=MAX_CACHED_MODELS_UPPER).contains(&n) => n,
266 Ok(n) => {
267 tracing::warn!(
268 env = MAX_CACHED_MODELS_ENV,
269 value = n,
270 lower = MAX_CACHED_MODELS_LOWER,
271 upper = MAX_CACHED_MODELS_UPPER,
272 "ignoring out-of-range cache cap; using default"
273 );
274 DEFAULT_MAX_CACHED_MODELS
275 }
276 Err(e) => {
277 tracing::warn!(
278 env = MAX_CACHED_MODELS_ENV,
279 raw = %raw,
280 error = %e,
281 "ignoring unparseable cache cap; using default"
282 );
283 DEFAULT_MAX_CACHED_MODELS
284 }
285 },
286 Err(_) => DEFAULT_MAX_CACHED_MODELS,
287 }
288}
289
290impl AppState {
291 pub fn new(
293 engine: Box<dyn InferenceEngine>,
294 config: Config,
295 queue: QueueHandle,
296 gpu_pool: Arc<GpuPool>,
297 queue_capacity: usize,
298 ) -> Self {
299 let mut cache = ModelCache::new(resolve_max_cached_models());
300 cache.insert(engine, 0);
301 Self {
302 gpu_pool,
303 queue_capacity,
304 model_cache: Arc::new(Mutex::new(cache)),
305 active_generation: Arc::new(RwLock::new(None)),
306 config: Arc::new(tokio::sync::RwLock::new(config)),
307 start_time: Instant::now(),
308 model_load_lock: Arc::new(Mutex::new(())),
309 pull_lock: Arc::new(Mutex::new(())),
310 queue,
311 job_registry: JobRegistry::new(),
312 shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
313 shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
314 upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
315 metadata_db: Arc::new(None),
316 chain_jobs: None,
317 downloads: DownloadQueue::new(),
318 resources: ResourceBroadcaster::new(),
319 catalog_live_cache: default_live_cache(),
320 catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
321 catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
322 }
323 }
324
325 pub fn empty(
327 config: Config,
328 queue: QueueHandle,
329 gpu_pool: Arc<GpuPool>,
330 queue_capacity: usize,
331 ) -> Self {
332 Self {
333 gpu_pool,
334 queue_capacity,
335 model_cache: Arc::new(Mutex::new(ModelCache::new(resolve_max_cached_models()))),
336 active_generation: Arc::new(RwLock::new(None)),
337 config: Arc::new(tokio::sync::RwLock::new(config)),
338 start_time: Instant::now(),
339 model_load_lock: Arc::new(Mutex::new(())),
340 pull_lock: Arc::new(Mutex::new(())),
341 queue,
342 job_registry: JobRegistry::new(),
343 shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
344 shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
345 upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
346 metadata_db: Arc::new(None),
347 chain_jobs: None,
348 downloads: DownloadQueue::new(),
349 resources: ResourceBroadcaster::new(),
350 catalog_live_cache: default_live_cache(),
351 catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
352 catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
353 }
354 }
355
356 #[cfg(test)]
358 pub(crate) fn empty_gpu_pool() -> Arc<GpuPool> {
359 Arc::new(GpuPool {
360 workers: Vec::new(),
361 })
362 }
363
364 #[cfg(test)]
368 pub(crate) fn empty_gpu_pool_for_test() -> Arc<GpuPool> {
369 Self::empty_gpu_pool()
370 }
371
372 #[cfg(test)]
373 pub fn with_engine(engine: impl InferenceEngine + 'static) -> Self {
374 let (tx, _rx) = tokio::sync::mpsc::channel(16);
375 let queue = QueueHandle::new(tx);
376 let mut cache = ModelCache::new(resolve_max_cached_models());
377 cache.insert(Box::new(engine), 0);
378 Self {
379 gpu_pool: Self::empty_gpu_pool(),
380 queue_capacity: 200,
381 model_cache: Arc::new(Mutex::new(cache)),
382 active_generation: Arc::new(RwLock::new(None)),
383 config: Arc::new(tokio::sync::RwLock::new(Config::default())),
384 start_time: Instant::now(),
385 model_load_lock: Arc::new(Mutex::new(())),
386 pull_lock: Arc::new(Mutex::new(())),
387 queue,
388 job_registry: JobRegistry::new(),
389 shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
390 shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
391 upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
392 metadata_db: Arc::new(None),
393 chain_jobs: None,
394 downloads: DownloadQueue::new(),
395 resources: ResourceBroadcaster::new(),
396 catalog_live_cache: default_live_cache(),
397 catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
398 catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
399 }
400 }
401
402 #[cfg(test)]
404 pub fn with_engine_and_queue(
405 engine: impl InferenceEngine + 'static,
406 ) -> (Self, tokio::sync::mpsc::Receiver<GenerationJob>) {
407 let (tx, rx) = tokio::sync::mpsc::channel(16);
408 let queue = QueueHandle::new(tx);
409 let mut cache = ModelCache::new(resolve_max_cached_models());
410 cache.insert(Box::new(engine), 0);
411 let state = Self {
412 gpu_pool: Self::empty_gpu_pool(),
413 queue_capacity: 200,
414 model_cache: Arc::new(Mutex::new(cache)),
415 active_generation: Arc::new(RwLock::new(None)),
416 config: Arc::new(tokio::sync::RwLock::new(Config::default())),
417 start_time: Instant::now(),
418 model_load_lock: Arc::new(Mutex::new(())),
419 pull_lock: Arc::new(Mutex::new(())),
420 queue,
421 job_registry: JobRegistry::new(),
422 shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
423 shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
424 upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
425 metadata_db: Arc::new(None),
426 chain_jobs: None,
427 downloads: DownloadQueue::new(),
428 resources: ResourceBroadcaster::new(),
429 catalog_live_cache: default_live_cache(),
430 catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
431 catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
432 };
433 (state, rx)
434 }
435
436 pub fn for_tests() -> Self {
440 let (tx, _rx) = tokio::sync::mpsc::channel(16);
441 let queue = QueueHandle::new(tx);
442 Self {
443 gpu_pool: Arc::new(GpuPool {
444 workers: Vec::new(),
445 }),
446 queue_capacity: 200,
447 model_cache: Arc::new(Mutex::new(ModelCache::new(resolve_max_cached_models()))),
448 active_generation: Arc::new(RwLock::new(None)),
449 config: Arc::new(tokio::sync::RwLock::new(Config::default())),
450 start_time: Instant::now(),
451 model_load_lock: Arc::new(Mutex::new(())),
452 pull_lock: Arc::new(Mutex::new(())),
453 queue,
454 job_registry: JobRegistry::new(),
455 shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
456 shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
457 upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
458 metadata_db: Arc::new(None),
459 chain_jobs: None,
460 downloads: DownloadQueue::new(),
461 resources: ResourceBroadcaster::new(),
462 catalog_live_cache: default_live_cache(),
463 catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
464 catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
465 }
466 }
467
468 pub fn with_civitai_base(mut self, base: impl Into<String>) -> Self {
472 self.catalog_live_civitai_base = Arc::new(base.into());
473 self
474 }
475}
476
477#[cfg(test)]
478mod tests {
479 use super::*;
480
481 #[test]
482 fn engine_snapshot_default_is_unloaded() {
483 let snap = EngineSnapshot::default();
484 assert!(snap.model_name.is_none());
485 assert!(!snap.is_loaded);
486 }
487
488 #[test]
489 fn active_generation_snapshot_stores_fields() {
490 let snap = ActiveGenerationSnapshot {
491 model: "flux-dev:q8".to_string(),
492 prompt_sha256: "abc123".to_string(),
493 started_at_unix_ms: 1700000000000,
494 started_at: std::time::Instant::now(),
495 };
496 assert_eq!(snap.model, "flux-dev:q8");
497 assert_eq!(snap.prompt_sha256, "abc123");
498 assert_eq!(snap.started_at_unix_ms, 1700000000000);
499 }
500
501 #[test]
502 fn queue_handle_pending_starts_at_zero() {
503 let (tx, _rx) = tokio::sync::mpsc::channel::<GenerationJob>(16);
504 let handle = QueueHandle::new(tx);
505 assert_eq!(handle.pending(), 0);
506 }
507
508 #[test]
509 fn upscaler_cache_starts_empty() {
510 let config = mold_core::Config::default();
511 let state = AppState::empty(
512 config,
513 QueueHandle::new(tokio::sync::mpsc::channel(1).0),
514 AppState::empty_gpu_pool(),
515 200,
516 );
517 let cache = state.upscaler_cache.lock().unwrap();
518 assert!(cache.is_none());
519 }
520
521 #[test]
522 fn upscaler_cache_cleared_by_setting_none() {
523 let config = mold_core::Config::default();
524 let state = AppState::empty(
525 config,
526 QueueHandle::new(tokio::sync::mpsc::channel(1).0),
527 AppState::empty_gpu_pool(),
528 200,
529 );
530 {
531 let mut cache = state.upscaler_cache.lock().unwrap();
532 *cache = None;
533 }
534 let cache = state.upscaler_cache.lock().unwrap();
535 assert!(cache.is_none());
536 }
537
538 #[test]
539 fn app_state_exposes_resources_broadcaster() {
540 let config = mold_core::Config::default();
541 let state = AppState::empty(
542 config,
543 QueueHandle::new(tokio::sync::mpsc::channel(1).0),
544 AppState::empty_gpu_pool(),
545 200,
546 );
547 assert!(state.resources.latest().is_none());
549 let _rx = state.resources.subscribe();
551 }
552
553 static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
557
558 fn with_env<R>(name: &str, value: Option<&str>, f: impl FnOnce() -> R) -> R {
562 let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
563 let prev = std::env::var(name).ok();
564 match value {
565 Some(v) => std::env::set_var(name, v),
566 None => std::env::remove_var(name),
567 }
568 let out = f();
569 match prev {
570 Some(v) => std::env::set_var(name, v),
571 None => std::env::remove_var(name),
572 }
573 out
574 }
575
576 #[test]
577 fn resolve_max_cached_uses_default_when_env_missing() {
578 let n = with_env(MAX_CACHED_MODELS_ENV, None, resolve_max_cached_models);
579 assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
580 }
581
582 #[test]
583 fn resolve_max_cached_honors_env_within_range() {
584 let n = with_env(MAX_CACHED_MODELS_ENV, Some("8"), resolve_max_cached_models);
585 assert_eq!(n, 8);
586 }
587
588 #[test]
589 fn resolve_max_cached_clamps_zero_back_to_default() {
590 let n = with_env(MAX_CACHED_MODELS_ENV, Some("0"), resolve_max_cached_models);
591 assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
592 }
593
594 #[test]
595 fn resolve_max_cached_clamps_overflow_back_to_default() {
596 let n = with_env(
597 MAX_CACHED_MODELS_ENV,
598 Some("999"),
599 resolve_max_cached_models,
600 );
601 assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
602 }
603
604 #[test]
605 fn resolve_max_cached_falls_back_when_env_unparseable() {
606 let n = with_env(
607 MAX_CACHED_MODELS_ENV,
608 Some("not-a-number"),
609 resolve_max_cached_models,
610 );
611 assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
612 }
613
614 #[test]
615 fn resolve_cache_idle_ttl_uses_default_when_env_missing() {
616 let n = with_env(CACHE_IDLE_TTL_ENV, None, resolve_cache_idle_ttl_secs);
617 assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
618 }
619
620 #[test]
621 fn resolve_cache_idle_ttl_honors_env_within_range() {
622 let n = with_env(CACHE_IDLE_TTL_ENV, Some("600"), resolve_cache_idle_ttl_secs);
624 assert_eq!(n, 600);
625 }
626
627 #[test]
628 fn resolve_cache_idle_ttl_clamps_zero_back_to_default() {
629 let n = with_env(CACHE_IDLE_TTL_ENV, Some("0"), resolve_cache_idle_ttl_secs);
631 assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
632 }
633
634 #[test]
635 fn resolve_cache_idle_ttl_clamps_overflow_back_to_default() {
636 let n = with_env(
638 CACHE_IDLE_TTL_ENV,
639 Some("100000"),
640 resolve_cache_idle_ttl_secs,
641 );
642 assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
643 }
644
645 #[test]
646 fn resolve_cache_idle_ttl_falls_back_when_env_unparseable() {
647 let n = with_env(
648 CACHE_IDLE_TTL_ENV,
649 Some("not-a-number"),
650 resolve_cache_idle_ttl_secs,
651 );
652 assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
653 }
654}