fix(channels): unify lifecycle ownership
This commit is contained in:
parent
63d20d1eb8
commit
a9c297764e
@ -8,7 +8,8 @@ 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::{RwLock, broadcast};
|
use tokio::sync::{Mutex, RwLock, broadcast};
|
||||||
|
use tokio::task::JoinHandle;
|
||||||
|
|
||||||
use crate::bus::{MediaItem, MessageBus, OutboundMessage};
|
use crate::bus::{MediaItem, MessageBus, OutboundMessage};
|
||||||
use crate::channels::base::{Channel, ChannelError};
|
use crate::channels::base::{Channel, ChannelError};
|
||||||
@ -152,6 +153,7 @@ pub struct FeishuChannel {
|
|||||||
http_client: reqwest::Client,
|
http_client: reqwest::Client,
|
||||||
running: Arc<RwLock<bool>>,
|
running: Arc<RwLock<bool>>,
|
||||||
shutdown_tx: Arc<RwLock<Option<broadcast::Sender<()>>>>,
|
shutdown_tx: Arc<RwLock<Option<broadcast::Sender<()>>>>,
|
||||||
|
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.
|
||||||
tenant_token: Arc<RwLock<Option<CachedTenantToken>>>,
|
tenant_token: Arc<RwLock<Option<CachedTenantToken>>>,
|
||||||
@ -185,6 +187,7 @@ impl FeishuChannel {
|
|||||||
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_tx: Arc::new(RwLock::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)),
|
||||||
seen_message_ids: Arc::new(RwLock::new(HashMap::new())),
|
seen_message_ids: Arc::new(RwLock::new(HashMap::new())),
|
||||||
@ -1953,6 +1956,14 @@ impl Channel for FeishuChannel {
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
let mut run_task = self.run_task.lock().await;
|
||||||
|
if run_task.as_ref().is_some_and(|task| !task.is_finished()) {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
if let Some(finished) = run_task.take() {
|
||||||
|
let _ = finished.await;
|
||||||
|
}
|
||||||
|
|
||||||
*self.running.write().await = true;
|
*self.running.write().await = true;
|
||||||
|
|
||||||
let (shutdown_tx, _) = broadcast::channel(1);
|
let (shutdown_tx, _) = broadcast::channel(1);
|
||||||
@ -1960,7 +1971,7 @@ impl Channel for FeishuChannel {
|
|||||||
|
|
||||||
let channel = self.clone();
|
let channel = self.clone();
|
||||||
let bus = bus.clone();
|
let bus = bus.clone();
|
||||||
tokio::spawn(async move {
|
*run_task = Some(tokio::spawn(async move {
|
||||||
let mut consecutive_failures = 0;
|
let mut consecutive_failures = 0;
|
||||||
let max_failures = 3;
|
let max_failures = 3;
|
||||||
|
|
||||||
@ -1994,7 +2005,7 @@ impl Channel for FeishuChannel {
|
|||||||
|
|
||||||
*channel.running.write().await = false;
|
*channel.running.write().await = false;
|
||||||
tracing::info!("Feishu channel stopped");
|
tracing::info!("Feishu channel stopped");
|
||||||
});
|
}));
|
||||||
|
|
||||||
tracing::info!("Feishu channel started");
|
tracing::info!("Feishu channel started");
|
||||||
Ok(())
|
Ok(())
|
||||||
@ -2007,6 +2018,12 @@ impl Channel for FeishuChannel {
|
|||||||
let _ = tx.send(());
|
let _ = tx.send(());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
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}"))
|
||||||
|
})?;
|
||||||
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -2,7 +2,7 @@ use std::collections::HashMap;
|
|||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tokio::sync::RwLock;
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
use crate::bus::{MessageBus, OutboundMessage};
|
use crate::bus::MessageBus;
|
||||||
use crate::channels::base::{Channel, ChannelError};
|
use crate::channels::base::{Channel, ChannelError};
|
||||||
use crate::channels::feishu::FeishuChannel;
|
use crate::channels::feishu::FeishuChannel;
|
||||||
use crate::config::Config;
|
use crate::config::Config;
|
||||||
@ -16,14 +16,6 @@ pub struct ChannelManager {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl ChannelManager {
|
impl ChannelManager {
|
||||||
pub fn new(cli_chat_channel: Arc<crate::channels::CliChatChannel>) -> Self {
|
|
||||||
Self {
|
|
||||||
channels: Arc::new(RwLock::new(HashMap::new())),
|
|
||||||
cli_chat_channel,
|
|
||||||
bus: MessageBus::new(100),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn with_bus(
|
pub fn with_bus(
|
||||||
cli_chat_channel: Arc<crate::channels::CliChatChannel>,
|
cli_chat_channel: Arc<crate::channels::CliChatChannel>,
|
||||||
bus: Arc<MessageBus>,
|
bus: Arc<MessageBus>,
|
||||||
@ -91,27 +83,53 @@ impl ChannelManager {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub async fn start_all(&self) -> Result<(), ChannelError> {
|
pub async fn start_all(&self) -> Result<(), ChannelError> {
|
||||||
let channels = self.channels.read().await;
|
let channels: Vec<_> = self
|
||||||
|
.channels
|
||||||
|
.read()
|
||||||
|
.await
|
||||||
|
.iter()
|
||||||
|
.map(|(name, channel)| (name.clone(), channel.clone()))
|
||||||
|
.collect();
|
||||||
let bus = self.bus.clone();
|
let bus = self.bus.clone();
|
||||||
for (name, channel) in channels.iter() {
|
let mut failures = Vec::new();
|
||||||
|
for (name, channel) in channels {
|
||||||
tracing::info!(channel = %name, "Starting channel");
|
tracing::info!(channel = %name, "Starting channel");
|
||||||
if let Err(e) = channel.start(bus.clone()).await {
|
if let Err(e) = channel.start(bus.clone()).await {
|
||||||
tracing::error!(channel = %name, error = %e, "Failed to start channel");
|
tracing::error!(channel = %name, error = %e, "Failed to start channel");
|
||||||
|
failures.push(format!("{name}: {e}"));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
if failures.is_empty() {
|
||||||
Ok(())
|
Ok(())
|
||||||
|
} else {
|
||||||
|
Err(ChannelError::Other(format!(
|
||||||
|
"failed to start channels: {}",
|
||||||
|
failures.join("; ")
|
||||||
|
)))
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn stop_all(&self) -> Result<(), ChannelError> {
|
pub async fn stop_all(&self) -> Result<(), ChannelError> {
|
||||||
let mut channels = self.channels.write().await;
|
let channels: Vec<_> = {
|
||||||
for (name, channel) in channels.iter() {
|
let mut registered = self.channels.write().await;
|
||||||
|
registered.drain().collect()
|
||||||
|
};
|
||||||
|
let mut failures = Vec::new();
|
||||||
|
for (name, channel) in channels {
|
||||||
tracing::info!(channel = %name, "Stopping channel");
|
tracing::info!(channel = %name, "Stopping channel");
|
||||||
if let Err(e) = channel.stop().await {
|
if let Err(e) = channel.stop().await {
|
||||||
tracing::error!(channel = %name, error = %e, "Error stopping channel");
|
tracing::error!(channel = %name, error = %e, "Error stopping channel");
|
||||||
|
failures.push(format!("{name}: {e}"));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
channels.clear();
|
if failures.is_empty() {
|
||||||
Ok(())
|
Ok(())
|
||||||
|
} else {
|
||||||
|
Err(ChannelError::Other(format!(
|
||||||
|
"failed to stop channels: {}",
|
||||||
|
failures.join("; ")
|
||||||
|
)))
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn get_channel(&self, name: &str) -> Option<Arc<dyn Channel + Send + Sync>> {
|
pub async fn get_channel(&self, name: &str) -> Option<Arc<dyn Channel + Send + Sync>> {
|
||||||
@ -122,17 +140,104 @@ impl ChannelManager {
|
|||||||
pub async fn list_channel_names(&self) -> Vec<String> {
|
pub async fn list_channel_names(&self) -> Vec<String> {
|
||||||
self.channels.read().await.keys().cloned().collect()
|
self.channels.read().await.keys().cloned().collect()
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Dispatch an outbound message to the appropriate channel
|
#[cfg(test)]
|
||||||
pub async fn dispatch(&self, msg: OutboundMessage) -> Result<(), ChannelError> {
|
mod tests {
|
||||||
let channel_name = &msg.channel;
|
use super::*;
|
||||||
if let Some(channel) = self.get_channel(channel_name).await {
|
use async_trait::async_trait;
|
||||||
channel.send(msg).await
|
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||||
} else {
|
|
||||||
Err(ChannelError::Other(format!(
|
struct TestChannel {
|
||||||
"Channel not found: {}",
|
name: &'static str,
|
||||||
channel_name
|
fail_start: bool,
|
||||||
)))
|
fail_stop: bool,
|
||||||
|
running: AtomicBool,
|
||||||
|
starts: AtomicUsize,
|
||||||
|
stops: AtomicUsize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TestChannel {
|
||||||
|
fn new(name: &'static str, fail_start: bool, fail_stop: bool) -> Self {
|
||||||
|
Self {
|
||||||
|
name,
|
||||||
|
fail_start,
|
||||||
|
fail_stop,
|
||||||
|
running: AtomicBool::new(false),
|
||||||
|
starts: AtomicUsize::new(0),
|
||||||
|
stops: AtomicUsize::new(0),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Channel for TestChannel {
|
||||||
|
fn name(&self) -> &str {
|
||||||
|
self.name
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_running(&self) -> bool {
|
||||||
|
self.running.load(Ordering::SeqCst)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn start(&self, _bus: Arc<MessageBus>) -> Result<(), ChannelError> {
|
||||||
|
self.starts.fetch_add(1, Ordering::SeqCst);
|
||||||
|
if self.fail_start {
|
||||||
|
return Err(ChannelError::ConnectionError("unavailable".into()));
|
||||||
|
}
|
||||||
|
self.running.store(true, Ordering::SeqCst);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn stop(&self) -> Result<(), ChannelError> {
|
||||||
|
self.stops.fetch_add(1, Ordering::SeqCst);
|
||||||
|
self.running.store(false, Ordering::SeqCst);
|
||||||
|
if self.fail_stop {
|
||||||
|
return Err(ChannelError::Other("stop failed".into()));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn send(&self, _msg: crate::bus::OutboundMessage) -> Result<(), ChannelError> {
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn manager() -> ChannelManager {
|
||||||
|
ChannelManager::with_bus(
|
||||||
|
Arc::new(crate::channels::CliChatChannel::new()),
|
||||||
|
MessageBus::new(8),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn start_all_reports_channel_failures_and_starts_the_rest() {
|
||||||
|
let manager = manager();
|
||||||
|
let healthy = Arc::new(TestChannel::new("healthy", false, false));
|
||||||
|
let broken = Arc::new(TestChannel::new("broken", true, false));
|
||||||
|
manager.register_channel("healthy", healthy.clone()).await;
|
||||||
|
manager.register_channel("broken", broken.clone()).await;
|
||||||
|
|
||||||
|
let error = manager.start_all().await.unwrap_err().to_string();
|
||||||
|
|
||||||
|
assert!(error.contains("broken"));
|
||||||
|
assert_eq!(healthy.starts.load(Ordering::SeqCst), 1);
|
||||||
|
assert_eq!(broken.starts.load(Ordering::SeqCst), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn stop_all_reports_failures_and_unregisters_every_channel() {
|
||||||
|
let manager = manager();
|
||||||
|
let healthy = Arc::new(TestChannel::new("healthy", false, false));
|
||||||
|
let broken = Arc::new(TestChannel::new("broken", false, true));
|
||||||
|
manager.register_channel("healthy", healthy.clone()).await;
|
||||||
|
manager.register_channel("broken", broken.clone()).await;
|
||||||
|
|
||||||
|
let error = manager.stop_all().await.unwrap_err().to_string();
|
||||||
|
|
||||||
|
assert!(error.contains("broken"));
|
||||||
|
assert_eq!(healthy.stops.load(Ordering::SeqCst), 1);
|
||||||
|
assert_eq!(broken.stops.load(Ordering::SeqCst), 1);
|
||||||
|
assert!(manager.list_channel_names().await.is_empty());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@ -6,7 +6,7 @@ use std::sync::Arc;
|
|||||||
use tokio::net::TcpListener;
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
use crate::bus::{ControlMessage, MessageBus, OutboundDispatcher};
|
use crate::bus::{ControlMessage, MessageBus, OutboundDispatcher};
|
||||||
use crate::channels::base::{Channel, ChannelError};
|
use crate::channels::base::ChannelError;
|
||||||
use crate::channels::{ChannelManager, CliChatChannel};
|
use crate::channels::{ChannelManager, CliChatChannel};
|
||||||
use crate::config::{Config, ensure_workspace_dir, expand_path};
|
use crate::config::{Config, ensure_workspace_dir, expand_path};
|
||||||
use crate::logging;
|
use crate::logging;
|
||||||
@ -198,12 +198,6 @@ impl GatewayState {
|
|||||||
let bus_for_outbound = bus.clone();
|
let bus_for_outbound = bus.clone();
|
||||||
let session_manager = self.session_manager.clone();
|
let session_manager = self.session_manager.clone();
|
||||||
|
|
||||||
// Start CLI Chat Channel (it's already registered in ChannelManager)
|
|
||||||
let cli_chat_channel = self.cli_chat_channel();
|
|
||||||
if let Err(e) = cli_chat_channel.start(bus.clone()).await {
|
|
||||||
tracing::error!(error = %e, "Failed to start CLI chat channel");
|
|
||||||
}
|
|
||||||
|
|
||||||
// Spawn unified message processor
|
// Spawn unified message processor
|
||||||
// This handles both inbound AI messages and control messages in one loop
|
// This handles both inbound AI messages and control messages in one loop
|
||||||
self.task_supervisor.spawn("message-processor", async move {
|
self.task_supervisor.spawn("message-processor", async move {
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user