use mold_core::Config;
use mold_inference::InferenceEngine;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, RwLock};
use std::time::Instant;
use tokio::sync::Mutex;
use mold_inference::shared_pool::SharedPool;
use crate::downloads::DownloadQueue;
use crate::gpu_pool::GpuPool;
use crate::job_registry::{JobRegistry, SharedJobRegistry};
use crate::model_cache::ModelCache;
use crate::resources::ResourceBroadcaster;
#[derive(Debug, Clone, Default)]
pub struct EngineSnapshot {
pub model_name: Option<String>,
pub is_loaded: bool,
pub cached_models: Vec<String>,
}
#[derive(Debug, Clone)]
pub struct ActiveGenerationSnapshot {
pub model: String,
pub prompt_sha256: String,
pub started_at_unix_ms: u64,
pub started_at: Instant,
}
pub enum SseMessage {
Progress(mold_core::SseProgressEvent),
Complete(mold_core::SseCompleteEvent),
UpscaleComplete(mold_core::SseUpscaleCompleteEvent),
Error(mold_core::SseErrorEvent),
}
pub struct GenerationJob {
pub id: String,
pub request: mold_core::GenerateRequest,
pub progress_tx: Option<tokio::sync::mpsc::UnboundedSender<SseMessage>>,
pub result_tx: tokio::sync::oneshot::Sender<Result<GenerationJobResult, String>>,
pub output_dir: Option<PathBuf>,
}
pub struct GenerationJobResult {
pub response: mold_core::GenerateResponse,
pub image: mold_core::ImageData,
}
#[derive(Clone)]
pub struct QueueHandle {
job_tx: tokio::sync::mpsc::Sender<GenerationJob>,
pending_count: Arc<AtomicUsize>,
}
#[derive(Debug)]
pub enum SubmitError {
Full { pending: usize, capacity: usize },
Shutdown,
}
impl QueueHandle {
pub fn new(job_tx: tokio::sync::mpsc::Sender<GenerationJob>) -> Self {
Self {
job_tx,
pending_count: Arc::new(AtomicUsize::new(0)),
}
}
pub async fn submit(&self, job: GenerationJob, capacity: usize) -> Result<usize, SubmitError> {
let prev = self.pending_count.fetch_add(1, Ordering::SeqCst);
if prev >= capacity {
self.pending_count.fetch_sub(1, Ordering::SeqCst);
return Err(SubmitError::Full {
pending: prev,
capacity,
});
}
if self.job_tx.send(job).await.is_err() {
self.pending_count.fetch_sub(1, Ordering::SeqCst);
return Err(SubmitError::Shutdown);
}
#[cfg(feature = "metrics")]
{
crate::metrics::record_queue_submit();
crate::metrics::record_queue_depth(self.pending_count.load(Ordering::SeqCst));
}
Ok(prev)
}
pub fn decrement(&self) {
self.pending_count.fetch_sub(1, Ordering::SeqCst);
}
pub fn pending(&self) -> usize {
self.pending_count.load(Ordering::SeqCst)
}
}
#[derive(Clone)]
pub struct AppState {
pub gpu_pool: Arc<GpuPool>,
pub queue_capacity: usize,
pub model_cache: Arc<Mutex<ModelCache>>,
pub active_generation: Arc<RwLock<Option<ActiveGenerationSnapshot>>>,
pub config: Arc<tokio::sync::RwLock<Config>>,
pub start_time: Instant,
pub model_load_lock: Arc<Mutex<()>>,
pub pull_lock: Arc<Mutex<()>>,
pub queue: QueueHandle,
pub job_registry: SharedJobRegistry,
pub shared_pool: Arc<std::sync::Mutex<SharedPool>>,
pub shutdown_tx: Arc<tokio::sync::Mutex<Option<tokio::sync::oneshot::Sender<()>>>>,
pub upscaler_cache: Arc<std::sync::Mutex<Option<Box<dyn mold_inference::UpscaleEngine>>>>,
pub metadata_db: Arc<Option<mold_db::MetadataDb>>,
pub chain_jobs: Option<Arc<crate::chain_job_runner::ChainJobRunnerHandle>>,
pub downloads: Arc<DownloadQueue>,
pub resources: Arc<ResourceBroadcaster>,
pub catalog_live_cache: mold_catalog::live::LiveCache,
pub catalog_live_civitai_base: Arc<String>,
pub catalog_intents: Arc<
tokio::sync::RwLock<
std::collections::HashMap<String, mold_catalog::synthesis::CatalogModelIntent>,
>,
>,
}
const CATALOG_LIVE_CACHE_TTL_SECS: u64 = 300;
const CATALOG_LIVE_CACHE_MAX_KEYS: usize = 256;
pub const CATALOG_LIVE_CIVITAI_BASE: &str = "https://civitai.com";
fn default_live_cache() -> mold_catalog::live::LiveCache {
mold_catalog::live::LiveCache::new(
std::time::Duration::from_secs(CATALOG_LIVE_CACHE_TTL_SECS),
CATALOG_LIVE_CACHE_MAX_KEYS,
)
}
pub const DEFAULT_MAX_CACHED_MODELS: usize = 3;
const MAX_CACHED_MODELS_LOWER: usize = 1;
const MAX_CACHED_MODELS_UPPER: usize = 16;
pub const MAX_CACHED_MODELS_ENV: &str = "MOLD_MAX_CACHED_MODELS";
pub const DEFAULT_CACHE_IDLE_TTL_SECS: u64 = 1800;
const CACHE_IDLE_TTL_LOWER_SECS: u64 = 60;
const CACHE_IDLE_TTL_UPPER_SECS: u64 = 86_400;
pub const CACHE_IDLE_TTL_ENV: &str = "MOLD_CACHE_IDLE_TTL_SECS";
pub fn resolve_cache_idle_ttl_secs() -> u64 {
match std::env::var(CACHE_IDLE_TTL_ENV) {
Ok(raw) => match raw.trim().parse::<u64>() {
Ok(n) if (CACHE_IDLE_TTL_LOWER_SECS..=CACHE_IDLE_TTL_UPPER_SECS).contains(&n) => n,
Ok(n) => {
tracing::warn!(
env = CACHE_IDLE_TTL_ENV,
value = n,
lower = CACHE_IDLE_TTL_LOWER_SECS,
upper = CACHE_IDLE_TTL_UPPER_SECS,
"ignoring out-of-range cache idle-TTL; using default"
);
DEFAULT_CACHE_IDLE_TTL_SECS
}
Err(e) => {
tracing::warn!(
env = CACHE_IDLE_TTL_ENV,
raw = %raw,
error = %e,
"ignoring unparseable cache idle-TTL; using default"
);
DEFAULT_CACHE_IDLE_TTL_SECS
}
},
Err(_) => DEFAULT_CACHE_IDLE_TTL_SECS,
}
}
pub fn resolve_max_cached_models() -> usize {
match std::env::var(MAX_CACHED_MODELS_ENV) {
Ok(raw) => match raw.trim().parse::<usize>() {
Ok(n) if (MAX_CACHED_MODELS_LOWER..=MAX_CACHED_MODELS_UPPER).contains(&n) => n,
Ok(n) => {
tracing::warn!(
env = MAX_CACHED_MODELS_ENV,
value = n,
lower = MAX_CACHED_MODELS_LOWER,
upper = MAX_CACHED_MODELS_UPPER,
"ignoring out-of-range cache cap; using default"
);
DEFAULT_MAX_CACHED_MODELS
}
Err(e) => {
tracing::warn!(
env = MAX_CACHED_MODELS_ENV,
raw = %raw,
error = %e,
"ignoring unparseable cache cap; using default"
);
DEFAULT_MAX_CACHED_MODELS
}
},
Err(_) => DEFAULT_MAX_CACHED_MODELS,
}
}
impl AppState {
pub fn new(
engine: Box<dyn InferenceEngine>,
config: Config,
queue: QueueHandle,
gpu_pool: Arc<GpuPool>,
queue_capacity: usize,
) -> Self {
let mut cache = ModelCache::new(resolve_max_cached_models());
cache.insert(engine, 0);
Self {
gpu_pool,
queue_capacity,
model_cache: Arc::new(Mutex::new(cache)),
active_generation: Arc::new(RwLock::new(None)),
config: Arc::new(tokio::sync::RwLock::new(config)),
start_time: Instant::now(),
model_load_lock: Arc::new(Mutex::new(())),
pull_lock: Arc::new(Mutex::new(())),
queue,
job_registry: JobRegistry::new(),
shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
metadata_db: Arc::new(None),
chain_jobs: None,
downloads: DownloadQueue::new(),
resources: ResourceBroadcaster::new(),
catalog_live_cache: default_live_cache(),
catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
}
}
pub fn empty(
config: Config,
queue: QueueHandle,
gpu_pool: Arc<GpuPool>,
queue_capacity: usize,
) -> Self {
Self {
gpu_pool,
queue_capacity,
model_cache: Arc::new(Mutex::new(ModelCache::new(resolve_max_cached_models()))),
active_generation: Arc::new(RwLock::new(None)),
config: Arc::new(tokio::sync::RwLock::new(config)),
start_time: Instant::now(),
model_load_lock: Arc::new(Mutex::new(())),
pull_lock: Arc::new(Mutex::new(())),
queue,
job_registry: JobRegistry::new(),
shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
metadata_db: Arc::new(None),
chain_jobs: None,
downloads: DownloadQueue::new(),
resources: ResourceBroadcaster::new(),
catalog_live_cache: default_live_cache(),
catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
}
}
#[cfg(test)]
pub(crate) fn empty_gpu_pool() -> Arc<GpuPool> {
Arc::new(GpuPool {
workers: Vec::new(),
})
}
#[cfg(test)]
pub(crate) fn empty_gpu_pool_for_test() -> Arc<GpuPool> {
Self::empty_gpu_pool()
}
#[cfg(test)]
pub fn with_engine(engine: impl InferenceEngine + 'static) -> Self {
let (tx, _rx) = tokio::sync::mpsc::channel(16);
let queue = QueueHandle::new(tx);
let mut cache = ModelCache::new(resolve_max_cached_models());
cache.insert(Box::new(engine), 0);
Self {
gpu_pool: Self::empty_gpu_pool(),
queue_capacity: 200,
model_cache: Arc::new(Mutex::new(cache)),
active_generation: Arc::new(RwLock::new(None)),
config: Arc::new(tokio::sync::RwLock::new(Config::default())),
start_time: Instant::now(),
model_load_lock: Arc::new(Mutex::new(())),
pull_lock: Arc::new(Mutex::new(())),
queue,
job_registry: JobRegistry::new(),
shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
metadata_db: Arc::new(None),
chain_jobs: None,
downloads: DownloadQueue::new(),
resources: ResourceBroadcaster::new(),
catalog_live_cache: default_live_cache(),
catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
}
}
#[cfg(test)]
pub fn with_engine_and_queue(
engine: impl InferenceEngine + 'static,
) -> (Self, tokio::sync::mpsc::Receiver<GenerationJob>) {
let (tx, rx) = tokio::sync::mpsc::channel(16);
let queue = QueueHandle::new(tx);
let mut cache = ModelCache::new(resolve_max_cached_models());
cache.insert(Box::new(engine), 0);
let state = Self {
gpu_pool: Self::empty_gpu_pool(),
queue_capacity: 200,
model_cache: Arc::new(Mutex::new(cache)),
active_generation: Arc::new(RwLock::new(None)),
config: Arc::new(tokio::sync::RwLock::new(Config::default())),
start_time: Instant::now(),
model_load_lock: Arc::new(Mutex::new(())),
pull_lock: Arc::new(Mutex::new(())),
queue,
job_registry: JobRegistry::new(),
shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
metadata_db: Arc::new(None),
chain_jobs: None,
downloads: DownloadQueue::new(),
resources: ResourceBroadcaster::new(),
catalog_live_cache: default_live_cache(),
catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
};
(state, rx)
}
pub fn for_tests() -> Self {
let (tx, _rx) = tokio::sync::mpsc::channel(16);
let queue = QueueHandle::new(tx);
Self {
gpu_pool: Arc::new(GpuPool {
workers: Vec::new(),
}),
queue_capacity: 200,
model_cache: Arc::new(Mutex::new(ModelCache::new(resolve_max_cached_models()))),
active_generation: Arc::new(RwLock::new(None)),
config: Arc::new(tokio::sync::RwLock::new(Config::default())),
start_time: Instant::now(),
model_load_lock: Arc::new(Mutex::new(())),
pull_lock: Arc::new(Mutex::new(())),
queue,
job_registry: JobRegistry::new(),
shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
metadata_db: Arc::new(None),
chain_jobs: None,
downloads: DownloadQueue::new(),
resources: ResourceBroadcaster::new(),
catalog_live_cache: default_live_cache(),
catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
}
}
pub fn with_civitai_base(mut self, base: impl Into<String>) -> Self {
self.catalog_live_civitai_base = Arc::new(base.into());
self
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn engine_snapshot_default_is_unloaded() {
let snap = EngineSnapshot::default();
assert!(snap.model_name.is_none());
assert!(!snap.is_loaded);
}
#[test]
fn active_generation_snapshot_stores_fields() {
let snap = ActiveGenerationSnapshot {
model: "flux-dev:q8".to_string(),
prompt_sha256: "abc123".to_string(),
started_at_unix_ms: 1700000000000,
started_at: std::time::Instant::now(),
};
assert_eq!(snap.model, "flux-dev:q8");
assert_eq!(snap.prompt_sha256, "abc123");
assert_eq!(snap.started_at_unix_ms, 1700000000000);
}
#[test]
fn queue_handle_pending_starts_at_zero() {
let (tx, _rx) = tokio::sync::mpsc::channel::<GenerationJob>(16);
let handle = QueueHandle::new(tx);
assert_eq!(handle.pending(), 0);
}
#[test]
fn upscaler_cache_starts_empty() {
let config = mold_core::Config::default();
let state = AppState::empty(
config,
QueueHandle::new(tokio::sync::mpsc::channel(1).0),
AppState::empty_gpu_pool(),
200,
);
let cache = state.upscaler_cache.lock().unwrap();
assert!(cache.is_none());
}
#[test]
fn upscaler_cache_cleared_by_setting_none() {
let config = mold_core::Config::default();
let state = AppState::empty(
config,
QueueHandle::new(tokio::sync::mpsc::channel(1).0),
AppState::empty_gpu_pool(),
200,
);
{
let mut cache = state.upscaler_cache.lock().unwrap();
*cache = None;
}
let cache = state.upscaler_cache.lock().unwrap();
assert!(cache.is_none());
}
#[test]
fn app_state_exposes_resources_broadcaster() {
let config = mold_core::Config::default();
let state = AppState::empty(
config,
QueueHandle::new(tokio::sync::mpsc::channel(1).0),
AppState::empty_gpu_pool(),
200,
);
assert!(state.resources.latest().is_none());
let _rx = state.resources.subscribe();
}
static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
fn with_env<R>(name: &str, value: Option<&str>, f: impl FnOnce() -> R) -> R {
let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prev = std::env::var(name).ok();
match value {
Some(v) => std::env::set_var(name, v),
None => std::env::remove_var(name),
}
let out = f();
match prev {
Some(v) => std::env::set_var(name, v),
None => std::env::remove_var(name),
}
out
}
#[test]
fn resolve_max_cached_uses_default_when_env_missing() {
let n = with_env(MAX_CACHED_MODELS_ENV, None, resolve_max_cached_models);
assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
}
#[test]
fn resolve_max_cached_honors_env_within_range() {
let n = with_env(MAX_CACHED_MODELS_ENV, Some("8"), resolve_max_cached_models);
assert_eq!(n, 8);
}
#[test]
fn resolve_max_cached_clamps_zero_back_to_default() {
let n = with_env(MAX_CACHED_MODELS_ENV, Some("0"), resolve_max_cached_models);
assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
}
#[test]
fn resolve_max_cached_clamps_overflow_back_to_default() {
let n = with_env(
MAX_CACHED_MODELS_ENV,
Some("999"),
resolve_max_cached_models,
);
assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
}
#[test]
fn resolve_max_cached_falls_back_when_env_unparseable() {
let n = with_env(
MAX_CACHED_MODELS_ENV,
Some("not-a-number"),
resolve_max_cached_models,
);
assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
}
#[test]
fn resolve_cache_idle_ttl_uses_default_when_env_missing() {
let n = with_env(CACHE_IDLE_TTL_ENV, None, resolve_cache_idle_ttl_secs);
assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
}
#[test]
fn resolve_cache_idle_ttl_honors_env_within_range() {
let n = with_env(CACHE_IDLE_TTL_ENV, Some("600"), resolve_cache_idle_ttl_secs);
assert_eq!(n, 600);
}
#[test]
fn resolve_cache_idle_ttl_clamps_zero_back_to_default() {
let n = with_env(CACHE_IDLE_TTL_ENV, Some("0"), resolve_cache_idle_ttl_secs);
assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
}
#[test]
fn resolve_cache_idle_ttl_clamps_overflow_back_to_default() {
let n = with_env(
CACHE_IDLE_TTL_ENV,
Some("100000"),
resolve_cache_idle_ttl_secs,
);
assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
}
#[test]
fn resolve_cache_idle_ttl_falls_back_when_env_unparseable() {
let n = with_env(
CACHE_IDLE_TTL_ENV,
Some("not-a-number"),
resolve_cache_idle_ttl_secs,
);
assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
}
}