66 lines
1.8 KiB
Rust
66 lines
1.8 KiB
Rust
use {
|
|
eyre::Result,
|
|
futures::Stream,
|
|
lool::fail,
|
|
std::{
|
|
pin::Pin,
|
|
task::{Context, Poll},
|
|
},
|
|
tokio::sync::broadcast::{self, Receiver, Sender},
|
|
};
|
|
|
|
use crate::bus::BusMessage;
|
|
|
|
pub struct SourceStream<RM: BusMessage> {
|
|
sender: Option<Sender<RM>>,
|
|
}
|
|
|
|
impl<RM: BusMessage> Default for SourceStream<RM> {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl<RM: BusMessage> SourceStream<RM> {
|
|
// Create a new SourceStream with a broadcast channel
|
|
pub fn new() -> Self {
|
|
let (sender, _) = broadcast::channel(100); // Adjust the buffer size as needed
|
|
SourceStream {
|
|
sender: Some(sender),
|
|
}
|
|
}
|
|
|
|
pub fn sender(&self) -> Option<Sender<RM>> {
|
|
self.sender.clone()
|
|
}
|
|
|
|
// Subscribe to the stream
|
|
pub fn subscribe(&self) -> Result<Pin<Box<dyn Stream<Item = RM> + Send + 'static>>> {
|
|
if let Some(sender) = &self.sender {
|
|
let receiver = sender.subscribe();
|
|
Ok(Box::pin(BroadcastStream { receiver }))
|
|
} else {
|
|
fail!("SourceStream has been consumed")
|
|
}
|
|
}
|
|
}
|
|
|
|
// Wrapper around Receiver to implement Stream
|
|
struct BroadcastStream<RM: BusMessage> {
|
|
receiver: Receiver<RM>,
|
|
}
|
|
|
|
impl<RM: BusMessage> Stream for BroadcastStream<RM> {
|
|
type Item = RM;
|
|
|
|
fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
|
let this = self.get_mut();
|
|
// Spawn an async task to receive the message
|
|
tokio::task::block_in_place(|| match futures::executor::block_on(this.receiver.recv()) {
|
|
Ok(msg) => Poll::Ready(Some(msg)),
|
|
Err(broadcast::error::RecvError::Closed) => Poll::Ready(None),
|
|
Err(broadcast::error::RecvError::Lagged(_)) => Poll::Pending,
|
|
})
|
|
}
|
|
}
|