From 18f1e47f77048e272b61d6a2f84a4d2d4593efbf Mon Sep 17 00:00:00 2001 From: xiaoxixi Date: Tue, 14 Jul 2026 12:43:02 +0800 Subject: [PATCH] fix(feishu): bound shutdown during connection --- src/channels/feishu.rs | 117 +++++++++++++++++++++++++++++++---------- 1 file changed, 89 insertions(+), 28 deletions(-) diff --git a/src/channels/feishu.rs b/src/channels/feishu.rs index b3ec68d..9f6b354 100644 --- a/src/channels/feishu.rs +++ b/src/channels/feishu.rs @@ -8,8 +8,9 @@ use futures_util::{SinkExt, StreamExt}; use prost::{Message as ProstMessage, bytes::Bytes}; use regex::Regex; use serde::Deserialize; -use tokio::sync::{Mutex, RwLock, broadcast}; +use tokio::sync::{Mutex, RwLock}; use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; use crate::bus::{MediaItem, MessageBus, OutboundMessage}; use crate::channels::base::{Channel, ChannelError}; @@ -21,6 +22,11 @@ const FEISHU_WS_BASE: &str = "https://open.feishu.cn"; /// Heartbeat timeout for WS connection — must be larger than ping_interval (default 120 s). /// If no binary frame (pong or event) is received within this window, reconnect. const WS_HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(300); +const CHANNEL_STOP_GRACE: Duration = if cfg!(test) { + Duration::from_millis(100) +} else { + Duration::from_secs(5) +}; /// Refresh tenant token this many seconds before the announced expiry. const TOKEN_REFRESH_SKEW: Duration = Duration::from_secs(120); /// Default tenant token TTL when `expire`/`expires_in` is absent. @@ -145,7 +151,7 @@ pub struct FeishuChannel { config: FeishuChannelConfig, http_client: reqwest::Client, running: Arc>, - shutdown_tx: Arc>>>, + shutdown: Arc>>, run_task: Arc>>>, connected: Arc>, /// Cached tenant access token with proactive refresh. @@ -179,7 +185,7 @@ impl FeishuChannel { config, http_client: reqwest::Client::new(), running: Arc::new(RwLock::new(false)), - shutdown_tx: Arc::new(RwLock::new(None)), + shutdown: Arc::new(RwLock::new(None)), run_task: Arc::new(Mutex::new(None)), connected: Arc::new(RwLock::new(false)), tenant_token: Arc::new(RwLock::new(None)), @@ -1166,18 +1172,22 @@ impl FeishuChannel { async fn run_ws_loop( &self, bus: Arc, - mut shutdown_rx: broadcast::Receiver<()>, + shutdown: CancellationToken, ) -> Result<(), ChannelError> { - let (wss_url, client_config) = self.get_ws_endpoint(&self.http_client).await?; + let (wss_url, client_config) = tokio::select! { + result = self.get_ws_endpoint(&self.http_client) => result?, + _ = shutdown.cancelled() => return Ok(()), + }; let service_id = Self::extract_service_id(&wss_url); tracing::info!(url = %wss_url, "Connecting to Feishu WebSocket"); - let (ws_stream, _) = tokio_tungstenite::connect_async(&wss_url) - .await - .map_err(|e| { + let (ws_stream, _) = tokio::select! { + result = tokio_tungstenite::connect_async(&wss_url) => result.map_err(|e| { ChannelError::ConnectionError(format!("WebSocket connection failed: {}", e)) - })?; + })?, + _ = shutdown.cancelled() => return Ok(()), + }; *self.connected.write().await = true; tracing::info!("Feishu WebSocket connected"); @@ -1196,14 +1206,14 @@ impl FeishuChannel { }], payload: None, }; - write - .send(tokio_tungstenite::tungstenite::Message::Binary( + tokio::select! { + result = write.send(tokio_tungstenite::tungstenite::Message::Binary( ping_frame.encode_to_vec().into(), - )) - .await - .map_err(|e| { + )) => result.map_err(|e| { ChannelError::ConnectionError(format!("Failed to send initial ping: {}", e)) - })?; + })?, + _ = shutdown.cancelled() => return Ok(()), + }; let ping_interval = client_config.ping_interval.unwrap_or(120).max(10); let mut ping_interval_tok = @@ -1336,7 +1346,7 @@ impl FeishuChannel { let mut seen = self.seen_message_ids.write().await; seen.retain(|_, ts| now.duration_since(*ts) < DEDUP_CACHE_TTL); } - _ = shutdown_rx.recv() => { + _ = shutdown.cancelled() => { tracing::info!("Feishu channel shutdown signal received"); break; } @@ -1954,8 +1964,8 @@ impl Channel for FeishuChannel { *self.running.write().await = true; - let (shutdown_tx, _) = broadcast::channel(1); - *self.shutdown_tx.write().await = Some(shutdown_tx.clone()); + let shutdown = CancellationToken::new(); + *self.shutdown.write().await = Some(shutdown.clone()); let channel = self.clone(); let bus = bus.clone(); @@ -1968,8 +1978,7 @@ impl Channel for FeishuChannel { break; } - let shutdown_rx = shutdown_tx.subscribe(); - match channel.run_ws_loop(bus.clone(), shutdown_rx).await { + match channel.run_ws_loop(bus.clone(), shutdown.clone()).await { Ok(_) => { tracing::info!("Feishu WebSocket disconnected"); } @@ -1983,12 +1992,15 @@ impl Channel for FeishuChannel { } } - if !*channel.running.read().await { + if !*channel.running.read().await || shutdown.is_cancelled() { break; } tracing::info!("Feishu channel retrying in 5s..."); - tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; + tokio::select! { + _ = tokio::time::sleep(tokio::time::Duration::from_secs(5)) => {} + _ = shutdown.cancelled() => break, + } } *channel.running.write().await = false; @@ -2001,15 +2013,27 @@ impl Channel for FeishuChannel { async fn stop(&self) -> Result<(), ChannelError> { *self.running.write().await = false; + *self.connected.write().await = false; - if let Some(tx) = self.shutdown_tx.write().await.take() { - let _ = tx.send(()); + if let Some(shutdown) = self.shutdown.write().await.take() { + shutdown.cancel(); } - if let Some(task) = self.run_task.lock().await.take() { - task.await.map_err(|error| { - ChannelError::Other(format!("Feishu channel task failed to join: {error}")) - })?; + let task = { self.run_task.lock().await.take() }; + if let Some(mut task) = task { + match tokio::time::timeout(CHANNEL_STOP_GRACE, &mut task).await { + Ok(result) => result.map_err(|error| { + ChannelError::Other(format!("Feishu channel task failed to join: {error}")) + })?, + Err(_) => { + tracing::warn!( + grace_ms = CHANNEL_STOP_GRACE.as_millis(), + "Feishu channel did not stop in time; aborting connection task" + ); + task.abort(); + let _ = task.await; + } + } } Ok(()) @@ -2229,6 +2253,43 @@ impl Channel for FeishuChannel { mod tests { use super::*; + fn test_channel() -> FeishuChannel { + FeishuChannel::new( + FeishuChannelConfig { + enabled: true, + app_id: "test-app".to_string(), + app_secret: "test-secret".to_string(), + allow_from: vec!["*".to_string()], + agent: String::new(), + media_dir: String::new(), + reaction_emoji: "THUMBSUP".to_string(), + }, + Path::new("/tmp"), + ) + .expect("test channel should be valid") + } + + #[tokio::test] + async fn stop_aborts_connection_task_that_ignores_cancellation() { + let channel = test_channel(); + let shutdown = CancellationToken::new(); + + *channel.running.write().await = true; + *channel.connected.write().await = true; + *channel.shutdown.write().await = Some(shutdown.clone()); + *channel.run_task.lock().await = Some(tokio::spawn(std::future::pending())); + + tokio::time::timeout(Duration::from_secs(1), channel.stop()) + .await + .expect("stop must have a hard deadline") + .expect("stop should succeed after aborting the stuck task"); + + assert!(shutdown.is_cancelled()); + assert!(!channel.is_running()); + assert!(!*channel.connected.read().await); + assert!(channel.run_task.lock().await.is_none()); + } + #[test] fn collect_post_image_keys_finds_nested_images() { let content = serde_json::json!({