diff --git a/crates/application/src/model_server.rs b/crates/application/src/model_server.rs index fd0ca51..f9abbf4 100644 --- a/crates/application/src/model_server.rs +++ b/crates/application/src/model_server.rs @@ -14,6 +14,7 @@ use domain::ports::{ ModelServerRegistry, ModelServerRuntime, ProcessStatus, ProfileStore, RemotePath, }; use domain::{LocalModelServerId, StopPolicy}; +use tokio::sync::{Mutex as AsyncMutex, Notify}; use crate::error::AppError; @@ -167,6 +168,35 @@ struct ActiveServer { stop_policy: StopPolicy, } +#[derive(Debug)] +struct InflightEnsure { + result: Mutex>>, + notify: Notify, +} + +impl InflightEnsure { + fn new() -> Self { + Self { + result: Mutex::new(None), + notify: Notify::new(), + } + } + + async fn wait(&self) -> Result { + loop { + if let Some(result) = self.result.lock().unwrap().clone() { + return result; + } + self.notify.notified().await; + } + } + + fn complete(&self, result: Result) { + *self.result.lock().unwrap() = Some(result); + self.notify.notify_waiters(); + } +} + /// Ensures a configured local model server is reachable, starting it when allowed. pub struct EnsureLocalModelServer { registry: Arc, @@ -176,6 +206,7 @@ pub struct EnsureLocalModelServer { fs: Arc, events: Arc, active: Mutex>, + inflight: AsyncMutex>>, readiness: ReadinessPolicy, } @@ -199,6 +230,7 @@ impl EnsureLocalModelServer { fs, events, active: Mutex::new(HashMap::new()), + inflight: AsyncMutex::new(HashMap::new()), readiness: ReadinessPolicy::default(), } } @@ -217,6 +249,38 @@ impl EnsureLocalModelServer { pub async fn execute( &self, input: EnsureLocalModelServerInput, + ) -> Result { + let server_id = input.server_id; + let (flight, is_leader) = { + let mut inflight = self.inflight.lock().await; + if let Some(flight) = inflight.get(&server_id) { + (Arc::clone(flight), false) + } else { + let flight = Arc::new(InflightEnsure::new()); + inflight.insert(server_id, Arc::clone(&flight)); + (flight, true) + } + }; + + if !is_leader { + return flight.wait().await; + } + + let result = self.execute_inner(input).await; + flight.complete(result.clone()); + let mut inflight = self.inflight.lock().await; + if inflight + .get(&server_id) + .is_some_and(|current| Arc::ptr_eq(current, &flight)) + { + inflight.remove(&server_id); + } + result + } + + async fn execute_inner( + &self, + input: EnsureLocalModelServerInput, ) -> Result { let config = self .registry diff --git a/crates/application/tests/model_server.rs b/crates/application/tests/model_server.rs index 3fd9981..0abb431 100644 --- a/crates/application/tests/model_server.rs +++ b/crates/application/tests/model_server.rs @@ -158,11 +158,15 @@ struct FakeProcess { spawns: Mutex>, kills: Mutex>, statuses: Mutex>, + spawn_delay: Duration, } #[async_trait] impl ManagedProcess for FakeProcess { async fn spawn(&self, spec: SpawnSpec) -> Result { + if !self.spawn_delay.is_zero() { + tokio::time::sleep(self.spawn_delay).await; + } self.spawns.lock().unwrap().push(spec); let handle = ManagedProcessHandle { id: format!("h{}", self.spawns.lock().unwrap().len()), @@ -373,6 +377,73 @@ async fn absent_auto_start_spawns_and_waits_until_ready() { ); } +#[tokio::test] +async fn concurrent_ensure_same_server_shares_one_start_attempt() { + let registry = Arc::new(FakeRegistry::default()); + registry + .save(config(sid(10), 8090, "/models/qwen.gguf", true)) + .await + .unwrap(); + let fs = Arc::new(FakeFs::default()); + fs.existing + .lock() + .unwrap() + .push("/models/qwen.gguf".to_owned()); + let process = Arc::new(FakeProcess { + spawn_delay: Duration::from_millis(25), + ..FakeProcess::default() + }); + let events = Arc::new(FakeEvents::default()); + let usecase = Arc::new(ensure( + Arc::clone(®istry), + Arc::new(FakeProbe::new(vec![ + ModelServerStatus::Unreachable, + ModelServerStatus::ReadyReused, + ])), + Arc::clone(&process), + fs, + Arc::clone(&events), + )); + + let mut tasks = Vec::new(); + for _ in 0..5 { + let usecase = Arc::clone(&usecase); + tasks.push(tokio::spawn(async move { + usecase + .execute(EnsureLocalModelServerInput { server_id: sid(10) }) + .await + .unwrap() + })); + } + + let mut outputs = Vec::new(); + for task in tasks { + outputs.push(task.await.unwrap()); + } + + assert!(outputs + .iter() + .all(|out| out.ready.status == ModelServerStatus::ReadyStarted)); + assert_eq!(process.spawns.lock().unwrap().len(), 1); + + let starting_events = events + .0 + .lock() + .unwrap() + .iter() + .filter(|event| { + matches!( + event, + DomainEvent::ModelServerStatusChanged { + status: domain::model_server::ModelServerLifecycleStatus::Starting, + .. + } + ) + }) + .count(); + assert_eq!(starting_events, 1); +} + #[tokio::test] async fn missing_model_path_is_path_not_accessible() { let registry = Arc::new(FakeRegistry::default());