use std::collections::HashMap;
use anyhow::Result;
use tokio::runtime::Builder;
pub struct Subscriber<T> {
id: String,
topic: String,
exec_fn: fn(msg: T) -> Result<()>,
}
impl<T> Subscriber<T> {
pub fn get_id(&self) -> String {
self.id.clone()
}
}
pub struct Publisher<T> {
rt: tokio::runtime::Runtime,
subs: HashMap<String, Subscriber<T>>,
}
impl<T: Clone + Send + 'static> Publisher<T> {
pub fn new(wt: usize) -> Self {
let rt = Builder::new_multi_thread()
.worker_threads(wt)
.enable_all()
.build()
.unwrap();
Self {
rt,
subs: HashMap::new(),
}
}
pub async fn sub(&mut self, topic: &str, exec_fn: fn(msg: T) -> Result<()>) -> String {
let sub_id = nanoid::nanoid!(10);
self.subs.insert(
sub_id.clone(),
Subscriber {
id: sub_id.clone(),
topic: topic.to_string(),
exec_fn,
},
);
sub_id
}
pub async fn remove(&mut self, id: &str) {
self.subs.remove(id).unwrap();
}
pub async fn publish_all(&self, msg: T) {
self.publish("", msg).await;
}
pub async fn publish(&self, topic: &str, msg: T) {
let mut hds = vec![];
for (_, sub) in self.subs.iter() {
if topic.eq("") || sub.topic.eq(topic) {
let exec_fn = sub.exec_fn;
let c_msg = msg.clone();
let spawn = self.rt.spawn(async move {
match exec_fn(c_msg) {
Ok(_) => {}
Err(err) => {
println!("{err}")
}
}
});
hds.push(spawn);
}
}
for ele in hds {
ele.await.unwrap()
}
}
pub fn shutdown(self) {
self.rt.shutdown_background();
}
}
#[tokio::test]
async fn test_pub_sub() {
let mut ps: Publisher<String> = Publisher::new(10);
ps.sub("hello", |msg| {
println!("hello:{msg}");
Ok(())
})
.await;
ps.sub("hello", |msg| {
println!("hello:{msg}");
Ok(())
})
.await;
ps.publish("hello", "hello niha".to_string()).await;
ps.shutdown();
println!("haha")
}