fix(feishu): bound shutdown during connection

This commit is contained in:
xiaoxixi 2026-07-14 12:43:02 +08:00
parent 3f1350c33b
commit 18f1e47f77

View File

@ -8,8 +8,9 @@ use futures_util::{SinkExt, StreamExt};
use prost::{Message as ProstMessage, bytes::Bytes}; use prost::{Message as ProstMessage, bytes::Bytes};
use regex::Regex; use regex::Regex;
use serde::Deserialize; use serde::Deserialize;
use tokio::sync::{Mutex, RwLock, broadcast}; use tokio::sync::{Mutex, RwLock};
use tokio::task::JoinHandle; use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use crate::bus::{MediaItem, MessageBus, OutboundMessage}; use crate::bus::{MediaItem, MessageBus, OutboundMessage};
use crate::channels::base::{Channel, ChannelError}; 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). /// 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. /// If no binary frame (pong or event) is received within this window, reconnect.
const WS_HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(300); 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. /// Refresh tenant token this many seconds before the announced expiry.
const TOKEN_REFRESH_SKEW: Duration = Duration::from_secs(120); const TOKEN_REFRESH_SKEW: Duration = Duration::from_secs(120);
/// Default tenant token TTL when `expire`/`expires_in` is absent. /// Default tenant token TTL when `expire`/`expires_in` is absent.
@ -145,7 +151,7 @@ pub struct FeishuChannel {
config: FeishuChannelConfig, config: FeishuChannelConfig,
http_client: reqwest::Client, http_client: reqwest::Client,
running: Arc<RwLock<bool>>, running: Arc<RwLock<bool>>,
shutdown_tx: Arc<RwLock<Option<broadcast::Sender<()>>>>, shutdown: Arc<RwLock<Option<CancellationToken>>>,
run_task: Arc<Mutex<Option<JoinHandle<()>>>>, run_task: Arc<Mutex<Option<JoinHandle<()>>>>,
connected: Arc<RwLock<bool>>, connected: Arc<RwLock<bool>>,
/// Cached tenant access token with proactive refresh. /// Cached tenant access token with proactive refresh.
@ -179,7 +185,7 @@ impl FeishuChannel {
config, config,
http_client: reqwest::Client::new(), http_client: reqwest::Client::new(),
running: Arc::new(RwLock::new(false)), 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)), run_task: Arc::new(Mutex::new(None)),
connected: Arc::new(RwLock::new(false)), connected: Arc::new(RwLock::new(false)),
tenant_token: Arc::new(RwLock::new(None)), tenant_token: Arc::new(RwLock::new(None)),
@ -1166,18 +1172,22 @@ impl FeishuChannel {
async fn run_ws_loop( async fn run_ws_loop(
&self, &self,
bus: Arc<MessageBus>, bus: Arc<MessageBus>,
mut shutdown_rx: broadcast::Receiver<()>, shutdown: CancellationToken,
) -> Result<(), ChannelError> { ) -> 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); let service_id = Self::extract_service_id(&wss_url);
tracing::info!(url = %wss_url, "Connecting to Feishu WebSocket"); tracing::info!(url = %wss_url, "Connecting to Feishu WebSocket");
let (ws_stream, _) = tokio_tungstenite::connect_async(&wss_url) let (ws_stream, _) = tokio::select! {
.await result = tokio_tungstenite::connect_async(&wss_url) => result.map_err(|e| {
.map_err(|e| {
ChannelError::ConnectionError(format!("WebSocket connection failed: {}", e)) ChannelError::ConnectionError(format!("WebSocket connection failed: {}", e))
})?; })?,
_ = shutdown.cancelled() => return Ok(()),
};
*self.connected.write().await = true; *self.connected.write().await = true;
tracing::info!("Feishu WebSocket connected"); tracing::info!("Feishu WebSocket connected");
@ -1196,14 +1206,14 @@ impl FeishuChannel {
}], }],
payload: None, payload: None,
}; };
write tokio::select! {
.send(tokio_tungstenite::tungstenite::Message::Binary( result = write.send(tokio_tungstenite::tungstenite::Message::Binary(
ping_frame.encode_to_vec().into(), ping_frame.encode_to_vec().into(),
)) )) => result.map_err(|e| {
.await
.map_err(|e| {
ChannelError::ConnectionError(format!("Failed to send initial ping: {}", 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 ping_interval = client_config.ping_interval.unwrap_or(120).max(10);
let mut ping_interval_tok = let mut ping_interval_tok =
@ -1336,7 +1346,7 @@ impl FeishuChannel {
let mut seen = self.seen_message_ids.write().await; let mut seen = self.seen_message_ids.write().await;
seen.retain(|_, ts| now.duration_since(*ts) < DEDUP_CACHE_TTL); seen.retain(|_, ts| now.duration_since(*ts) < DEDUP_CACHE_TTL);
} }
_ = shutdown_rx.recv() => { _ = shutdown.cancelled() => {
tracing::info!("Feishu channel shutdown signal received"); tracing::info!("Feishu channel shutdown signal received");
break; break;
} }
@ -1954,8 +1964,8 @@ impl Channel for FeishuChannel {
*self.running.write().await = true; *self.running.write().await = true;
let (shutdown_tx, _) = broadcast::channel(1); let shutdown = CancellationToken::new();
*self.shutdown_tx.write().await = Some(shutdown_tx.clone()); *self.shutdown.write().await = Some(shutdown.clone());
let channel = self.clone(); let channel = self.clone();
let bus = bus.clone(); let bus = bus.clone();
@ -1968,8 +1978,7 @@ impl Channel for FeishuChannel {
break; break;
} }
let shutdown_rx = shutdown_tx.subscribe(); match channel.run_ws_loop(bus.clone(), shutdown.clone()).await {
match channel.run_ws_loop(bus.clone(), shutdown_rx).await {
Ok(_) => { Ok(_) => {
tracing::info!("Feishu WebSocket disconnected"); 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; break;
} }
tracing::info!("Feishu channel retrying in 5s..."); 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; *channel.running.write().await = false;
@ -2001,15 +2013,27 @@ impl Channel for FeishuChannel {
async fn stop(&self) -> Result<(), ChannelError> { async fn stop(&self) -> Result<(), ChannelError> {
*self.running.write().await = false; *self.running.write().await = false;
*self.connected.write().await = false;
if let Some(tx) = self.shutdown_tx.write().await.take() { if let Some(shutdown) = self.shutdown.write().await.take() {
let _ = tx.send(()); shutdown.cancel();
} }
if let Some(task) = self.run_task.lock().await.take() { let task = { self.run_task.lock().await.take() };
task.await.map_err(|error| { if let Some(mut task) = task {
ChannelError::Other(format!("Feishu channel task failed to join: {error}")) 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(()) Ok(())
@ -2229,6 +2253,43 @@ impl Channel for FeishuChannel {
mod tests { mod tests {
use super::*; 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] #[test]
fn collect_post_image_keys_finds_nested_images() { fn collect_post_image_keys_finds_nested_images() {
let content = serde_json::json!({ let content = serde_json::json!({