407 lines
14 KiB
Rust
407 lines
14 KiB
Rust
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use tokio::sync::mpsc;
|
|
|
|
use crate::bus::{MessageBus, OutboundMessage};
|
|
use crate::channels::ChannelManager;
|
|
use crate::channels::base::{Channel, ChannelError};
|
|
use crate::delivery::ConversationWriteLocks;
|
|
use crate::task_supervisor::TaskSupervisor;
|
|
|
|
const LANE_CAPACITY: usize = 64;
|
|
const LANE_IDLE_TIMEOUT: Duration = Duration::from_secs(300);
|
|
const SEND_TIMEOUT: Duration = Duration::from_secs(30);
|
|
|
|
/// Dispatches outbound messages through independent per-conversation lanes.
|
|
/// Messages to the same channel/chat remain ordered, while a slow destination
|
|
/// cannot block delivery to unrelated conversations.
|
|
pub struct OutboundDispatcher {
|
|
bus: Arc<MessageBus>,
|
|
channel_manager: ChannelManager,
|
|
task_supervisor: TaskSupervisor,
|
|
write_locks: ConversationWriteLocks,
|
|
}
|
|
|
|
impl OutboundDispatcher {
|
|
pub fn new(
|
|
bus: Arc<MessageBus>,
|
|
channel_manager: ChannelManager,
|
|
task_supervisor: TaskSupervisor,
|
|
write_locks: ConversationWriteLocks,
|
|
) -> Self {
|
|
Self {
|
|
bus,
|
|
channel_manager,
|
|
task_supervisor,
|
|
write_locks,
|
|
}
|
|
}
|
|
|
|
pub async fn run(&self) {
|
|
tracing::info!(lane_capacity = LANE_CAPACITY, "OutboundDispatcher started");
|
|
let mut lanes: HashMap<String, mpsc::Sender<OutboundMessage>> = HashMap::new();
|
|
let mut messages_seen = 0_u64;
|
|
|
|
loop {
|
|
let Some(msg) = self.bus.consume_outbound().await else {
|
|
tracing::warn!("OutboundDispatcher stopping because outbound bus closed");
|
|
break;
|
|
};
|
|
|
|
messages_seen = messages_seen.wrapping_add(1);
|
|
if messages_seen.is_multiple_of(128) {
|
|
lanes.retain(|_, sender| !sender.is_closed());
|
|
}
|
|
|
|
let lane_key = format!("{}\0{}", msg.channel, msg.chat_id);
|
|
let mut sender = lanes.get(&lane_key).cloned();
|
|
if sender.as_ref().is_none_or(mpsc::Sender::is_closed) {
|
|
let Some(channel) = self.channel_manager.get_channel(&msg.channel).await else {
|
|
tracing::warn!(channel = %msg.channel, "No channel found for message");
|
|
msg.complete_delivery(Err(format!("channel not found: {}", msg.channel)));
|
|
continue;
|
|
};
|
|
let (new_sender, receiver) = mpsc::channel(LANE_CAPACITY);
|
|
if !self.spawn_lane(channel, receiver, msg.channel.clone(), msg.chat_id.clone()) {
|
|
msg.complete_delivery(Err("dispatcher is shutting down".to_string()));
|
|
continue;
|
|
}
|
|
lanes.insert(lane_key.clone(), new_sender.clone());
|
|
sender = Some(new_sender);
|
|
}
|
|
|
|
let Some(sender) = sender else {
|
|
tracing::error!("Outbound lane creation did not produce a sender");
|
|
continue;
|
|
};
|
|
match sender.try_send(msg) {
|
|
Ok(()) => {}
|
|
Err(mpsc::error::TrySendError::Full(msg)) => {
|
|
tracing::error!(
|
|
channel = %msg.channel,
|
|
chat_id = %msg.chat_id,
|
|
capacity = LANE_CAPACITY,
|
|
"Outbound lane full; rejecting message instead of blocking other destinations"
|
|
);
|
|
msg.complete_delivery(Err("outbound lane is full".to_string()));
|
|
}
|
|
Err(mpsc::error::TrySendError::Closed(msg)) => {
|
|
// The lane may have expired between the closed check and
|
|
// enqueue. Recreate it once and preserve this message.
|
|
lanes.remove(&lane_key);
|
|
let Some(channel) = self.channel_manager.get_channel(&msg.channel).await else {
|
|
tracing::warn!(channel = %msg.channel, "No channel found for message");
|
|
msg.complete_delivery(Err(format!("channel not found: {}", msg.channel)));
|
|
continue;
|
|
};
|
|
let (new_sender, receiver) = mpsc::channel(LANE_CAPACITY);
|
|
if !self.spawn_lane(channel, receiver, msg.channel.clone(), msg.chat_id.clone())
|
|
{
|
|
msg.complete_delivery(Err("dispatcher is shutting down".to_string()));
|
|
continue;
|
|
}
|
|
match new_sender.try_send(msg) {
|
|
Ok(()) => {
|
|
lanes.insert(lane_key, new_sender);
|
|
}
|
|
Err(error) => {
|
|
error.into_inner().complete_delivery(Err(
|
|
"outbound lane could not be restarted during shutdown".to_string(),
|
|
));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn spawn_lane(
|
|
&self,
|
|
channel: Arc<dyn Channel + Send + Sync>,
|
|
mut receiver: mpsc::Receiver<OutboundMessage>,
|
|
channel_name: String,
|
|
chat_id: String,
|
|
) -> bool {
|
|
let target_lock = self.write_locks.for_target(&channel_name, &chat_id);
|
|
self.task_supervisor.spawn(
|
|
format!("outbound-lane:{channel_name}:{chat_id}"),
|
|
async move {
|
|
loop {
|
|
let msg = match tokio::time::timeout(LANE_IDLE_TIMEOUT, receiver.recv()).await {
|
|
Ok(Some(msg)) => msg,
|
|
Ok(None) | Err(_) => break,
|
|
};
|
|
let result = Self::send_with_retry(&*channel, &msg, &target_lock).await;
|
|
if let Err(error) = &result {
|
|
tracing::error!(
|
|
channel = %channel_name,
|
|
chat_id = %chat_id,
|
|
error = %error,
|
|
"Failed to send message after retries"
|
|
);
|
|
}
|
|
msg.complete_delivery(result.map_err(|error| error.to_string()));
|
|
}
|
|
},
|
|
)
|
|
}
|
|
|
|
async fn send_with_retry(
|
|
channel: &dyn Channel,
|
|
msg: &OutboundMessage,
|
|
target_lock: &tokio::sync::Mutex<()>,
|
|
) -> Result<(), ChannelError> {
|
|
let _guard = target_lock.lock().await;
|
|
const DELAYS: &[u64] = &[1, 2, 4];
|
|
|
|
for (attempt, &delay) in DELAYS.iter().enumerate() {
|
|
let result = tokio::time::timeout(SEND_TIMEOUT, channel.send(msg.clone())).await;
|
|
match result {
|
|
Ok(Ok(())) => return Ok(()),
|
|
Ok(Err(error)) if attempt < DELAYS.len() - 1 && error.is_transient() => {
|
|
tracing::warn!(attempt = attempt + 1, delay, error = %error, "Send failed, retrying");
|
|
}
|
|
Ok(Err(error)) => return Err(error),
|
|
Err(_) if attempt < DELAYS.len() - 1 => {
|
|
tracing::warn!(attempt = attempt + 1, delay, "Send timed out, retrying");
|
|
}
|
|
Err(_) => {
|
|
return Err(ChannelError::Other(format!(
|
|
"send timed out after {} seconds",
|
|
SEND_TIMEOUT.as_secs()
|
|
)));
|
|
}
|
|
}
|
|
tokio::time::sleep(Duration::from_secs(delay)).await;
|
|
}
|
|
unreachable!()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use async_trait::async_trait;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use tokio::sync::{Mutex, Notify};
|
|
|
|
struct RecordingChannel {
|
|
sent: Mutex<Vec<String>>,
|
|
notify: Notify,
|
|
}
|
|
|
|
struct PermanentFailureChannel {
|
|
attempts: AtomicUsize,
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Channel for PermanentFailureChannel {
|
|
fn name(&self) -> &str {
|
|
"permanent-failure"
|
|
}
|
|
fn is_running(&self) -> bool {
|
|
true
|
|
}
|
|
async fn start(&self, _bus: Arc<MessageBus>) -> Result<(), ChannelError> {
|
|
Ok(())
|
|
}
|
|
async fn stop(&self) -> Result<(), ChannelError> {
|
|
Ok(())
|
|
}
|
|
async fn send(&self, _msg: OutboundMessage) -> Result<(), ChannelError> {
|
|
self.attempts.fetch_add(1, Ordering::SeqCst);
|
|
Err(ChannelError::Other("invalid destination".to_string()))
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Channel for RecordingChannel {
|
|
fn name(&self) -> &str {
|
|
"recording"
|
|
}
|
|
|
|
fn is_running(&self) -> bool {
|
|
true
|
|
}
|
|
|
|
async fn start(&self, _bus: Arc<MessageBus>) -> Result<(), ChannelError> {
|
|
Ok(())
|
|
}
|
|
|
|
async fn stop(&self) -> Result<(), ChannelError> {
|
|
Ok(())
|
|
}
|
|
|
|
async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> {
|
|
if msg.chat_id == "slow" {
|
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
|
}
|
|
self.sent.lock().await.push(msg.content);
|
|
self.notify.notify_waiters();
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
fn outbound(chat_id: &str, content: &str) -> OutboundMessage {
|
|
OutboundMessage {
|
|
channel: "recording".to_string(),
|
|
chat_id: chat_id.to_string(),
|
|
content: content.to_string(),
|
|
reply_to: None,
|
|
media: vec![],
|
|
metadata: HashMap::new(),
|
|
delivery: None,
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn slow_conversation_does_not_block_other_conversations() {
|
|
let bus = MessageBus::new(8);
|
|
let manager = ChannelManager::with_bus(
|
|
Arc::new(crate::channels::CliChatChannel::new()),
|
|
bus.clone(),
|
|
);
|
|
let channel = Arc::new(RecordingChannel {
|
|
sent: Mutex::new(Vec::new()),
|
|
notify: Notify::new(),
|
|
});
|
|
manager.register_channel("recording", channel.clone()).await;
|
|
|
|
let supervisor = TaskSupervisor::new();
|
|
let dispatcher = OutboundDispatcher::new(
|
|
bus.clone(),
|
|
manager,
|
|
supervisor.clone(),
|
|
ConversationWriteLocks::default(),
|
|
);
|
|
let task = tokio::spawn(async move { dispatcher.run().await });
|
|
bus.publish_outbound(outbound("slow", "slow-1"))
|
|
.await
|
|
.unwrap();
|
|
bus.publish_outbound(outbound("slow", "slow-2"))
|
|
.await
|
|
.unwrap();
|
|
bus.publish_outbound(outbound("fast", "fast-1"))
|
|
.await
|
|
.unwrap();
|
|
|
|
tokio::time::timeout(Duration::from_secs(1), async {
|
|
loop {
|
|
if channel.sent.lock().await.len() == 3 {
|
|
break;
|
|
}
|
|
tokio::time::sleep(Duration::from_millis(5)).await;
|
|
}
|
|
})
|
|
.await
|
|
.unwrap();
|
|
|
|
let sent = channel.sent.lock().await.clone();
|
|
assert_eq!(sent[0], "fast-1");
|
|
assert_eq!(&sent[1..], &["slow-1", "slow-2"]);
|
|
task.abort();
|
|
supervisor.shutdown(Duration::from_secs(1)).await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn confirmed_delivery_reports_missing_channel() {
|
|
let bus = MessageBus::new(8);
|
|
let manager = ChannelManager::with_bus(
|
|
Arc::new(crate::channels::CliChatChannel::new()),
|
|
bus.clone(),
|
|
);
|
|
let supervisor = TaskSupervisor::new();
|
|
let dispatcher = OutboundDispatcher::new(
|
|
bus.clone(),
|
|
manager,
|
|
supervisor.clone(),
|
|
ConversationWriteLocks::default(),
|
|
);
|
|
let task = tokio::spawn(async move { dispatcher.run().await });
|
|
|
|
let mut message = outbound("missing", "not delivered");
|
|
message.channel = "missing".to_string();
|
|
let error = bus.deliver_outbound(message).await.unwrap_err();
|
|
|
|
assert!(matches!(error, crate::bus::BusError::DeliveryFailed(_)));
|
|
task.abort();
|
|
supervisor.shutdown(Duration::from_secs(1)).await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn confirmed_delivery_waits_for_channel_send() {
|
|
let bus = MessageBus::new(8);
|
|
let manager = ChannelManager::with_bus(
|
|
Arc::new(crate::channels::CliChatChannel::new()),
|
|
bus.clone(),
|
|
);
|
|
let channel = Arc::new(RecordingChannel {
|
|
sent: Mutex::new(Vec::new()),
|
|
notify: Notify::new(),
|
|
});
|
|
manager.register_channel("recording", channel.clone()).await;
|
|
let supervisor = TaskSupervisor::new();
|
|
let dispatcher = OutboundDispatcher::new(
|
|
bus.clone(),
|
|
manager,
|
|
supervisor.clone(),
|
|
ConversationWriteLocks::default(),
|
|
);
|
|
let task = tokio::spawn(async move { dispatcher.run().await });
|
|
|
|
bus.deliver_outbound(outbound("confirmed", "delivered"))
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(channel.sent.lock().await.as_slice(), &["delivered"]);
|
|
task.abort();
|
|
supervisor.shutdown(Duration::from_secs(1)).await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn permanent_send_failure_is_not_retried() {
|
|
let channel = PermanentFailureChannel {
|
|
attempts: AtomicUsize::new(0),
|
|
};
|
|
|
|
let target_lock = tokio::sync::Mutex::new(());
|
|
let error = OutboundDispatcher::send_with_retry(
|
|
&channel,
|
|
&outbound("invalid", "message"),
|
|
&target_lock,
|
|
)
|
|
.await
|
|
.unwrap_err();
|
|
|
|
assert!(matches!(error, ChannelError::Other(_)));
|
|
assert_eq!(channel.attempts.load(Ordering::SeqCst), 1);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn shared_write_lock_orders_dispatcher_with_live_turn_writes() {
|
|
let channel = RecordingChannel {
|
|
sent: Mutex::new(Vec::new()),
|
|
notify: Notify::new(),
|
|
};
|
|
let write_locks = ConversationWriteLocks::default();
|
|
let target_lock = write_locks.for_target("recording", "same-chat");
|
|
let live_write = target_lock.lock().await;
|
|
let message = outbound("same-chat", "after-live-update");
|
|
|
|
let send = OutboundDispatcher::send_with_retry(&channel, &message, &target_lock);
|
|
tokio::pin!(send);
|
|
assert!(
|
|
tokio::time::timeout(Duration::from_millis(10), &mut send)
|
|
.await
|
|
.is_err()
|
|
);
|
|
assert!(channel.sent.lock().await.is_empty());
|
|
|
|
drop(live_write);
|
|
send.await.unwrap();
|
|
assert_eq!(channel.sent.lock().await.as_slice(), &["after-live-update"]);
|
|
}
|
|
}
|