use crate::messages::Message;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
pub enum ReceiverType<T> {
Unbounded(mpsc::UnboundedReceiver<T>),
Bounded(mpsc::Receiver<T>),
}
pub struct AggregatedReceiver<T: Message> {
pattern: String,
topics: Vec<String>,
rx: mpsc::UnboundedReceiver<(String, T)>,
_tasks: Vec<JoinHandle<()>>,
}
impl<T: Message> AggregatedReceiver<T> {
pub fn new(pattern: String, receivers: Vec<(String, mpsc::UnboundedReceiver<T>)>) -> Self {
let receivers = receivers
.into_iter()
.map(|(topic_path, receiver)| (topic_path, ReceiverType::Unbounded(receiver)))
.collect();
Self::new_mixed(pattern, receivers)
}
pub fn new_mixed(pattern: String, receivers: Vec<(String, ReceiverType<T>)>) -> Self {
let (tx, rx) = mpsc::unbounded_channel();
let mut tasks = Vec::new();
let mut topics = Vec::new();
for (topic_path, mut receiver) in receivers {
topics.push(topic_path.clone());
let tx = tx.clone();
let topic_path_clone = topic_path.clone();
let task = tokio::spawn(async move {
loop {
let msg_opt = match &mut receiver {
ReceiverType::Unbounded(rx) => rx.recv().await,
ReceiverType::Bounded(rx) => rx.recv().await,
};
match msg_opt {
Some(message) => {
if tx.send((topic_path_clone.clone(), message)).is_err() {
break;
}
}
None => {
break;
}
}
}
});
tasks.push(task);
}
Self {
pattern,
topics,
rx,
_tasks: tasks,
}
}
pub async fn recv(&mut self) -> Option<(String, T)> {
self.rx.recv().await
}
pub fn pattern(&self) -> &str {
&self.pattern
}
pub fn topics(&self) -> &[String] {
&self.topics
}
pub fn topic_count(&self) -> usize {
self.topics.len()
}
}
impl<T: Message> Unpin for AggregatedReceiver<T> {}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn new_forwards_messages_from_unbounded_receivers() {
let (tx_a, rx_a) = mpsc::unbounded_channel::<i32>();
let (tx_b, rx_b) = mpsc::unbounded_channel::<i32>();
let mut aggregated = AggregatedReceiver::new(
"topic/*".to_string(),
vec![("topic/a".to_string(), rx_a), ("topic/b".to_string(), rx_b)],
);
assert_eq!(aggregated.pattern(), "topic/*");
assert_eq!(aggregated.topic_count(), 2);
tx_a.send(1).unwrap();
tx_b.send(2).unwrap();
drop(tx_a);
drop(tx_b);
let mut received = vec![aggregated.recv().await.unwrap(), aggregated.recv().await.unwrap()];
received.sort();
assert_eq!(received, vec![("topic/a".to_string(), 1), ("topic/b".to_string(), 2)]);
assert_eq!(aggregated.recv().await, None);
}
}
pub mod pattern {
pub fn matches_pattern(path: &str, pattern: &str) -> bool {
use globset::GlobBuilder;
let glob = GlobBuilder::new(pattern)
.literal_separator(false)
.backslash_escape(true)
.build();
match glob {
Ok(glob_pattern) => {
let matcher = glob_pattern.compile_matcher();
matcher.is_match(path)
}
Err(_) => path == pattern,
}
}
pub fn convert_glob_to_redis(pattern: &str) -> String {
format!("mecha10:{}", pattern)
}
}