simple_pub_sub/
topics.rs

1use crate::message::{self, Msg};
2use log::{error, info, trace};
3use serde_json::json;
4use std::collections::{BTreeMap, HashMap};
5use tokio;
6use tokio::sync::broadcast::Sender;
7
8type ClientChannelMap = HashMap<String, Sender<Msg>>;
9
10/// The `TopicMap` struct is used to store the channels for a given topic.
11#[derive(Debug, Clone)]
12pub struct TopicMap {
13    pub map: BTreeMap<String, ClientChannelMap>,
14}
15impl TopicMap {
16    /// Returns the number of connected clients for a given topic.
17    pub fn query(&self, topic: String) -> String {
18        let v: Vec<String>;
19        if topic == "*" {
20            v = self
21                .map
22                .iter()
23                .map(|(k, v)| format!("{}: {}", k, v.len()))
24                .collect();
25            json!({topic: v}).to_string()
26        } else if self.map.contains_key(&topic) {
27            v = vec![format!("{}", self.map.get(&topic).unwrap().len())];
28            json!({topic: v}).to_string()
29        } else {
30            "".to_string()
31        }
32    }
33    /// Adds a channel to the map.
34    pub fn add_channel(&mut self, topic: String, client_id: String, channel: Sender<Msg>) {
35        if self.map.contains_key(&topic.clone()) {
36            if let Some(channels) = self.map.get_mut(&topic.clone()) {
37                channels.entry(client_id).or_insert(channel);
38                // Not sure if the channel should be replaced if the key is already present.
39            }
40        } else {
41            let mut client_map = ClientChannelMap::new();
42            client_map.insert(client_id, channel);
43            self.map.insert(topic, client_map);
44        }
45    }
46    /// Removes a channel from the map.
47    pub fn remove_channel(&mut self, topic: String, client_id: String) {
48        if self.map.contains_key(&topic) {
49            if let Some(channels) = self.map.get_mut(&topic) {
50                channels.remove(&client_id);
51            }
52            trace!("channels: {:?}", self.map);
53        }
54    }
55    /// adds a topic to the map.
56    pub fn add_topic(&mut self, topic: String) {
57        self.map.entry(topic).or_default();
58    }
59
60    /// Publishes the message to the channels.
61    pub async fn publish(&mut self, msg: Msg) {
62        if !self.map.contains_key(&msg.topic) {
63            return;
64        }
65        let topic = msg.topic.clone();
66
67        if let Some(channels) = self.map.get_mut(&topic.clone()) {
68            let dead_channels = channels
69                .iter()
70                .map(|(client_id, channel)| {
71                    info!("sending msg to the {}", client_id);
72                    match channel.send(msg.clone()) {
73                        Ok(_n) => "".to_string(),
74                        Err(e) => {
75                            error!(
76                                "error occurred: {} while sending the message to the channel {}",
77                                e.to_string(),
78                                client_id
79                            );
80                            error!("cleaning up");
81                            client_id.clone()
82                        }
83                    }
84                })
85                .collect::<Vec<_>>();
86            info!("dead_channels: {:?}", dead_channels);
87            let _ = dead_channels
88                .iter()
89                .map(|client_id| {
90                    self.remove_channel(topic.clone(), client_id.clone());
91                })
92                .collect::<Vec<_>>();
93        }
94    }
95}
96
97/// returns a global broadcaster.
98pub fn get_global_broadcaster() -> tokio::sync::broadcast::Sender<Msg> {
99    let channel_capacity_str = std::env::var("PUB_SUB_CAP").unwrap_or("1024".to_string());
100    let channel_capacity = channel_capacity_str.parse::<u16>().unwrap_or_else(|x| {
101        error!("invalid channel channel_capacity: {}", x);
102        error!("using the default value 1024");
103        1024
104    });
105    info!(
106        "creating broadcast channel with capacity: {}",
107        channel_capacity
108    );
109    let (glob_tx, _) = tokio::sync::broadcast::channel(channel_capacity.into());
110    glob_tx
111}
112
113/// Handles the incoming and out-going messages for each topic.
114pub async fn topic_manager(chan: Sender<Msg>) {
115    let mut map: TopicMap = TopicMap {
116        map: BTreeMap::new(),
117    };
118    let mut rx = chan.subscribe();
119    loop {
120        match rx.recv().await {
121            Ok(msg) => {
122                if !msg.topic.is_empty() {
123                    info!("topic received: {}", msg.topic);
124                    match msg.header.pkt_type {
125                        message::PktType::PUBLISH => {
126                            trace!("publishing to map:{:?}", map);
127                            map.publish(msg).await;
128                        }
129                        message::PktType::SUBSCRIBE => {
130                            if msg.client_id.is_some() || msg.channel.is_some() {
131                                map.add_channel(
132                                    msg.topic,
133                                    msg.client_id.unwrap(),
134                                    msg.channel.unwrap(),
135                                );
136                                trace!("map: {:?}", map);
137                            }
138                        }
139                        message::PktType::UNSUBSCRIBE => {
140                            if msg.client_id.is_some() {
141                                info!("unsubscribing:");
142                                map.remove_channel(msg.topic, msg.client_id.unwrap());
143                            }
144                        }
145                        message::PktType::QUERY => {
146                            info!("querying");
147                            let query_resp = map.query(msg.topic.clone());
148                            info!("query_resp: {}", query_resp.clone());
149                            let resp_msg = match msg.response_msg(query_resp.into_bytes()) {
150                                Ok(rm) => rm,
151                                Err(e) => {
152                                    error!(
153                                        "error while getting the response to the query message: {}",
154                                        e.to_string()
155                                    );
156                                    continue;
157                                }
158                            };
159                            info!("generated query resp: {:?}", resp_msg);
160                            match msg.channel.unwrap().send(resp_msg) {
161                                Ok(n) => n,
162                                Err(e) => {
163                                    error!(
164                                        "error while sending the query response: {}",
165                                        e.to_string()
166                                    );
167                                    0
168                                }
169                            };
170                        }
171                        _ => {}
172                    };
173                }
174            }
175            Err(e) => {
176                error!(
177                    "error occurred while receiving the topic: {}",
178                    e.to_string()
179                );
180                // "".to_string()
181            }
182        };
183    }
184}
simple_pub_sub/topics.rs

simple_pub_sub/
topics.rs
