mold_server/
state.rs

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::model_cache::ModelCache;
12
13#[derive(Debug, Clone, Default)]
14pub struct EngineSnapshot {
15    /// Currently GPU-loaded model (None if no model on GPU).
16    pub model_name: Option<String>,
17    pub is_loaded: bool,
18    /// All models in the cache (loaded + unloaded), for status display.
19    pub cached_models: Vec<String>,
20}
21
22#[derive(Debug, Clone)]
23pub struct ActiveGenerationSnapshot {
24    pub model: String,
25    pub prompt_sha256: String,
26    pub started_at_unix_ms: u64,
27    pub started_at: Instant,
28}
29
30// ── Generation queue types ──────────────────────────────────────────────────
31
32/// Internal SSE message type used by both the queue worker and SSE streams.
33pub enum SseMessage {
34    Progress(mold_core::SseProgressEvent),
35    Complete(mold_core::SseCompleteEvent),
36    UpscaleComplete(mold_core::SseUpscaleCompleteEvent),
37    Error(mold_core::SseErrorEvent),
38}
39
40/// A generation job submitted to the queue worker.
41pub struct GenerationJob {
42    pub request: mold_core::GenerateRequest,
43    /// Channel to send SSE progress/complete/error events (None for non-streaming).
44    pub progress_tx: Option<tokio::sync::mpsc::UnboundedSender<SseMessage>>,
45    /// Oneshot to return the final result for non-streaming callers.
46    pub result_tx: tokio::sync::oneshot::Sender<Result<GenerationJobResult, String>>,
47    /// Pre-resolved output directory for server-side image saving.
48    pub output_dir: Option<PathBuf>,
49}
50
51pub struct GenerationJobResult {
52    pub response: mold_core::GenerateResponse,
53    pub image: mold_core::ImageData,
54}
55
56/// Handle for submitting jobs to the generation queue.
57#[derive(Clone)]
58pub struct QueueHandle {
59    job_tx: tokio::sync::mpsc::Sender<GenerationJob>,
60    pending_count: Arc<AtomicUsize>,
61}
62
63impl QueueHandle {
64    pub fn new(job_tx: tokio::sync::mpsc::Sender<GenerationJob>) -> Self {
65        Self {
66            job_tx,
67            pending_count: Arc::new(AtomicUsize::new(0)),
68        }
69    }
70
71    /// Submit a generation job. Returns the queue position (0-based).
72    pub async fn submit(&self, job: GenerationJob) -> Result<usize, String> {
73        let position = self.pending_count.fetch_add(1, Ordering::SeqCst);
74        if let Err(_e) = self.job_tx.send(job).await {
75            self.pending_count.fetch_sub(1, Ordering::SeqCst);
76            return Err("generation queue shut down".to_string());
77        }
78        #[cfg(feature = "metrics")]
79        {
80            crate::metrics::record_queue_submit();
81            crate::metrics::record_queue_depth(self.pending_count.load(Ordering::SeqCst));
82        }
83        Ok(position)
84    }
85
86    pub fn decrement(&self) {
87        self.pending_count.fetch_sub(1, Ordering::SeqCst);
88    }
89
90    pub fn pending(&self) -> usize {
91        self.pending_count.load(Ordering::SeqCst)
92    }
93}
94
95// ── AppState ────────────────────────────────────────────────────────────────
96
97#[derive(Clone)]
98pub struct AppState {
99    pub model_cache: Arc<Mutex<ModelCache>>,
100    pub engine_snapshot: Arc<tokio::sync::RwLock<EngineSnapshot>>,
101    /// Uses std::sync::RwLock (not tokio) because it's only accessed from
102    /// synchronous contexts (inside spawn_blocking closures and brief reads).
103    /// Must never be held across an .await point.
104    pub active_generation: Arc<RwLock<Option<ActiveGenerationSnapshot>>>,
105    pub config: Arc<tokio::sync::RwLock<Config>>,
106    pub start_time: Instant,
107    /// Guards concurrent model loads and hot-swaps.
108    pub model_load_lock: Arc<Mutex<()>>,
109    /// Guards concurrent pulls — only one download at a time.
110    pub pull_lock: Arc<Mutex<()>>,
111    /// Generation request queue.
112    pub queue: QueueHandle,
113    /// Shared tokenizer pool for cross-engine caching.
114    pub shared_pool: Arc<std::sync::Mutex<SharedPool>>,
115    /// Shutdown trigger for graceful shutdown via `/api/shutdown` endpoint.
116    pub shutdown_tx: Arc<tokio::sync::Mutex<Option<tokio::sync::oneshot::Sender<()>>>>,
117    /// Cached upscaler engine to avoid recreating per request. Small models (2-64MB), single slot.
118    pub upscaler_cache: Arc<std::sync::Mutex<Option<Box<dyn mold_inference::UpscaleEngine>>>>,
119}
120
121/// Default maximum number of cached models (loaded + unloaded engine structs).
122const DEFAULT_MAX_CACHED_MODELS: usize = 3;
123
124impl AppState {
125    /// Create state with a pre-loaded engine (server starts with a configured model).
126    pub fn new(engine: Box<dyn InferenceEngine>, config: Config, queue: QueueHandle) -> Self {
127        let name = engine.model_name().to_string();
128        let loaded = engine.is_loaded();
129        let mut cache = ModelCache::new(DEFAULT_MAX_CACHED_MODELS);
130        cache.insert(engine, 0);
131        let snapshot = EngineSnapshot {
132            model_name: Some(name),
133            is_loaded: loaded,
134            cached_models: cache.cached_model_names(),
135        };
136        Self {
137            model_cache: Arc::new(Mutex::new(cache)),
138            engine_snapshot: Arc::new(tokio::sync::RwLock::new(snapshot)),
139            active_generation: Arc::new(RwLock::new(None)),
140            config: Arc::new(tokio::sync::RwLock::new(config)),
141            start_time: Instant::now(),
142            model_load_lock: Arc::new(Mutex::new(())),
143            pull_lock: Arc::new(Mutex::new(())),
144            queue,
145            shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
146            shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
147            upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
148        }
149    }
150
151    /// Create state with no engine (zero-config startup, models pulled on demand).
152    pub fn empty(config: Config, queue: QueueHandle) -> Self {
153        Self {
154            model_cache: Arc::new(Mutex::new(ModelCache::new(DEFAULT_MAX_CACHED_MODELS))),
155            engine_snapshot: Arc::new(tokio::sync::RwLock::new(EngineSnapshot::default())),
156            active_generation: Arc::new(RwLock::new(None)),
157            config: Arc::new(tokio::sync::RwLock::new(config)),
158            start_time: Instant::now(),
159            model_load_lock: Arc::new(Mutex::new(())),
160            pull_lock: Arc::new(Mutex::new(())),
161            queue,
162            shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
163            shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
164            upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
165        }
166    }
167
168    #[cfg(test)]
169    pub fn with_engine(engine: impl InferenceEngine + 'static) -> Self {
170        let (tx, _rx) = tokio::sync::mpsc::channel(16);
171        let queue = QueueHandle::new(tx);
172        let name = engine.model_name().to_string();
173        let loaded = engine.is_loaded();
174        let mut cache = ModelCache::new(DEFAULT_MAX_CACHED_MODELS);
175        cache.insert(Box::new(engine), 0);
176        let snapshot = EngineSnapshot {
177            model_name: Some(name),
178            is_loaded: loaded,
179            cached_models: cache.cached_model_names(),
180        };
181        Self {
182            model_cache: Arc::new(Mutex::new(cache)),
183            engine_snapshot: Arc::new(tokio::sync::RwLock::new(snapshot)),
184            active_generation: Arc::new(RwLock::new(None)),
185            config: Arc::new(tokio::sync::RwLock::new(Config::default())),
186            start_time: Instant::now(),
187            model_load_lock: Arc::new(Mutex::new(())),
188            pull_lock: Arc::new(Mutex::new(())),
189            queue,
190            shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
191            shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
192            upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
193        }
194    }
195
196    /// Create state with a queue whose receiver is returned for testing.
197    #[cfg(test)]
198    pub fn with_engine_and_queue(
199        engine: impl InferenceEngine + 'static,
200    ) -> (Self, tokio::sync::mpsc::Receiver<GenerationJob>) {
201        let (tx, rx) = tokio::sync::mpsc::channel(16);
202        let queue = QueueHandle::new(tx);
203        let name = engine.model_name().to_string();
204        let loaded = engine.is_loaded();
205        let mut cache = ModelCache::new(DEFAULT_MAX_CACHED_MODELS);
206        cache.insert(Box::new(engine), 0);
207        let snapshot = EngineSnapshot {
208            model_name: Some(name),
209            is_loaded: loaded,
210            cached_models: cache.cached_model_names(),
211        };
212        let state = Self {
213            model_cache: Arc::new(Mutex::new(cache)),
214            engine_snapshot: Arc::new(tokio::sync::RwLock::new(snapshot)),
215            active_generation: Arc::new(RwLock::new(None)),
216            config: Arc::new(tokio::sync::RwLock::new(Config::default())),
217            start_time: Instant::now(),
218            model_load_lock: Arc::new(Mutex::new(())),
219            pull_lock: Arc::new(Mutex::new(())),
220            queue,
221            shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
222            shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
223            upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
224        };
225        (state, rx)
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232
233    #[test]
234    fn engine_snapshot_default_is_unloaded() {
235        let snap = EngineSnapshot::default();
236        assert!(snap.model_name.is_none());
237        assert!(!snap.is_loaded);
238    }
239
240    #[test]
241    fn active_generation_snapshot_stores_fields() {
242        let snap = ActiveGenerationSnapshot {
243            model: "flux-dev:q8".to_string(),
244            prompt_sha256: "abc123".to_string(),
245            started_at_unix_ms: 1700000000000,
246            started_at: std::time::Instant::now(),
247        };
248        assert_eq!(snap.model, "flux-dev:q8");
249        assert_eq!(snap.prompt_sha256, "abc123");
250        assert_eq!(snap.started_at_unix_ms, 1700000000000);
251    }
252
253    #[test]
254    fn queue_handle_pending_starts_at_zero() {
255        let (tx, _rx) = tokio::sync::mpsc::channel::<GenerationJob>(16);
256        let handle = QueueHandle::new(tx);
257        assert_eq!(handle.pending(), 0);
258    }
259
260    #[test]
261    fn upscaler_cache_starts_empty() {
262        let config = mold_core::Config::default();
263        let state = AppState::empty(config, QueueHandle::new(tokio::sync::mpsc::channel(1).0));
264        let cache = state.upscaler_cache.lock().unwrap();
265        assert!(cache.is_none());
266    }
267
268    #[test]
269    fn upscaler_cache_cleared_by_setting_none() {
270        let config = mold_core::Config::default();
271        let state = AppState::empty(config, QueueHandle::new(tokio::sync::mpsc::channel(1).0));
272        {
273            let mut cache = state.upscaler_cache.lock().unwrap();
274            *cache = None;
275        }
276        let cache = state.upscaler_cache.lock().unwrap();
277        assert!(cache.is_none());
278    }
279}
mold_server/state.rs

mold_server/
state.rs
