Skip to main content

mold_server/
state.rs

1use mold_core::Config;
2use mold_inference::InferenceEngine;
3use std::sync::Arc;
4use std::time::Instant;
5use tokio::sync::Mutex;
6
7#[derive(Clone)]
8pub struct AppState {
9    pub engine: Arc<Mutex<Option<Box<dyn InferenceEngine>>>>,
10    pub config: Arc<tokio::sync::RwLock<Config>>,
11    pub start_time: Instant,
12    /// Guards concurrent model loads and hot-swaps.
13    pub model_load_lock: Arc<Mutex<()>>,
14    /// Guards concurrent pulls — only one download at a time.
15    pub pull_lock: Arc<Mutex<()>>,
16}
17
18impl AppState {
19    /// Create state with a pre-loaded engine (server starts with a configured model).
20    pub fn new(engine: Box<dyn InferenceEngine>, config: Config) -> Self {
21        Self {
22            engine: Arc::new(Mutex::new(Some(engine))),
23            config: Arc::new(tokio::sync::RwLock::new(config)),
24            start_time: Instant::now(),
25            model_load_lock: Arc::new(Mutex::new(())),
26            pull_lock: Arc::new(Mutex::new(())),
27        }
28    }
29
30    /// Create state with no engine (zero-config startup, models pulled on demand).
31    pub fn empty(config: Config) -> Self {
32        Self {
33            engine: Arc::new(Mutex::new(None)),
34            config: Arc::new(tokio::sync::RwLock::new(config)),
35            start_time: Instant::now(),
36            model_load_lock: Arc::new(Mutex::new(())),
37            pull_lock: Arc::new(Mutex::new(())),
38        }
39    }
40
41    #[cfg(test)]
42    pub fn with_engine(engine: impl InferenceEngine + 'static) -> Self {
43        Self {
44            engine: Arc::new(Mutex::new(Some(Box::new(engine)))),
45            config: Arc::new(tokio::sync::RwLock::new(Config::default())),
46            start_time: Instant::now(),
47            model_load_lock: Arc::new(Mutex::new(())),
48            pull_lock: Arc::new(Mutex::new(())),
49        }
50    }
51}