hyperlane_plugin_websocket/websocket/
impl.rs

1use crate::*;
2
3/// Allows `String` to be used as a broadcast identifier.
4impl BroadcastTypeTrait for String {}
5/// Allows string slices to be used as broadcast identifiers.
6impl BroadcastTypeTrait for &str {}
7/// Allows `char` to be used as a broadcast identifier.
8impl BroadcastTypeTrait for char {}
9/// Allows `bool` to be used as a broadcast identifier.
10impl BroadcastTypeTrait for bool {}
11/// Allows `i8` to be used as a broadcast identifier.
12impl BroadcastTypeTrait for i8 {}
13/// Allows `i16` to be used as a broadcast identifier.
14impl BroadcastTypeTrait for i16 {}
15/// Allows `i32` to be used as a broadcast identifier.
16impl BroadcastTypeTrait for i32 {}
17/// Allows `i64` to be used as a broadcast identifier.
18impl BroadcastTypeTrait for i64 {}
19/// Allows `i128` to be used as a broadcast identifier.
20impl BroadcastTypeTrait for i128 {}
21/// Allows `isize` to be used as a broadcast identifier.
22impl BroadcastTypeTrait for isize {}
23/// Allows `u8` to be used as a broadcast identifier.
24impl BroadcastTypeTrait for u8 {}
25/// Allows `u16` to be used as a broadcast identifier.
26impl BroadcastTypeTrait for u16 {}
27/// Allows `u32` to be used as a broadcast identifier.
28impl BroadcastTypeTrait for u32 {}
29/// Allows `u64` to be used as a broadcast identifier.
30impl BroadcastTypeTrait for u64 {}
31/// Allows `u128` to be used as a broadcast identifier.
32impl BroadcastTypeTrait for u128 {}
33/// Allows `usize` to be used as a broadcast identifier.
34impl BroadcastTypeTrait for usize {}
35/// Allows `f32` to be used as a broadcast identifier.
36impl BroadcastTypeTrait for f32 {}
37/// Allows `f64` to be used as a broadcast identifier.
38impl BroadcastTypeTrait for f64 {}
39/// Allows `IpAddr` to be used as a broadcast identifier.
40impl BroadcastTypeTrait for IpAddr {}
41/// Allows `Ipv4Addr` to be used as a broadcast identifier.
42impl BroadcastTypeTrait for Ipv4Addr {}
43/// Allows `Ipv6Addr` to be used as a broadcast identifier.
44impl BroadcastTypeTrait for Ipv6Addr {}
45/// Allows `SocketAddr` to be used as a broadcast identifier.
46impl BroadcastTypeTrait for SocketAddr {}
47/// Allows `NonZeroU8` to be used as a broadcast identifier.
48impl BroadcastTypeTrait for NonZeroU8 {}
49/// Allows `NonZeroU16` to be used as a broadcast identifier.
50impl BroadcastTypeTrait for NonZeroU16 {}
51/// Allows `NonZeroU32` to be used as a broadcast identifier.
52impl BroadcastTypeTrait for NonZeroU32 {}
53/// Allows `NonZeroU64` to be used as a broadcast identifier.
54impl BroadcastTypeTrait for NonZeroU64 {}
55/// Allows `NonZeroU128` to be used as a broadcast identifier.
56impl BroadcastTypeTrait for NonZeroU128 {}
57/// Allows `NonZeroUsize` to be used as a broadcast identifier.
58impl BroadcastTypeTrait for NonZeroUsize {}
59/// Allows `NonZeroI8` to be used as a broadcast identifier.
60impl BroadcastTypeTrait for NonZeroI8 {}
61/// Allows `NonZeroI16` to be used as a broadcast identifier.
62impl BroadcastTypeTrait for NonZeroI16 {}
63/// Allows `NonZeroI32` to be used as a broadcast identifier.
64impl BroadcastTypeTrait for NonZeroI32 {}
65/// Allows `NonZeroI64` to be used as a broadcast identifier.
66impl BroadcastTypeTrait for NonZeroI64 {}
67/// Allows `NonZeroI128` to be used as a broadcast identifier.
68impl BroadcastTypeTrait for NonZeroI128 {}
69/// Allows `NonZeroIsize` to be used as a broadcast identifier.
70impl BroadcastTypeTrait for NonZeroIsize {}
71/// Allows `Infallible` to be used as a broadcast identifier.
72impl BroadcastTypeTrait for Infallible {}
73
74/// Allows references to `String` to be used as broadcast identifiers.
75impl BroadcastTypeTrait for &String {}
76/// Allows double references to string slices to be used as broadcast identifiers.
77impl BroadcastTypeTrait for &&str {}
78/// Allows references to `char` to be used as broadcast identifiers.
79impl BroadcastTypeTrait for &char {}
80/// Allows references to `bool` to be used as broadcast identifiers.
81impl BroadcastTypeTrait for &bool {}
82/// Allows references to `i8` to be used as broadcast identifiers.
83impl BroadcastTypeTrait for &i8 {}
84/// Allows references to `i16` to be used as broadcast identifiers.
85impl BroadcastTypeTrait for &i16 {}
86/// Allows references to `i32` to be used as broadcast identifiers.
87impl BroadcastTypeTrait for &i32 {}
88/// Allows references to `i64` to be used as broadcast identifiers.
89impl BroadcastTypeTrait for &i64 {}
90/// Allows references to `i128` to be used as broadcast identifiers.
91impl BroadcastTypeTrait for &i128 {}
92/// Allows references to `isize` to be used as broadcast identifiers.
93impl BroadcastTypeTrait for &isize {}
94/// Allows references to `u8` to be used as broadcast identifiers.
95impl BroadcastTypeTrait for &u8 {}
96/// Allows references to `u16` to be used as broadcast identifiers.
97impl BroadcastTypeTrait for &u16 {}
98/// Allows references to `u32` to be used as broadcast identifiers.
99impl BroadcastTypeTrait for &u32 {}
100/// Allows references to `u64` to be used as
101/// Implements `BroadcastTypeTrait` for `&u128`.
102///
103/// This allows references to `u128` to be used as a broadcast identifier.
104impl BroadcastTypeTrait for &u128 {}
105/// Implements `BroadcastTypeTrait` for `&usize`.
106///
107/// This allows references to `usize` to be used as a broadcast identifier.
108impl BroadcastTypeTrait for &usize {}
109/// Implements `BroadcastTypeTrait` for `&f32`.
110///
111/// This allows references to `f32` to be used as a broadcast identifier.
112impl BroadcastTypeTrait for &f32 {}
113/// Implements `BroadcastTypeTrait` for `&f64`.
114///
115/// This allows references to `f64` to be used as a broadcast identifier.
116impl BroadcastTypeTrait for &f64 {}
117/// Implements `BroadcastTypeTrait` for `&IpAddr`.
118///
119/// This allows references to `IpAddr` to be used as a broadcast identifier.
120impl BroadcastTypeTrait for &IpAddr {}
121/// Implements `BroadcastTypeTrait` for `&Ipv4Addr`.
122///
123/// This allows references to `Ipv4Addr` to be used as a broadcast identifier.
124impl BroadcastTypeTrait for &Ipv4Addr {}
125/// Implements `BroadcastTypeTrait` for `&Ipv6Addr`.
126///
127/// This allows references to `Ipv6Addr` to be used as a broadcast identifier.
128impl BroadcastTypeTrait for &Ipv6Addr {}
129/// Implements `BroadcastTypeTrait` for `&SocketAddr`.
130///
131/// This allows references to `SocketAddr` to be used as a broadcast identifier.
132impl BroadcastTypeTrait for &SocketAddr {}
133/// Implements `BroadcastTypeTrait` for `&NonZeroU8`.
134///
135/// This allows references to `NonZeroU8` to be used as a broadcast identifier.
136impl BroadcastTypeTrait for &NonZeroU8 {}
137/// Implements `BroadcastTypeTrait` for `&NonZeroU16`.
138///
139/// This allows references to `NonZeroU16` to be used as a broadcast identifier.
140impl BroadcastTypeTrait for &NonZeroU16 {}
141/// Implements `BroadcastTypeTrait` for `&NonZeroU32`.
142///
143/// This allows references to `NonZeroU32` to be used as a broadcast identifier.
144impl BroadcastTypeTrait for &NonZeroU32 {}
145/// Implements `BroadcastTypeTrait` for `&NonZeroU64`.
146///
147/// This allows references to `NonZeroU64` to be used as a broadcast identifier.
148impl BroadcastTypeTrait for &NonZeroU64 {}
149/// Implements `BroadcastTypeTrait` for `&NonZeroU128`.
150///
151/// This allows references to `NonZeroU128` to be used as a broadcast identifier.
152impl BroadcastTypeTrait for &NonZeroU128 {}
153/// Implements `BroadcastTypeTrait` for `&NonZeroUsize`.
154///
155/// This allows references to `NonZeroUsize` to be used as a broadcast identifier.
156impl BroadcastTypeTrait for &NonZeroUsize {}
157/// Implements `BroadcastTypeTrait` for `&NonZeroI8`.
158///
159/// This allows references to `NonZeroI8` to be used as a broadcast identifier.
160impl BroadcastTypeTrait for &NonZeroI8 {}
161/// Implements `BroadcastTypeTrait` for `&NonZeroI16`.
162///
163/// This allows references to `NonZeroI16` to be used as a broadcast identifier.
164impl BroadcastTypeTrait for &NonZeroI16 {}
165/// Implements `BroadcastTypeTrait` for `&NonZeroI32`.
166///
167/// This allows references to `NonZeroI32` to be used as a broadcast identifier.
168impl BroadcastTypeTrait for &NonZeroI32 {}
169/// Implements `BroadcastTypeTrait` for `&NonZeroI64`.
170///
171/// This allows references to `NonZeroI64` to be used as a broadcast identifier.
172impl BroadcastTypeTrait for &NonZeroI64 {}
173/// Implements `BroadcastTypeTrait` for `&NonZeroI128`.
174///
175/// This allows references to `NonZeroI128` to be used as a broadcast identifier.
176impl BroadcastTypeTrait for &NonZeroI128 {}
177/// Implements `BroadcastTypeTrait` for `&NonZeroIsize`.
178///
179/// This allows references to `NonZeroIsize` to be used as a broadcast identifier.
180impl BroadcastTypeTrait for &NonZeroIsize {}
181/// Implements `BroadcastTypeTrait` for `&Infallible`.
182///
183/// This allows references to `Infallible` to be used as a broadcast identifier.
184impl BroadcastTypeTrait for &Infallible {}
185
186/// Implements the `Default` trait for `BroadcastType`.
187///
188/// The default value is `BroadcastType::Unknown`.
189///
190/// # Type Parameters
191///
192/// - `B`: The type parameter for `BroadcastType`, which must implement `BroadcastTypeTrait`.
193impl<B: BroadcastTypeTrait> Default for BroadcastType<B> {
194    #[inline(always)]
195    fn default() -> Self {
196        BroadcastType::Unknown
197    }
198}
199
200impl<B: BroadcastTypeTrait> BroadcastType<B> {
201    /// Generates a unique key string for a given broadcast type.
202    ///
203    /// For point-to-point types, the keys are sorted to ensure consistent key generation
204    /// regardless of the order of the input keys.
205    ///
206    /// # Arguments
207    ///
208    /// - `BroadcastType<B>` - The broadcast type for which to generate the key.
209    ///
210    /// # Returns
211    ///
212    /// - `String` - The unique key string for the broadcast type.
213    #[inline(always)]
214    pub fn get_key(broadcast_type: BroadcastType<B>) -> String {
215        match broadcast_type {
216            BroadcastType::PointToPoint(key1, key2) => {
217                let (first_key, second_key) = if key1 <= key2 {
218                    (key1, key2)
219                } else {
220                    (key2, key1)
221                };
222                format!(
223                    "{}-{}-{}",
224                    POINT_TO_POINT_KEY,
225                    first_key.to_string(),
226                    second_key.to_string()
227                )
228            }
229            BroadcastType::PointToGroup(key) => {
230                format!("{}-{}", POINT_TO_GROUP_KEY, key.to_string())
231            }
232            BroadcastType::Unknown => String::new(),
233        }
234    }
235}
236
237/// Implements the `Default` trait for `WebSocketConfig`.
238///
239/// Provides a default configuration for WebSocket connections, including
240/// default hook types that do nothing.
241///
242/// # Type Parameters
243///
244/// - `B`: The type parameter for `WebSocketConfig`, which must implement `BroadcastTypeTrait`.
245impl<B: BroadcastTypeTrait> Default for WebSocketConfig<B> {
246    #[inline(always)]
247    fn default() -> Self {
248        let default_hook: ServerHookHandler = Arc::new(|_ctx| Box::pin(async {}));
249        Self {
250            context: Context::default(),
251            buffer_size: DEFAULT_BUFFER_SIZE,
252            capacity: DEFAULT_BROADCAST_SENDER_CAPACITY,
253            broadcast_type: BroadcastType::default(),
254            connected_hook: default_hook.clone(),
255            request_hook: default_hook.clone(),
256            sended_hook: default_hook.clone(),
257            closed_hook: default_hook,
258        }
259    }
260}
261
262impl<B: BroadcastTypeTrait> WebSocketConfig<B> {
263    /// Creates a new WebSocket configuration with default values.
264    ///
265    /// # Returns
266    ///
267    /// - `WebSocketConfig<B>` - A new WebSocket configuration instance.
268    #[inline(always)]
269    pub fn new() -> Self {
270        Self::default()
271    }
272}
273
274impl<B: BroadcastTypeTrait> WebSocketConfig<B> {
275    /// Sets the buffer size for the WebSocket connection.
276    ///
277    /// # Arguments
278    ///
279    /// - `usize` - The desired buffer size in bytes.
280    ///
281    /// # Returns
282    ///
283    /// - `WebSocketConfig<B>` - The modified WebSocket configuration instance.
284    #[inline(always)]
285    pub fn set_buffer_size(mut self, buffer_size: usize) -> Self {
286        self.buffer_size = buffer_size;
287        self
288    }
289
290    /// Sets the capacity for the broadcast sender.
291    ///
292    /// # Arguments
293    ///
294    /// - `Capacity` - The desired capacity.
295    ///
296    /// # Returns
297    ///
298    /// - `WebSocketConfig<B>` - The modified WebSocket configuration instance.
299    #[inline(always)]
300    pub fn set_capacity(mut self, capacity: Capacity) -> Self {
301        self.capacity = capacity;
302        self
303    }
304
305    /// Sets the context for the WebSocket connection.
306    ///
307    /// # Arguments
308    ///
309    /// - `Context` - The context object to associate with the WebSocket.
310    ///
311    /// # Returns
312    ///
313    /// - `WebSocketConfig<B>` - The modified WebSocket configuration instance.
314    #[inline(always)]
315    pub fn set_context(mut self, context: Context) -> Self {
316        self.context = context;
317        self
318    }
319
320    /// Sets the broadcast type for the WebSocket connection.
321    ///
322    /// # Arguments
323    ///
324    /// - `BroadcastType<B>` - The broadcast type to use for this WebSocket.
325    ///
326    /// # Returns
327    ///
328    /// - `WebSocketConfig<B>` - The modified WebSocket configuration instance.
329    #[inline(always)]
330    pub fn set_broadcast_type(mut self, broadcast_type: BroadcastType<B>) -> Self {
331        self.broadcast_type = broadcast_type;
332        self
333    }
334
335    /// Retrieves a reference to the context associated with this configuration.
336    ///
337    /// # Returns
338    ///
339    /// - `&Context` - A reference to the context object.
340    #[inline(always)]
341    pub fn get_context(&self) -> &Context {
342        &self.context
343    }
344
345    /// Retrieves the buffer size configured for the WebSocket connection.
346    ///
347    /// # Returns
348    ///
349    /// - `usize` - The buffer size in bytes.
350    #[inline(always)]
351    pub fn get_buffer_size(&self) -> usize {
352        self.buffer_size
353    }
354
355    /// Retrieves the capacity configured for the broadcast sender.
356    ///
357    /// # Returns
358    ///
359    /// - `Capacity` - The capacity.
360    #[inline(always)]
361    pub fn get_capacity(&self) -> Capacity {
362        self.capacity
363    }
364
365    /// Retrieves a reference to the broadcast type configured for this WebSocket.
366    ///
367    /// # Returns
368    ///
369    /// - `&BroadcastType<B>` - A reference to the broadcast type object.
370    #[inline(always)]
371    pub fn get_broadcast_type(&self) -> &BroadcastType<B> {
372        &self.broadcast_type
373    }
374
375    /// Sets the connected hook handler.
376    ///
377    /// This hook is executed when the WebSocket connection is established.
378    ///
379    /// # Type Parameters
380    ///
381    /// - `S`: The hook type, which must implement `ServerHook`.
382    ///
383    /// # Returns
384    ///
385    /// The modified `WebSocketConfig` instance.
386    ///
387    /// # Examples
388    ///
389    /// ```rust,ignore
390    /// struct MyConnectedHook;
391    /// impl ServerHook for MyConnectedHook {
392    ///     async fn new(_ctx: &Context) -> Self { Self }
393    ///     async fn handle(self, ctx: &Context) { /* ... */ }
394    /// }
395    ///
396    /// let config = WebSocketConfig::new()
397    ///     .set_connected_hook::<MyConnectedHook>();
398    /// ```
399    #[inline(always)]
400    pub fn set_connected_hook<S>(mut self) -> Self
401    where
402        S: ServerHook,
403    {
404        self.connected_hook = server_hook_factory::<S>();
405        self
406    }
407
408    /// Sets the request hook handler.
409    ///
410    /// This hook is executed when a new request is received on the WebSocket.
411    ///
412    /// # Type Parameters
413    ///
414    /// - `S`: The hook type, which must implement `ServerHook`.
415    ///
416    /// # Returns
417    ///
418    /// The modified `WebSocketConfig` instance.
419    ///
420    /// # Examples
421    ///
422    /// ```rust,ignore
423    /// struct MyRequestHook;
424    /// impl ServerHook for MyRequestHook {
425    ///     async fn new(_ctx: &Context) -> Self { Self }
426    ///     async fn handle(self, ctx: &Context) { /* ... */ }
427    /// }
428    ///
429    /// let config = WebSocketConfig::new()
430    ///     .set_request_hook::<MyRequestHook>();
431    /// ```
432    #[inline(always)]
433    pub fn set_request_hook<S>(mut self) -> Self
434    where
435        S: ServerHook,
436    {
437        self.request_hook = server_hook_factory::<S>();
438        self
439    }
440
441    /// Sets the sended hook handler.
442    ///
443    /// This hook is executed after a message has been successfully sent over the WebSocket.
444    ///
445    /// # Type Parameters
446    ///
447    /// - `S`: The hook type, which must implement `ServerHook`.
448    ///
449    /// # Returns
450    ///
451    /// The modified `WebSocketConfig` instance.
452    ///
453    /// # Examples
454    ///
455    /// ```rust,ignore
456    /// struct MySendedHook;
457    /// impl ServerHook for MySendedHook {
458    ///     async fn new(_ctx: &Context) -> Self { Self }
459    ///     async fn handle(self, ctx: &Context) { /* ... */ }
460    /// }
461    ///
462    /// let config = WebSocketConfig::new()
463    ///     .set_sended_hook::<MySendedHook>();
464    /// ```
465    #[inline(always)]
466    pub fn set_sended_hook<S>(mut self) -> Self
467    where
468        S: ServerHook,
469    {
470        self.sended_hook = server_hook_factory::<S>();
471        self
472    }
473
474    /// Sets the closed hook handler.
475    ///
476    /// This hook is executed when the WebSocket connection is closed.
477    ///
478    /// # Type Parameters
479    ///
480    /// - `S`: The hook type, which must implement `ServerHook`.
481    ///
482    /// # Returns
483    ///
484    /// The modified `WebSocketConfig` instance.
485    ///
486    /// # Examples
487    ///
488    /// ```rust,ignore
489    /// struct MyClosedHook;
490    /// impl ServerHook for MyClosedHook {
491    ///     async fn new(_ctx: &Context) -> Self { Self }
492    ///     async fn handle(self, ctx: &Context) { /* ... */ }
493    /// }
494    ///
495    /// let config = WebSocketConfig::new()
496    ///     .set_closed_hook::<MyClosedHook>();
497    /// ```
498    #[inline(always)]
499    pub fn set_closed_hook<S>(mut self) -> Self
500    where
501        S: ServerHook,
502    {
503        self.closed_hook = server_hook_factory::<S>();
504        self
505    }
506
507    /// Retrieves a reference to the connected hook handler.
508    ///
509    /// # Returns
510    ///
511    /// - `&ServerHookHandler` - A reference to the connected hook handler.
512    #[inline(always)]
513    pub fn get_connected_hook(&self) -> &ServerHookHandler {
514        &self.connected_hook
515    }
516
517    /// Retrieves a reference to the request hook handler.
518    ///
519    /// # Returns
520    ///
521    /// - `&ServerHookHandler` - A reference to the request hook handler.
522    #[inline(always)]
523    pub fn get_request_hook(&self) -> &ServerHookHandler {
524        &self.request_hook
525    }
526
527    /// Retrieves a reference to the sended hook handler.
528    ///
529    /// # Returns
530    ///
531    /// - `&ServerHookHandler` - A reference to the sended hook handler.
532    #[inline(always)]
533    pub fn get_sended_hook(&self) -> &ServerHookHandler {
534        &self.sended_hook
535    }
536
537    /// Retrieves a reference to the closed hook handler.
538    ///
539    /// # Returns
540    ///
541    /// - `&ServerHookHandler` - A reference to the closed hook handler.
542    #[inline(always)]
543    pub fn get_closed_hook(&self) -> &ServerHookHandler {
544        &self.closed_hook
545    }
546}
547
548impl WebSocket {
549    /// Creates a new WebSocket instance.
550    ///
551    /// Initializes with a default broadcast map.
552    ///
553    /// # Returns
554    ///
555    /// - `WebSocket` - A new WebSocket instance.
556    #[inline(always)]
557    pub fn new() -> Self {
558        Self::default()
559    }
560
561    /// Subscribes to a broadcast type or inserts a new one if it doesn't exist.
562    ///
563    /// # Type Parameters
564    ///
565    /// - `B`: The type implementing `BroadcastTypeTrait`.
566    ///
567    /// # Arguments
568    ///
569    /// - `BroadcastType<B>` - The broadcast type to subscribe to.
570    /// - `Capacity` - The capacity for the broadcast sender if a new one is inserted.
571    ///
572    /// # Returns
573    ///
574    /// - `BroadcastMapReceiver<Vec<u8>>` - A broadcast map receiver for the specified broadcast type.
575    #[inline(always)]
576    fn subscribe_unwrap_or_insert<B: BroadcastTypeTrait>(
577        &self,
578        broadcast_type: BroadcastType<B>,
579        capacity: Capacity,
580    ) -> BroadcastMapReceiver<Vec<u8>> {
581        let key: String = BroadcastType::get_key(broadcast_type);
582        self.broadcast_map.subscribe_or_insert(&key, capacity)
583    }
584
585    /// Subscribes to a point-to-point broadcast.
586    ///
587    /// # Type Parameters
588    ///
589    /// - `B`: The type implementing `BroadcastTypeTrait`.
590    ///
591    /// # Arguments
592    ///
593    /// - `&B` - The first identifier for the point-to-point communication.
594    /// - `&B` - The second identifier for the point-to-point communication.
595    /// - `Capacity` - The capacity for the broadcast sender.
596    ///
597    /// # Returns
598    ///
599    /// - `BroadcastMapReceiver<Vec<u8>>` - A broadcast map receiver for the point-to-point broadcast.
600    #[inline(always)]
601    fn point_to_point<B: BroadcastTypeTrait>(
602        &self,
603        key1: &B,
604        key2: &B,
605        capacity: Capacity,
606    ) -> BroadcastMapReceiver<Vec<u8>> {
607        self.subscribe_unwrap_or_insert(
608            BroadcastType::PointToPoint(key1.clone(), key2.clone()),
609            capacity,
610        )
611    }
612
613    /// Subscribes to a point-to-group broadcast.
614    ///
615    /// # Type Parameters
616    ///
617    /// - `B`: The type implementing `BroadcastTypeTrait`.
618    ///
619    /// # Arguments
620    ///
621    /// - `&B` - The identifier for the group.
622    /// - `Capacity` - The capacity for the broadcast sender.
623    ///
624    /// # Returns
625    ///
626    /// - `BroadcastMapReceiver<Vec<u8>>` - A broadcast map receiver for the point-to-group broadcast.
627    #[inline(always)]
628    fn point_to_group<B: BroadcastTypeTrait>(
629        &self,
630        key: &B,
631        capacity: Capacity,
632    ) -> BroadcastMapReceiver<Vec<u8>> {
633        self.subscribe_unwrap_or_insert(BroadcastType::PointToGroup(key.clone()), capacity)
634    }
635
636    /// Retrieves the current receiver count for a given broadcast type.
637    ///
638    /// # Type Parameters
639    ///
640    /// - `B`: The type implementing `BroadcastTypeTrait`.
641    ///
642    /// # Arguments
643    ///
644    /// - `BroadcastType<B>` - The broadcast type for which to get the receiver count.
645    ///
646    /// # Returns
647    ///
648    /// - `ReceiverCount` - The number of active receivers for the broadcast type, or 0 if not found.
649    #[inline(always)]
650    pub fn receiver_count<B: BroadcastTypeTrait>(
651        &self,
652        broadcast_type: BroadcastType<B>,
653    ) -> ReceiverCount {
654        let key: String = BroadcastType::get_key(broadcast_type);
655        self.broadcast_map.receiver_count(&key).unwrap_or(0)
656    }
657
658    /// Calculates the receiver count before a connection is established.
659    ///
660    /// Ensures the count does not exceed the maximum allowed value minus one.
661    ///
662    /// # Type Parameters
663    ///
664    /// - `B`: The type implementing `BroadcastTypeTrait`.
665    ///
666    /// # Arguments
667    ///
668    /// - `BroadcastType<B>` - The broadcast type for which to get the receiver count.
669    ///
670    /// # Returns
671    ///
672    /// - `ReceiverCount` - The receiver count after the connection is established.
673    #[inline(always)]
674    pub fn receiver_count_before_connected<B: BroadcastTypeTrait>(
675        &self,
676        broadcast_type: BroadcastType<B>,
677    ) -> ReceiverCount {
678        let count: ReceiverCount = self.receiver_count(broadcast_type);
679        count.clamp(0, ReceiverCount::MAX - 1) + 1
680    }
681
682    /// Calculates the receiver count after a connection is closed.
683    ///
684    /// Ensures the count does not go below 0.
685    ///
686    /// # Type Parameters
687    ///
688    /// - `B`: The type implementing `BroadcastTypeTrait`.
689    ///
690    /// # Arguments
691    ///
692    /// - `BroadcastType<B>` - The broadcast type for which to get the receiver count.
693    ///
694    /// # Returns
695    ///
696    /// - `ReceiverCount` - The receiver count after the connection is closed.
697    #[inline(always)]
698    pub fn receiver_count_after_closed<B: BroadcastTypeTrait>(
699        &self,
700        broadcast_type: BroadcastType<B>,
701    ) -> ReceiverCount {
702        let count: ReceiverCount = self.receiver_count(broadcast_type);
703        count.clamp(1, ReceiverCount::MAX) - 1
704    }
705
706    /// Sends data to all active receivers for a given broadcast type.
707    ///
708    /// # Type Parameters
709    ///
710    /// - `T`: The type of data to send, which must be convertible to `Vec<u8>`.
711    /// - `B`: The type implementing `BroadcastTypeTrait`.
712    ///
713    /// # Arguments
714    ///
715    /// - `BroadcastType<B>` - The broadcast type to which to send the data.
716    /// - `T` - The data to send.
717    ///
718    /// # Returns
719    ///
720    /// - `BroadcastMapSendResult<Vec<u8>>` - A result indicating the success or failure of the send operation.
721    #[inline(always)]
722    pub fn send<T, B>(
723        &self,
724        broadcast_type: BroadcastType<B>,
725        data: T,
726    ) -> BroadcastMapSendResult<Vec<u8>>
727    where
728        T: Into<Vec<u8>>,
729        B: BroadcastTypeTrait,
730    {
731        let key: String = BroadcastType::get_key(broadcast_type);
732        self.broadcast_map.send(&key, data.into())
733    }
734
735    /// Runs the WebSocket connection, handling incoming requests and outgoing messages.
736    ///
737    /// This asynchronous function continuously monitors for new WebSocket requests
738    /// and incoming broadcast messages, processing them according to the configured hooks.
739    ///
740    /// # Type Parameters
741    ///
742    /// - `B`: The type implementing `BroadcastTypeTrait`.
743    ///
744    /// # Arguments
745    ///
746    /// - `WebSocketConfig<B>` - The WebSocket configuration containing the configuration for this WebSocket instance.
747    ///
748    /// # Panics
749    ///
750    /// Panics if the context in the WebSocket configuration is not set (i.e., it's the default context).
751    /// Panics if the broadcast type in the WebSocket configuration is `BroadcastType::Unknown`.
752    pub async fn run<B: BroadcastTypeTrait>(&self, config: WebSocketConfig<B>) {
753        let ctx: Context = config.get_context().clone();
754        if ctx.to_string() == Context::default().to_string() {
755            panic!("Context must be set");
756        }
757        let buffer_size: usize = config.get_buffer_size();
758        let capacity: Capacity = config.get_capacity();
759        let broadcast_type: BroadcastType<B> = config.get_broadcast_type().clone();
760        let mut receiver: Receiver<Vec<u8>> = match &broadcast_type {
761            BroadcastType::PointToPoint(key1, key2) => self.point_to_point(key1, key2, capacity),
762            BroadcastType::PointToGroup(key) => self.point_to_group(key, capacity),
763            BroadcastType::Unknown => panic!("BroadcastType must be PointToPoint or PointToGroup"),
764        };
765        let key: String = BroadcastType::get_key(broadcast_type);
766        config.get_connected_hook()(&ctx).await;
767        let result_handle = || async {
768            ctx.aborted().await;
769            ctx.closed().await;
770        };
771        loop {
772            tokio::select! {
773                request_res = ctx.ws_from_stream(buffer_size) => {
774                    if request_res.is_ok() {
775                        config.get_request_hook()(&ctx).await;
776                    } else {
777                        config.get_closed_hook()(&ctx).await;
778                    }
779                    if ctx.get_aborted().await {
780                        continue;
781                    }
782                    if ctx.get_closed().await {
783                        break;
784                    }
785                    let body: ResponseBody = ctx.get_response_body().await;
786                    let is_err: bool = self.broadcast_map.send(&key, body).is_err();
787                    config.get_sended_hook()(&ctx).await;
788                    if is_err || ctx.get_closed().await{
789                        break;
790                    }
791                },
792                msg_res = receiver.recv() => {
793                    if let Ok(msg) = &msg_res {
794                        let frame_list: Vec<ResponseBody> = WebSocketFrame::create_frame_list(msg);
795                        if ctx.send_body_list_with_data(&frame_list).await.is_ok() {
796                            continue;
797                        }
798                    }
799                    break;
800                }
801            }
802        }
803        result_handle().await;
804    }
805}
hyperlane_plugin_websocket/websocket/impl.rs

hyperlane_plugin_websocket/websocket/
impl.rs
