use super::{AgentEvent, EventChannel, EventReceiver};
use futures::future::BoxFuture;
use std::sync::Arc;
#[derive(Debug)]
pub struct MpscEventChannel {
tx: tokio::sync::mpsc::Sender<Arc<dyn AgentEvent>>,
rx: std::sync::Mutex<Option<tokio::sync::mpsc::Receiver<Arc<dyn AgentEvent>>>>,
}
impl Default for MpscEventChannel {
fn default() -> Self {
Self::new(256)
}
}
impl MpscEventChannel {
pub fn new(capacity: usize) -> Self {
let (tx, rx) = tokio::sync::mpsc::channel(capacity);
Self {
tx,
rx: std::sync::Mutex::new(Some(rx)),
}
}
}
impl EventChannel for MpscEventChannel {
fn publish(&self, event: Arc<dyn AgentEvent>) {
let _ = self.tx.try_send(event);
}
fn subscribe(&self) -> Box<dyn EventReceiver> {
let mut guard = self
.rx
.lock()
.expect("MpscEventChannel is single-consumer; subscribe may only be called once");
let rx = guard
.take()
.expect("MpscEventChannel is single-consumer; subscribe may only be called once");
Box::new(MpscEventReceiver { rx })
}
}
struct MpscEventReceiver {
rx: tokio::sync::mpsc::Receiver<Arc<dyn AgentEvent>>,
}
impl EventReceiver for MpscEventReceiver {
fn recv(&mut self) -> BoxFuture<'_, Option<Arc<dyn AgentEvent>>> {
Box::pin(async move { self.rx.recv().await })
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::agent::{AgentEvent, ReActEvent};
use futures::FutureExt;
fn ev(n: &str) -> Arc<dyn AgentEvent> {
Arc::new(ReActEvent::RunStarted {
run_id: "test".into(),
input: n.into(),
})
}
fn input_of(ev: &dyn AgentEvent) -> &str {
match ev.as_any().downcast_ref::<ReActEvent>() {
Some(ReActEvent::RunStarted { input, .. }) => input.as_str(),
_ => panic!("test event must be RunStarted"),
}
}
#[tokio::test]
async fn publish_subscribe_buffered() {
let ch = MpscEventChannel::new(16);
ch.publish(ev("1"));
ch.publish(ev("2"));
let mut rx = ch.subscribe();
assert_eq!(input_of(&*rx.recv().await.unwrap()), "1");
assert_eq!(input_of(&*rx.recv().await.unwrap()), "2");
}
#[tokio::test]
#[should_panic(expected = "MpscEventChannel is single-consumer")]
async fn subscribe_twice_panics() {
let ch = MpscEventChannel::new(8);
let _rx = ch.subscribe();
let _rx2 = ch.subscribe();
}
#[tokio::test]
async fn full_drops_new() {
let ch = MpscEventChannel::new(2);
ch.publish(ev("1"));
ch.publish(ev("2"));
ch.publish(ev("3")); let mut rx = ch.subscribe();
assert_eq!(input_of(&*rx.recv().await.unwrap()), "1");
assert_eq!(input_of(&*rx.recv().await.unwrap()), "2");
assert!(rx.recv().now_or_never().is_none()); }
#[tokio::test]
async fn closed_ends_stream() {
let mut rx = {
let ch = MpscEventChannel::new(16);
ch.publish(ev("1"));
ch.subscribe()
};
assert_eq!(input_of(&*rx.recv().await.unwrap()), "1");
assert!(rx.recv().await.is_none());
}
}