fix(feishu): bound shutdown during connection
This commit is contained in:
parent
3f1350c33b
commit
18f1e47f77
@ -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 {
|
||||||
|
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}"))
|
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!({
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user