diff --git a/Cargo.lock b/Cargo.lock index 5c691c5..d529321 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2136,6 +2136,27 @@ version = "2.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d7a1e2f27636f116493b8b860f5546edb47c8d8f8ea73e1d2a20be88e28d1fea" +[[package]] +name = "deadpool" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb84100978c1c7b37f09ed3ce3e5f843af02c2a2c431bae5b19230dad2c1b490" +dependencies = [ + "async-trait", + "deadpool-runtime", + "num_cpus", + "tokio", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b" +dependencies = [ + "tokio", +] + [[package]] name = "deflate64" version = "0.1.10" @@ -2286,6 +2307,7 @@ version = "0.3.2" dependencies = [ "aho-corasick", "anyhow", + "async-trait", "axum", "axum-extra", "base64", @@ -2294,6 +2316,7 @@ dependencies = [ "cached", "chrono", "clap", + "deadpool", "env_logger", "futures-util", "hanconv", diff --git a/Cargo.toml b/Cargo.toml index 02d1016..1bfd2f1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,8 @@ serde_json = "1.0" toml = "0.8" rmp-serde = "1" +deadpool = { version = "0.10", features = ["rt_tokio_1"] } +async-trait = "0.1" cached = { version = "0.56.0", features = ["async"] } anyhow = "1.0" diff --git a/src/services/ws/stable/claude.rs b/src/services/ws/stable/claude.rs index 33a34bb..6091ccb 100644 --- a/src/services/ws/stable/claude.rs +++ b/src/services/ws/stable/claude.rs @@ -223,17 +223,17 @@ impl RunSessionState { let mut rx_list = LinkedList::new(); for chunk in finished_output { - let (tts_response_tx, tts_response_rx) = tokio::sync::mpsc::unbounded_channel(); - if let Err(e) = tts_req_tx.send((chunk.to_string(), tts_response_tx)).await { - log::error!( - "{}:{:x} error sending tts request: {}", - self.session.id, - self.session.request_id, - e - ); - } else { - rx_list.push_back((chunk, tts_response_rx)); - } + let tts_response_rx = super::tts::submit_request(tts_req_tx, chunk.to_string()) + .await + .map_err(|e| { + anyhow::anyhow!( + "{}:{:x} error sending tts request: {}", + self.session.id, + self.session.request_id, + e + ) + })?; + rx_list.push_back((chunk, tts_response_rx)); } for (text_chunk, mut tts_response_rx) in rx_list { @@ -746,7 +746,12 @@ pub async fn run_session_manager( mut session_rx: tokio::sync::mpsc::UnboundedReceiver, notifications: Arc>, ) -> anyhow::Result<()> { - let mut tts_session_pool = super::tts::TTSSessionPool::new(tts.clone(), 4); + let mut tts_session_pool = super::tts::TTSSessionPool::new( + tts.clone(), + super::tts::DEFAULT_TTS_IDLE_WORKERS, + super::tts::DEFAULT_TTS_MAX_WORKERS, + super::tts::DEFAULT_TTS_IDLE_TIMEOUT, + ); let (tts_req_tx, tts_req_rx) = tokio::sync::mpsc::channel(128); let mut sessions: HashMap< diff --git a/src/services/ws/stable/gemini.rs b/src/services/ws/stable/gemini.rs index 1e30857..19468e1 100644 --- a/src/services/ws/stable/gemini.rs +++ b/src/services/ws/stable/gemini.rs @@ -27,7 +27,12 @@ pub async fn run_session_manager( > = HashMap::new(); let tts_req_tx = if let Some(tts) = tts { - let mut tts_session_pool = super::tts::TTSSessionPool::new(tts.clone(), 4); + let mut tts_session_pool = super::tts::TTSSessionPool::new( + tts.clone(), + super::tts::DEFAULT_TTS_IDLE_WORKERS, + super::tts::DEFAULT_TTS_MAX_WORKERS, + super::tts::DEFAULT_TTS_IDLE_TIMEOUT, + ); let (tts_req_tx, tts_req_rx) = tokio::sync::mpsc::channel(128); tokio::spawn(async move { @@ -471,8 +476,7 @@ async fn run_session_with_tts( gemini::types::ServerContent::Interrupted(_) => {} gemini::types::ServerContent::TurnComplete(_) => { let (chunks_tx, chunks_rx) = tokio::sync::mpsc::unbounded_channel(); - let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); - tts_req_tx.send((llm_text.clone(), tx)).await?; + let rx = super::tts::submit_request(tts_req_tx, llm_text.clone()).await?; chunks_tx.send((llm_text.clone(), rx))?; asr_text.clear(); llm_text = String::with_capacity(1024); diff --git a/src/services/ws/stable/llm.rs b/src/services/ws/stable/llm.rs index be70f3a..3ce83c9 100644 --- a/src/services/ws/stable/llm.rs +++ b/src/services/ws/stable/llm.rs @@ -11,6 +11,19 @@ pub type ChunksRx = tokio::sync::mpsc::UnboundedReceiver<(String, super::tts::TT use tokio::time::Duration; +async fn queue_tts_request( + tts_tx: &mut super::tts::TTSRequestTx, + chunks_tx: &ChunksTx, + text: String, +) -> anyhow::Result<()> { + let tts_resp_rx = super::tts::submit_request(tts_tx, text.clone()).await?; + + chunks_tx + .send((text, tts_resp_rx)) + .map_err(|e| anyhow::anyhow!("error sending tts chunks receiver: {e}"))?; + Ok(()) +} + #[cached::proc_macro::cached(time = 60, size = 100, result = true)] async fn load_url_content(url: String) -> anyhow::Result { let client = reqwest::Client::new(); @@ -214,18 +227,7 @@ pub async fn chat( continue; } - let (tts_resp_tx, tts_resp_rx) = tokio::sync::mpsc::unbounded_channel(); - - tts_tx - .send((chunk_.to_string(), tts_resp_tx)) - .await - .map_err(|e| anyhow::anyhow!("error sending tts request for llm chunk: {e}"))?; - - chunks_tx - .send((chunk_.to_string(), tts_resp_rx)) - .map_err(|e| { - anyhow::anyhow!("error sending tts chunks receiver for llm chunk: {e}") - })?; + queue_tts_request(tts_tx, &chunks_tx, chunk_.to_string()).await?; } Ok(StableLLMResponseChunk::Functions(functions)) => { log::info!("llm functions: {:#?}", functions); @@ -234,20 +236,7 @@ pub async fn chat( if let Some(message) = chat_session.get_tool_call_message(&function) { log::info!("tool {} call message: {}", &function.function.name, message); if !message.is_empty() { - let (tts_resp_tx, tts_resp_rx) = tokio::sync::mpsc::unbounded_channel(); - - tts_tx - .send((message.to_string(), tts_resp_tx)) - .await - .map_err(|e| { - anyhow::anyhow!("error sending tts request for llm chunk: {e}") - })?; - - chunks_tx.send((message, tts_resp_rx)).map_err(|e| { - anyhow::anyhow!( - "error sending tts chunks receiver for llm chunk: {e}" - ) - })?; + queue_tts_request(tts_tx, &chunks_tx, message).await?; } } chat_session.execute_tool(&function).await? @@ -381,22 +370,7 @@ pub async fn responses( continue; } - let (tts_resp_tx, tts_resp_rx) = tokio::sync::mpsc::unbounded_channel(); - - tts_tx - .send((chunk_.to_string(), tts_resp_tx)) - .await - .map_err(|e| { - anyhow::anyhow!("error sending tts request for llm responses chunk: {e}") - })?; - - chunks_tx - .send((chunk_.to_string(), tts_resp_rx)) - .map_err(|e| { - anyhow::anyhow!( - "error sending tts chunks receiver for llm responses chunk: {e}" - ) - })?; + queue_tts_request(tts_tx, &chunks_tx, chunk_.to_string()).await?; } LLMResponsesChunk::Functions(functions) => { log::info!("llm responses functions: {:#?}", functions); @@ -405,20 +379,7 @@ pub async fn responses( if let Some(message) = responses_session.get_tool_call_message(&function) { log::info!("tool {} call message: {}", &function.function.name, message); if !message.is_empty() { - let (tts_resp_tx, tts_resp_rx) = tokio::sync::mpsc::unbounded_channel(); - - tts_tx - .send((message.to_string(), tts_resp_tx)) - .await - .map_err(|e| { - anyhow::anyhow!("error sending tts request for llm chunk: {e}") - })?; - - chunks_tx.send((message, tts_resp_rx)).map_err(|e| { - anyhow::anyhow!( - "error sending tts chunks receiver for llm chunk: {e}" - ) - })?; + queue_tts_request(tts_tx, &chunks_tx, message).await?; } } let result = responses_session.execute_tool(&function).await; diff --git a/src/services/ws/stable/mod.rs b/src/services/ws/stable/mod.rs index 0b01237..2ed1294 100644 --- a/src/services/ws/stable/mod.rs +++ b/src/services/ws/stable/mod.rs @@ -380,7 +380,12 @@ pub async fn run_session_manager( tokio::sync::mpsc::UnboundedSender<(Session, Option)>, > = HashMap::new(); - let mut tts_session_pool = tts::TTSSessionPool::new(tts.clone(), 4); + let mut tts_session_pool = tts::TTSSessionPool::new( + tts.clone(), + tts::DEFAULT_TTS_IDLE_WORKERS, + tts::DEFAULT_TTS_MAX_WORKERS, + tts::DEFAULT_TTS_IDLE_TIMEOUT, + ); let (tts_req_tx, tts_req_rx) = tokio::sync::mpsc::channel(128); tokio::spawn(async move { diff --git a/src/services/ws/stable/tts.rs b/src/services/ws/stable/tts.rs index f169cf4..769eefe 100644 --- a/src/services/ws/stable/tts.rs +++ b/src/services/ws/stable/tts.rs @@ -1,15 +1,34 @@ use bytes::{BufMut, Bytes}; +use std::time::{Duration, Instant}; +use std::{cell::Cell, sync::Mutex}; use crate::config::{ElevenlabsTTS, FishTTS, GSVTTS, GroqTTS, OpenaiTTS, StreamGSV}; -pub type TTSRequest = (String, TTSResponseTx); +pub type TTSRequest = (String, TTSResponseTx, TTSRequestAckTx); pub type TTSRequestTx = tokio::sync::mpsc::Sender; pub type TTSRequestRx = tokio::sync::mpsc::Receiver; +pub type TTSRequestAckTx = tokio::sync::oneshot::Sender>; pub type TTSResponseRx = tokio::sync::mpsc::UnboundedReceiver>; pub type TTSResponseTx = tokio::sync::mpsc::UnboundedSender>; +pub async fn submit_request(tts_tx: &TTSRequestTx, text: String) -> anyhow::Result { + let (tts_resp_tx, tts_resp_rx) = tokio::sync::mpsc::unbounded_channel(); + let (request_ack_tx, request_ack_rx) = tokio::sync::oneshot::channel(); + + tts_tx + .send((text, tts_resp_tx, request_ack_tx)) + .await + .map_err(|e| anyhow::anyhow!("error sending tts request: {e}"))?; + + request_ack_rx + .await + .map_err(|e| anyhow::anyhow!("TTS pool dropped request acknowledgement: {e}"))??; + + Ok(tts_resp_rx) +} + pub enum TTSSession { GsvStable { config: GSVTTS, @@ -122,90 +141,352 @@ impl TTSSession { } } -pub struct TTSSessionPool { - pub config: crate::config::TTSConfig, - pub workers: usize, - pub pool: tokio::sync::mpsc::UnboundedReceiver>, - pub tx: tokio::sync::mpsc::UnboundedSender>, +/// Default number of sessions kept ready while the pool is idle. +pub const DEFAULT_TTS_IDLE_WORKERS: usize = 1; + +/// Default upper bound for concurrently leased TTS sessions. +pub const DEFAULT_TTS_MAX_WORKERS: usize = 4; + +/// Default time after which an idle session is considered stale and reaped. +pub const DEFAULT_TTS_IDLE_TIMEOUT: Duration = Duration::from_secs(300); + +/// Default hard cap on how long a session may live before being discarded on +/// the next recycle check. +pub const DEFAULT_TTS_MAX_LIFETIME: Duration = Duration::from_secs(30 * 60); + +/// Pooled entry wrapping a `TTSSession` together with the timestamps used to +/// enforce idle / max-lifetime eviction. +pub struct TtsSessionEntry { + pub session: TTSSession, + pub created_at: Instant, + pub last_used: Mutex, } -impl TTSSessionPool { - pub fn new(config: crate::config::TTSConfig, workers: usize) -> Self { - let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); - TTSSessionPool { +impl TtsSessionEntry { + pub fn new(session: TTSSession) -> Self { + let now = Instant::now(); + Self { + session, + created_at: now, + last_used: Mutex::new(now), + } + } + + pub fn touch(&self) { + if let Ok(mut last) = self.last_used.lock() { + *last = Instant::now(); + } + } +} + +pub struct TTSManager { + config: crate::config::TTSConfig, + idle_timeout: Duration, + max_lifetime: Duration, +} + +impl TTSManager { + fn new( + config: crate::config::TTSConfig, + idle_timeout: Duration, + max_lifetime: Duration, + ) -> Self { + Self { config, - workers, - pool: rx, - tx, + idle_timeout, + max_lifetime, } } +} + +#[async_trait::async_trait] +impl deadpool::managed::Manager for TTSManager { + type Type = TtsSessionEntry; + type Error = anyhow::Error; - pub async fn create_session(&self) -> anyhow::Result { - TTSSession::new_from_config(&self.config).await + async fn create(&self) -> Result { + let session = TTSSession::new_from_config(&self.config).await?; + Ok(TtsSessionEntry::new(session)) } - pub async fn run_session( - id: u128, - mut session: TTSSession, - tx: tokio::sync::mpsc::UnboundedSender>, - ) -> anyhow::Result<()> { - log::info!("{} starting TTS session worker", id); - loop { - let (resp_tx, resp_rx) = tokio::sync::oneshot::channel(); - tx.send(resp_tx) - .map_err(|e| anyhow::anyhow!("send session request error: {}", e))?; + async fn recycle( + &self, + obj: &mut TtsSessionEntry, + _metrics: &deadpool::managed::Metrics, + ) -> deadpool::managed::RecycleResult { + let now = Instant::now(); + let age = now.saturating_duration_since(obj.created_at); + if age >= self.max_lifetime { + log::info!("TTS session exceeded max lifetime ({:?}); discarding", age); + return Err(deadpool::managed::RecycleError::Message( + "max lifetime exceeded".to_string(), + )); + } - let (text, tts_resp_tx) = resp_rx - .await - .map_err(|e| anyhow::anyhow!("receive session request error: {}", e))?; + let last_used = *obj + .last_used + .lock() + .map_err(|e| deadpool::managed::RecycleError::Message(e.to_string()))?; + let idle = now.saturating_duration_since(last_used); + if idle >= self.idle_timeout { + log::info!("TTS session idle for {:?}; discarding", idle); + return Err(deadpool::managed::RecycleError::Message( + "idle timeout exceeded".to_string(), + )); + } + + Ok(()) + } +} - log::info!("{} processing TTS request: {}", id, text); +pub struct TTSSessionPool { + pool: deadpool::managed::Pool, + idle_workers: usize, + idle_timeout: Duration, +} - if let Err(e) = session.synthesize(&text, &tts_resp_tx).await { - log::error!("{} TTS synthesis error: {}", id, e); - } +impl TTSSessionPool { + pub fn new( + config: crate::config::TTSConfig, + idle_workers: usize, + max_workers: usize, + idle_timeout: Duration, + ) -> Self { + Self::with_timeouts( + config, + idle_workers, + max_workers, + idle_timeout, + DEFAULT_TTS_MAX_LIFETIME, + ) + } + + pub fn with_timeouts( + config: crate::config::TTSConfig, + idle_workers: usize, + max_workers: usize, + idle_timeout: Duration, + max_lifetime: Duration, + ) -> Self { + let max_workers = max_workers.max(idle_workers); + let manager = TTSManager::new(config, idle_timeout, max_lifetime); + let pool = deadpool::managed::Pool::builder(manager) + .max_size(max_workers) + .timeouts(deadpool::managed::Timeouts { + wait: Some(Duration::from_secs(30)), + create: Some(Duration::from_secs(30)), + // Cap how long a single recycle() call may take; should be + // cheap but a stuck call must not stall pool.get(). + recycle: Some(Duration::from_secs(5)), + }) + .runtime(deadpool::Runtime::Tokio1) + .build() + .expect("Failed to create TTS session pool"); + TTSSessionPool { + pool, + idle_workers, + idle_timeout, } } - async fn get_req_tx(&mut self) -> anyhow::Result> { - let req_tx = self - .pool - .recv() - .await - .ok_or_else(|| anyhow::anyhow!("no available tts session"))?; - Ok(req_tx) + async fn prewarm(&self) -> anyhow::Result<()> { + let mut entries = Vec::with_capacity(self.idle_workers); + for worker in 0..self.idle_workers { + let entry = self + .pool + .get() + .await + .map_err(|e| anyhow::anyhow!("create idle TTS session[{worker}] error: {e}"))?; + entries.push(entry); + } + drop(entries); + Ok(()) } - pub async fn run_loop(&mut self, mut rx: TTSRequestRx) -> anyhow::Result<()> { - let mut sucess_workers = 0; - for i in 0..self.workers { - match self.create_session().await { - Ok(session) => { - tokio::spawn(Self::run_session(i as u128, session, self.tx.clone())); - sucess_workers += 1; - } - Err(e) => { - log::error!("create tts session[{i}] error: {}", e); - continue; - } - }; + fn reap_idle(&self) { + let status = self.pool.status(); + let removable = Cell::new(status.available.saturating_sub(self.idle_workers)); + if removable.get() == 0 { + return; } - if sucess_workers == 0 { - return Err(anyhow::anyhow!("no available tts session worker")); + let now = Instant::now(); + let idle_timeout = self.idle_timeout; + self.pool.retain(|entry, _| { + let remaining = removable.get(); + if remaining == 0 { + return true; + } + + let idle = entry + .last_used + .lock() + .map(|last_used| now.saturating_duration_since(*last_used) >= idle_timeout) + .unwrap_or(true); + if idle { + removable.set(remaining - 1); + false + } else { + true + } + }); + } + + pub async fn run_loop(&mut self, mut rx: TTSRequestRx) -> anyhow::Result<()> { + if let Err(e) = self.prewarm().await { + log::warn!("initial TTS pool prewarm failed; retrying on demand: {e}"); } - while let Some(tts_req) = rx.recv().await { - let req_tx = self.get_req_tx().await?; + let reap_period = self.idle_timeout.max(Duration::from_millis(1)); + let mut reap_interval = tokio::time::interval(reap_period); + reap_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - if let Err(e) = req_tx.send(tts_req) { - log::error!("send tts request to session error: {}", e.0); + loop { + tokio::select! { + request = rx.recv() => { + let Some((text, tts_resp_tx, request_ack_tx)) = request else { + break; + }; + + match self.pool.get().await { + Ok(mut entry) => { + let _ = request_ack_tx.send(Ok(())); + tokio::spawn(async move { + log::info!("Processing TTS request: {}", text); + let result = entry.session.synthesize(&text, &tts_resp_tx).await; + entry.touch(); + if let Err(e) = result { + log::error!("TTS synthesis error: {}", e); + } + }); + } + Err(e) => { + let message = format!("Failed to get TTS session from pool: {e}"); + log::error!("{message}"); + let _ = request_ack_tx.send(Err(anyhow::anyhow!(message))); + } + } + } + _ = reap_interval.tick() => self.reap_idle(), } } Ok(()) } } +#[cfg(test)] +mod tests { + use super::*; + use crate::config::{GSVTTS, TTSConfig}; + + fn test_config() -> TTSConfig { + TTSConfig::GSV(GSVTTS { + api_key: String::new(), + url: String::new(), + speaker: String::new(), + timeout_sec: None, + text_optimization: None, + }) + } + + #[tokio::test] + async fn prewarms_idle_workers_and_scales_to_max_workers() { + let pool = TTSSessionPool::with_timeouts( + test_config(), + 2, + 3, + Duration::from_secs(60), + Duration::from_secs(60), + ); + + pool.prewarm().await.unwrap(); + assert_eq!(pool.pool.status().size, 2); + + let (first, second, third) = + tokio::join!(pool.pool.get(), pool.pool.get(), pool.pool.get(),); + assert!(first.is_ok()); + assert!(second.is_ok()); + assert!(third.is_ok()); + assert_eq!(pool.pool.status().size, 3); + } + + #[tokio::test] + async fn reaper_keeps_idle_workers_and_removes_excess_idle_sessions() { + let pool = TTSSessionPool::with_timeouts( + test_config(), + 1, + 3, + Duration::from_millis(10), + Duration::from_secs(60), + ); + let (first, second, third) = + tokio::join!(pool.pool.get(), pool.pool.get(), pool.pool.get(),); + drop((first, second, third)); + + tokio::time::sleep(Duration::from_millis(20)).await; + pool.reap_idle(); + assert_eq!(pool.pool.status().size, 1); + } + + #[tokio::test] + async fn touching_a_session_after_use_prevents_premature_idle_reaping() { + let pool = TTSSessionPool::with_timeouts( + test_config(), + 0, + 1, + Duration::from_millis(100), + Duration::from_secs(60), + ); + let entry = pool.pool.get().await.unwrap(); + tokio::time::sleep(Duration::from_millis(20)).await; + entry.touch(); + drop(entry); + + tokio::time::sleep(Duration::from_millis(20)).await; + pool.reap_idle(); + assert_eq!(pool.pool.status().size, 1); + } + + #[tokio::test] + async fn recycle_replaces_expired_sessions() { + let pool = TTSSessionPool::with_timeouts( + test_config(), + 0, + 1, + Duration::from_secs(60), + Duration::from_millis(10), + ); + let entry = pool.pool.get().await.unwrap(); + let created_at = entry.created_at; + drop(entry); + + tokio::time::sleep(Duration::from_millis(20)).await; + let replacement = pool.pool.get().await.unwrap(); + assert!(replacement.created_at > created_at); + } + + #[tokio::test] + async fn submit_request_propagates_pool_acquisition_errors() { + let (tts_tx, mut tts_rx) = tokio::sync::mpsc::channel(1); + let request = + tokio::spawn(async move { submit_request(&tts_tx, "hello".to_string()).await }); + + let (_, _, request_ack_tx) = tts_rx.recv().await.unwrap(); + request_ack_tx + .send(Err(anyhow::anyhow!("pool unavailable"))) + .unwrap(); + + let result = request.await.unwrap(); + assert!( + result + .err() + .unwrap() + .to_string() + .contains("pool unavailable") + ); + } +} + async fn retry_gsv_tts( client: &reqwest::Client, url: &str,