1use super::{BroadcastMsg, ConnectionId, GlobalRegistry, PatternTrie, PubSubMessage, ThreadId};
2use bytes::Bytes;
3use crossbeam_channel::Receiver;
4use std::collections::HashMap;
5use std::sync::Arc;
6
7pub struct ThreadLocalPubSub {
8 thread_id: ThreadId,
9 exact_subs: HashMap<Vec<u8>, Vec<ConnectionId>>,
10 pattern_trie: PatternTrie,
11 pattern_subs: HashMap<Vec<u8>, Vec<ConnectionId>>,
12 inbox: Receiver<BroadcastMsg>,
13 registry: Arc<GlobalRegistry>,
14 connection_channels: HashMap<ConnectionId, Vec<Vec<u8>>>,
15 connection_patterns: HashMap<ConnectionId, Vec<Vec<u8>>>,
16}
17
18impl ThreadLocalPubSub {
19 pub fn new(
20 thread_id: ThreadId,
21 inbox: Receiver<BroadcastMsg>,
22 registry: Arc<GlobalRegistry>,
23 ) -> Self {
24 Self {
25 thread_id,
26 exact_subs: HashMap::new(),
27 pattern_trie: PatternTrie::new(),
28 pattern_subs: HashMap::new(),
29 inbox,
30 registry,
31 connection_channels: HashMap::new(),
32 connection_patterns: HashMap::new(),
33 }
34 }
35
36 pub fn subscribe(
37 &mut self,
38 conn_id: ConnectionId,
39 channels: Vec<Vec<u8>>,
40 ) -> Vec<PubSubMessage> {
41 let mut messages = Vec::new();
42
43 for channel in channels {
44 let is_new_channel = !self.exact_subs.contains_key(&channel);
45 let already_subscribed = self
46 .exact_subs
47 .get(&channel)
48 .map(|subs| subs.contains(&conn_id))
49 .unwrap_or(false);
50
51 if !already_subscribed {
52 self.exact_subs
53 .entry(channel.clone())
54 .or_default()
55 .push(conn_id);
56 self.connection_channels
57 .entry(conn_id)
58 .or_default()
59 .push(channel.clone());
60
61 self.registry.increment_channel_subscribers(&channel);
63 }
64
65 if is_new_channel {
66 self.registry
67 .add_channel_interest(channel.clone(), self.thread_id);
68 }
69
70 let channel_count = self
71 .connection_channels
72 .get(&conn_id)
73 .map(|c| c.len())
74 .unwrap_or(0);
75 let pattern_count = self
76 .connection_patterns
77 .get(&conn_id)
78 .map(|p| p.len())
79 .unwrap_or(0);
80 let total_count = channel_count + pattern_count;
81
82 messages.push(PubSubMessage::Subscribe {
83 channel,
84 count: total_count,
85 });
86 }
87
88 messages
89 }
90
91 pub fn unsubscribe(
92 &mut self,
93 conn_id: ConnectionId,
94 channels: Option<Vec<Vec<u8>>>,
95 ) -> Vec<PubSubMessage> {
96 let mut messages = Vec::new();
97
98 if let Some(channels) = channels {
99 for channel in channels {
100 let mut should_remove_interest = false;
101
102 if let Some(subs) = self.exact_subs.get_mut(&channel) {
103 let was_subscribed = subs.contains(&conn_id);
104 subs.retain(|&id| id != conn_id);
105 if was_subscribed {
106 self.registry.decrement_channel_subscribers(&channel);
107 }
108 if subs.is_empty() {
109 self.exact_subs.remove(&channel);
110 should_remove_interest = true;
111 }
112 }
113
114 if let Some(conn_channels) = self.connection_channels.get_mut(&conn_id) {
115 conn_channels.retain(|c| c != &channel);
116 }
117
118 if should_remove_interest {
119 self.registry
120 .remove_channel_interest(&channel, self.thread_id);
121 }
122
123 let total_count = self.get_connection_subscription_count(conn_id);
124 messages.push(PubSubMessage::Unsubscribe {
125 channel: Some(channel),
126 count: total_count,
127 });
128 }
129 } else if let Some(conn_channels) = self.connection_channels.remove(&conn_id) {
130 for channel in conn_channels {
131 if let Some(subs) = self.exact_subs.get_mut(&channel) {
132 let was_subscribed = subs.contains(&conn_id);
133 subs.retain(|&id| id != conn_id);
134 if was_subscribed {
135 self.registry.decrement_channel_subscribers(&channel);
136 }
137 if subs.is_empty() {
138 self.exact_subs.remove(&channel);
139 self.registry
140 .remove_channel_interest(&channel, self.thread_id);
141 }
142 }
143
144 let total_count = self.get_connection_subscription_count(conn_id);
145 messages.push(PubSubMessage::Unsubscribe {
146 channel: Some(channel),
147 count: total_count,
148 });
149 }
150 }
151
152 messages
153 }
154
155 pub fn psubscribe(
156 &mut self,
157 conn_id: ConnectionId,
158 patterns: Vec<Vec<u8>>,
159 ) -> Vec<PubSubMessage> {
160 let mut messages = Vec::new();
161
162 for pattern in patterns {
163 let is_new_pattern = !self.pattern_subs.contains_key(&pattern);
164 let already_subscribed = self
165 .pattern_subs
166 .get(&pattern)
167 .map(|subs| subs.contains(&conn_id))
168 .unwrap_or(false);
169
170 if !already_subscribed {
171 self.pattern_trie.insert(&pattern, conn_id);
172 self.pattern_subs
173 .entry(pattern.clone())
174 .or_default()
175 .push(conn_id);
176 self.connection_patterns
177 .entry(conn_id)
178 .or_default()
179 .push(pattern.clone());
180
181 self.registry.increment_pattern_subscribers(&pattern);
183 }
184
185 if is_new_pattern {
186 self.registry
187 .add_pattern_interest(pattern.clone(), self.thread_id);
188 }
189
190 let channel_count = self
191 .connection_channels
192 .get(&conn_id)
193 .map(|c| c.len())
194 .unwrap_or(0);
195 let pattern_count = self
196 .connection_patterns
197 .get(&conn_id)
198 .map(|p| p.len())
199 .unwrap_or(0);
200 let total_count = channel_count + pattern_count;
201
202 messages.push(PubSubMessage::PSubscribe {
203 pattern,
204 count: total_count,
205 });
206 }
207
208 messages
209 }
210
211 pub fn punsubscribe(
212 &mut self,
213 conn_id: ConnectionId,
214 patterns: Option<Vec<Vec<u8>>>,
215 ) -> Vec<PubSubMessage> {
216 let mut messages = Vec::new();
217
218 if let Some(patterns) = patterns {
219 for pattern in patterns {
220 let mut should_remove_interest = false;
221
222 self.pattern_trie.remove(&pattern, conn_id);
223
224 if let Some(subs) = self.pattern_subs.get_mut(&pattern) {
225 let was_subscribed = subs.contains(&conn_id);
226 subs.retain(|&id| id != conn_id);
227 if was_subscribed {
228 self.registry.decrement_pattern_subscribers(&pattern);
229 }
230 if subs.is_empty() {
231 self.pattern_subs.remove(&pattern);
232 should_remove_interest = true;
233 }
234 }
235
236 if let Some(conn_patterns) = self.connection_patterns.get_mut(&conn_id) {
237 conn_patterns.retain(|p| p != &pattern);
238 }
239
240 if should_remove_interest {
241 self.registry
242 .remove_pattern_interest(&pattern, self.thread_id);
243 }
244
245 let total_count = self.get_connection_subscription_count(conn_id);
246 messages.push(PubSubMessage::PUnsubscribe {
247 pattern: Some(pattern),
248 count: total_count,
249 });
250 }
251 } else if let Some(conn_patterns) = self.connection_patterns.remove(&conn_id) {
252 for pattern in conn_patterns {
253 self.pattern_trie.remove(&pattern, conn_id);
254
255 if let Some(subs) = self.pattern_subs.get_mut(&pattern) {
256 let was_subscribed = subs.contains(&conn_id);
257 subs.retain(|&id| id != conn_id);
258 if was_subscribed {
259 self.registry.decrement_pattern_subscribers(&pattern);
260 }
261 if subs.is_empty() {
262 self.pattern_subs.remove(&pattern);
263 self.registry
264 .remove_pattern_interest(&pattern, self.thread_id);
265 }
266 }
267
268 let total_count = self.get_connection_subscription_count(conn_id);
269 messages.push(PubSubMessage::PUnsubscribe {
270 pattern: Some(pattern),
271 count: total_count,
272 });
273 }
274 }
275
276 messages
277 }
278
279 pub fn publish_local(
280 &self,
281 channel: &[u8],
282 message: &Bytes,
283 ) -> Vec<(ConnectionId, PubSubMessage)> {
284 let mut deliveries = Vec::new();
285
286 if let Some(subs) = self.exact_subs.get(channel) {
287 for &conn_id in subs {
288 deliveries.push((
289 conn_id,
290 PubSubMessage::Message {
291 channel: channel.to_vec(),
292 payload: message.clone(),
293 },
294 ));
295 }
296 }
297
298 let pattern_matches = self.pattern_trie.find_matches(channel);
299 for (pattern, conn_id) in pattern_matches {
300 deliveries.push((
301 conn_id,
302 PubSubMessage::PatternMessage {
303 pattern,
304 channel: channel.to_vec(),
305 payload: message.clone(),
306 },
307 ));
308 }
309
310 deliveries
311 }
312
313 pub fn process_inbox(&mut self) -> Vec<(ConnectionId, PubSubMessage)> {
314 let mut deliveries = Vec::new();
315
316 while let Ok(msg) = self.inbox.try_recv() {
317 match msg {
318 BroadcastMsg::Publish {
319 channel,
320 message,
321 exclude_thread,
322 } => {
323 if Some(self.thread_id) != exclude_thread {
324 let local_deliveries = self.publish_local(&channel, &message);
325 deliveries.extend(local_deliveries);
326 }
327 }
328 BroadcastMsg::PatternPublish {
329 channel,
330 message,
331 exclude_thread,
332 } => {
333 if Some(self.thread_id) != exclude_thread {
334 let pattern_matches = self.pattern_trie.find_matches(&channel);
335 for (pattern, conn_id) in pattern_matches {
336 deliveries.push((
337 conn_id,
338 PubSubMessage::PatternMessage {
339 pattern,
340 channel: channel.clone(),
341 payload: message.clone(),
342 },
343 ));
344 }
345 }
346 }
347 }
348 }
349
350 deliveries
351 }
352
353 pub fn connection_dropped(&mut self, conn_id: ConnectionId) {
354 self.unsubscribe(conn_id, None);
355 self.punsubscribe(conn_id, None);
356 }
357
358 pub fn get_connection_subscription_count(&self, conn_id: ConnectionId) -> usize {
359 let channel_count = self
360 .connection_channels
361 .get(&conn_id)
362 .map(|c| c.len())
363 .unwrap_or(0);
364 let pattern_count = self
365 .connection_patterns
366 .get(&conn_id)
367 .map(|p| p.len())
368 .unwrap_or(0);
369 channel_count + pattern_count
370 }
371
372 pub fn is_connection_subscribed(&self, conn_id: ConnectionId) -> bool {
373 self.connection_channels.contains_key(&conn_id)
374 || self.connection_patterns.contains_key(&conn_id)
375 }
376
377 pub fn get_all_channels(&self) -> Vec<Vec<u8>> {
378 self.exact_subs.keys().cloned().collect()
379 }
380
381 pub fn get_all_patterns(&self) -> Vec<Vec<u8>> {
382 self.pattern_subs.keys().cloned().collect()
383 }
384}