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    fn default() -> Self {
195        BroadcastType::Unknown
196    }
197}
198
199impl<B: BroadcastTypeTrait> BroadcastType<B> {
200    /// Generates a unique key string for a given broadcast type.
201    ///
202    /// For point-to-point types, the keys are sorted to ensure consistent key generation
203    /// regardless of the order of the input keys.
204    ///
205    /// # Arguments
206    ///
207    /// - `BroadcastType<B>` - The broadcast type for which to generate the key.
208    ///
209    /// # Returns
210    ///
211    /// - `String` - The unique key string for the broadcast type.
212    #[inline]
213    pub fn get_key(broadcast_type: BroadcastType<B>) -> String {
214        match broadcast_type {
215            BroadcastType::PointToPoint(key1, key2) => {
216                let (first_key, second_key) = if key1 <= key2 {
217                    (key1, key2)
218                } else {
219                    (key2, key1)
220                };
221                format!(
222                    "{}-{}-{}",
223                    POINT_TO_POINT_KEY,
224                    first_key.to_string(),
225                    second_key.to_string()
226                )
227            }
228            BroadcastType::PointToGroup(key) => {
229                format!("{}-{}", POINT_TO_GROUP_KEY, key.to_string())
230            }
231            BroadcastType::Unknown => String::new(),
232        }
233    }
234}
235
236/// Implements the `Default` trait for `WebSocketConfig`.
237///
238/// Provides a default configuration for WebSocket connections, including
239/// default hook types that do nothing.
240///
241/// # Type Parameters
242///
243/// - `B`: The type parameter for `WebSocketConfig`, which must implement `BroadcastTypeTrait`.
244impl<B: BroadcastTypeTrait> Default for WebSocketConfig<B> {
245    fn default() -> Self {
246        let default_hook: ServerHookHandler = Arc::new(|_ctx| Box::pin(async {}));
247        Self {
248            context: Context::default(),
249            buffer_size: DEFAULT_BUFFER_SIZE,
250            capacity: DEFAULT_BROADCAST_SENDER_CAPACITY,
251            broadcast_type: BroadcastType::default(),
252            connected_hook: default_hook.clone(),
253            request_hook: default_hook.clone(),
254            sended_hook: default_hook.clone(),
255            closed_hook: default_hook,
256        }
257    }
258}
259
260impl<B: BroadcastTypeTrait> WebSocketConfig<B> {
261    /// Creates a new WebSocket configuration with default values.
262    ///
263    /// # Returns
264    ///
265    /// - `WebSocketConfig<B>` - A new WebSocket configuration instance.
266    #[inline]
267    pub fn new() -> Self {
268        Self::default()
269    }
270}
271
272impl<B: BroadcastTypeTrait> WebSocketConfig<B> {
273    /// Sets the buffer size for the WebSocket connection.
274    ///
275    /// # Arguments
276    ///
277    /// - `usize` - The desired buffer size in bytes.
278    ///
279    /// # Returns
280    ///
281    /// - `WebSocketConfig<B>` - The modified WebSocket configuration instance.
282    #[inline]
283    pub fn set_buffer_size(mut self, buffer_size: usize) -> Self {
284        self.buffer_size = buffer_size;
285        self
286    }
287
288    /// Sets the capacity for the broadcast sender.
289    ///
290    /// # Arguments
291    ///
292    /// - `Capacity` - The desired capacity.
293    ///
294    /// # Returns
295    ///
296    /// - `WebSocketConfig<B>` - The modified WebSocket configuration instance.
297    #[inline]
298    pub fn set_capacity(mut self, capacity: Capacity) -> Self {
299        self.capacity = capacity;
300        self
301    }
302
303    /// Sets the context for the WebSocket connection.
304    ///
305    /// # Arguments
306    ///
307    /// - `Context` - The context object to associate with the WebSocket.
308    ///
309    /// # Returns
310    ///
311    /// - `WebSocketConfig<B>` - The modified WebSocket configuration instance.
312    #[inline]
313    pub fn set_context(mut self, context: Context) -> Self {
314        self.context = context;
315        self
316    }
317
318    /// Sets the broadcast type for the WebSocket connection.
319    ///
320    /// # Arguments
321    ///
322    /// - `BroadcastType<B>` - The broadcast type to use for this WebSocket.
323    ///
324    /// # Returns
325    ///
326    /// - `WebSocketConfig<B>` - The modified WebSocket configuration instance.
327    #[inline]
328    pub fn set_broadcast_type(mut self, broadcast_type: BroadcastType<B>) -> Self {
329        self.broadcast_type = broadcast_type;
330        self
331    }
332
333    /// Retrieves a reference to the context associated with this configuration.
334    ///
335    /// # Returns
336    ///
337    /// - `&Context` - A reference to the context object.
338    #[inline]
339    pub fn get_context(&self) -> &Context {
340        &self.context
341    }
342
343    /// Retrieves the buffer size configured for the WebSocket connection.
344    ///
345    /// # Returns
346    ///
347    /// - `usize` - The buffer size in bytes.
348    #[inline]
349    pub fn get_buffer_size(&self) -> usize {
350        self.buffer_size
351    }
352
353    /// Retrieves the capacity configured for the broadcast sender.
354    ///
355    /// # Returns
356    ///
357    /// - `Capacity` - The capacity.
358    #[inline]
359    pub fn get_capacity(&self) -> Capacity {
360        self.capacity
361    }
362
363    /// Retrieves a reference to the broadcast type configured for this WebSocket.
364    ///
365    /// # Returns
366    ///
367    /// - `&BroadcastType<B>` - A reference to the broadcast type object.
368    #[inline]
369    pub fn get_broadcast_type(&self) -> &BroadcastType<B> {
370        &self.broadcast_type
371    }
372
373    /// Sets the connected hook handler.
374    ///
375    /// This hook is executed when the WebSocket connection is established.
376    ///
377    /// # Type Parameters
378    ///
379    /// - `S`: The hook type, which must implement `ServerHook`.
380    ///
381    /// # Returns
382    ///
383    /// The modified `WebSocketConfig` instance.
384    ///
385    /// # Examples
386    ///
387    /// ```rust,ignore
388    /// struct MyConnectedHook;
389    /// impl ServerHook for MyConnectedHook {
390    ///     async fn new(_ctx: &Context) -> Self { Self }
391    ///     async fn handle(self, ctx: &Context) { /* ... */ }
392    /// }
393    ///
394    /// let config = WebSocketConfig::new()
395    ///     .set_connected_hook::<MyConnectedHook>();
396    /// ```
397    #[inline]
398    pub fn set_connected_hook<S>(mut self) -> Self
399    where
400        S: ServerHook,
401    {
402        self.connected_hook = server_hook_factory::<S>();
403        self
404    }
405
406    /// Sets the request hook handler.
407    ///
408    /// This hook is executed when a new request is received on the WebSocket.
409    ///
410    /// # Type Parameters
411    ///
412    /// - `S`: The hook type, which must implement `ServerHook`.
413    ///
414    /// # Returns
415    ///
416    /// The modified `WebSocketConfig` instance.
417    ///
418    /// # Examples
419    ///
420    /// ```rust,ignore
421    /// struct MyRequestHook;
422    /// impl ServerHook for MyRequestHook {
423    ///     async fn new(_ctx: &Context) -> Self { Self }
424    ///     async fn handle(self, ctx: &Context) { /* ... */ }
425    /// }
426    ///
427    /// let config = WebSocketConfig::new()
428    ///     .set_request_hook::<MyRequestHook>();
429    /// ```
430    #[inline]
431    pub fn set_request_hook<S>(mut self) -> Self
432    where
433        S: ServerHook,
434    {
435        self.request_hook = server_hook_factory::<S>();
436        self
437    }
438
439    /// Sets the sended hook handler.
440    ///
441    /// This hook is executed after a message has been successfully sent over the WebSocket.
442    ///
443    /// # Type Parameters
444    ///
445    /// - `S`: The hook type, which must implement `ServerHook`.
446    ///
447    /// # Returns
448    ///
449    /// The modified `WebSocketConfig` instance.
450    ///
451    /// # Examples
452    ///
453    /// ```rust,ignore
454    /// struct MySendedHook;
455    /// impl ServerHook for MySendedHook {
456    ///     async fn new(_ctx: &Context) -> Self { Self }
457    ///     async fn handle(self, ctx: &Context) { /* ... */ }
458    /// }
459    ///
460    /// let config = WebSocketConfig::new()
461    ///     .set_sended_hook::<MySendedHook>();
462    /// ```
463    #[inline]
464    pub fn set_sended_hook<S>(mut self) -> Self
465    where
466        S: ServerHook,
467    {
468        self.sended_hook = server_hook_factory::<S>();
469        self
470    }
471
472    /// Sets the closed hook handler.
473    ///
474    /// This hook is executed when the WebSocket connection is closed.
475    ///
476    /// # Type Parameters
477    ///
478    /// - `S`: The hook type, which must implement `ServerHook`.
479    ///
480    /// # Returns
481    ///
482    /// The modified `WebSocketConfig` instance.
483    ///
484    /// # Examples
485    ///
486    /// ```rust,ignore
487    /// struct MyClosedHook;
488    /// impl ServerHook for MyClosedHook {
489    ///     async fn new(_ctx: &Context) -> Self { Self }
490    ///     async fn handle(self, ctx: &Context) { /* ... */ }
491    /// }
492    ///
493    /// let config = WebSocketConfig::new()
494    ///     .set_closed_hook::<MyClosedHook>();
495    /// ```
496    #[inline]
497    pub fn set_closed_hook<S>(mut self) -> Self
498    where
499        S: ServerHook,
500    {
501        self.closed_hook = server_hook_factory::<S>();
502        self
503    }
504
505    /// Retrieves a reference to the connected hook handler.
506    ///
507    /// # Returns
508    ///
509    /// - `&ServerHookHandler` - A reference to the connected hook handler.
510    #[inline]
511    pub fn get_connected_hook(&self) -> &ServerHookHandler {
512        &self.connected_hook
513    }
514
515    /// Retrieves a reference to the request hook handler.
516    ///
517    /// # Returns
518    ///
519    /// - `&ServerHookHandler` - A reference to the request hook handler.
520    #[inline]
521    pub fn get_request_hook(&self) -> &ServerHookHandler {
522        &self.request_hook
523    }
524
525    /// Retrieves a reference to the sended hook handler.
526    ///
527    /// # Returns
528    ///
529    /// - `&ServerHookHandler` - A reference to the sended hook handler.
530    #[inline]
531    pub fn get_sended_hook(&self) -> &ServerHookHandler {
532        &self.sended_hook
533    }
534
535    /// Retrieves a reference to the closed hook handler.
536    ///
537    /// # Returns
538    ///
539    /// - `&ServerHookHandler` - A reference to the closed hook handler.
540    #[inline]
541    pub fn get_closed_hook(&self) -> &ServerHookHandler {
542        &self.closed_hook
543    }
544}
545
546impl WebSocket {
547    /// Creates a new WebSocket instance.
548    ///
549    /// Initializes with a default broadcast map.
550    ///
551    /// # Returns
552    ///
553    /// - `WebSocket` - A new WebSocket instance.
554    #[inline]
555    pub fn new() -> Self {
556        Self::default()
557    }
558
559    /// Subscribes to a broadcast type or inserts a new one if it doesn't exist.
560    ///
561    /// # Type Parameters
562    ///
563    /// - `B`: The type implementing `BroadcastTypeTrait`.
564    ///
565    /// # Arguments
566    ///
567    /// - `BroadcastType<B>` - The broadcast type to subscribe to.
568    /// - `Capacity` - The capacity for the broadcast sender if a new one is inserted.
569    ///
570    /// # Returns
571    ///
572    /// - `BroadcastMapReceiver<Vec<u8>>` - A broadcast map receiver for the specified broadcast type.
573    #[inline]
574    fn subscribe_unwrap_or_insert<B: BroadcastTypeTrait>(
575        &self,
576        broadcast_type: BroadcastType<B>,
577        capacity: Capacity,
578    ) -> BroadcastMapReceiver<Vec<u8>> {
579        let key: String = BroadcastType::get_key(broadcast_type);
580        self.broadcast_map.subscribe_or_insert(&key, capacity)
581    }
582
583    /// Subscribes to a point-to-point broadcast.
584    ///
585    /// # Type Parameters
586    ///
587    /// - `B`: The type implementing `BroadcastTypeTrait`.
588    ///
589    /// # Arguments
590    ///
591    /// - `&B` - The first identifier for the point-to-point communication.
592    /// - `&B` - The second identifier for the point-to-point communication.
593    /// - `Capacity` - The capacity for the broadcast sender.
594    ///
595    /// # Returns
596    ///
597    /// - `BroadcastMapReceiver<Vec<u8>>` - A broadcast map receiver for the point-to-point broadcast.
598    #[inline]
599    fn point_to_point<B: BroadcastTypeTrait>(
600        &self,
601        key1: &B,
602        key2: &B,
603        capacity: Capacity,
604    ) -> BroadcastMapReceiver<Vec<u8>> {
605        self.subscribe_unwrap_or_insert(
606            BroadcastType::PointToPoint(key1.clone(), key2.clone()),
607            capacity,
608        )
609    }
610
611    /// Subscribes to a point-to-group broadcast.
612    ///
613    /// # Type Parameters
614    ///
615    /// - `B`: The type implementing `BroadcastTypeTrait`.
616    ///
617    /// # Arguments
618    ///
619    /// - `&B` - The identifier for the group.
620    /// - `Capacity` - The capacity for the broadcast sender.
621    ///
622    /// # Returns
623    ///
624    /// - `BroadcastMapReceiver<Vec<u8>>` - A broadcast map receiver for the point-to-group broadcast.
625    #[inline]
626    fn point_to_group<B: BroadcastTypeTrait>(
627        &self,
628        key: &B,
629        capacity: Capacity,
630    ) -> BroadcastMapReceiver<Vec<u8>> {
631        self.subscribe_unwrap_or_insert(BroadcastType::PointToGroup(key.clone()), capacity)
632    }
633
634    /// Retrieves the current receiver count for a given broadcast type.
635    ///
636    /// # Type Parameters
637    ///
638    /// - `B`: The type implementing `BroadcastTypeTrait`.
639    ///
640    /// # Arguments
641    ///
642    /// - `BroadcastType<B>` - The broadcast type for which to get the receiver count.
643    ///
644    /// # Returns
645    ///
646    /// - `ReceiverCount` - The number of active receivers for the broadcast type, or 0 if not found.
647    #[inline]
648    pub fn receiver_count<B: BroadcastTypeTrait>(
649        &self,
650        broadcast_type: BroadcastType<B>,
651    ) -> ReceiverCount {
652        let key: String = BroadcastType::get_key(broadcast_type);
653        self.broadcast_map.receiver_count(&key).unwrap_or(0)
654    }
655
656    /// Calculates the receiver count before a connection is established.
657    ///
658    /// Ensures the count does not exceed the maximum allowed value minus one.
659    ///
660    /// # Type Parameters
661    ///
662    /// - `B`: The type implementing `BroadcastTypeTrait`.
663    ///
664    /// # Arguments
665    ///
666    /// - `BroadcastType<B>` - The broadcast type for which to get the receiver count.
667    ///
668    /// # Returns
669    ///
670    /// - `ReceiverCount` - The receiver count after the connection is established.
671    #[inline]
672    pub fn receiver_count_before_connected<B: BroadcastTypeTrait>(
673        &self,
674        broadcast_type: BroadcastType<B>,
675    ) -> ReceiverCount {
676        let count: ReceiverCount = self.receiver_count(broadcast_type);
677        count.clamp(0, ReceiverCount::MAX - 1) + 1
678    }
679
680    /// Calculates the receiver count after a connection is closed.
681    ///
682    /// Ensures the count does not go below 0.
683    ///
684    /// # Type Parameters
685    ///
686    /// - `B`: The type implementing `BroadcastTypeTrait`.
687    ///
688    /// # Arguments
689    ///
690    /// - `BroadcastType<B>` - The broadcast type for which to get the receiver count.
691    ///
692    /// # Returns
693    ///
694    /// - `ReceiverCount` - The receiver count after the connection is closed.
695    #[inline]
696    pub fn receiver_count_after_closed<B: BroadcastTypeTrait>(
697        &self,
698        broadcast_type: BroadcastType<B>,
699    ) -> ReceiverCount {
700        let count: ReceiverCount = self.receiver_count(broadcast_type);
701        count.clamp(1, ReceiverCount::MAX) - 1
702    }
703
704    /// Sends data to all active receivers for a given broadcast type.
705    ///
706    /// # Type Parameters
707    ///
708    /// - `T`: The type of data to send, which must be convertible to `Vec<u8>`.
709    /// - `B`: The type implementing `BroadcastTypeTrait`.
710    ///
711    /// # Arguments
712    ///
713    /// - `BroadcastType<B>` - The broadcast type to which to send the data.
714    /// - `T` - The data to send.
715    ///
716    /// # Returns
717    ///
718    /// - `BroadcastMapSendResult<Vec<u8>>` - A result indicating the success or failure of the send operation.
719    #[inline]
720    pub fn send<T, B>(
721        &self,
722        broadcast_type: BroadcastType<B>,
723        data: T,
724    ) -> BroadcastMapSendResult<Vec<u8>>
725    where
726        T: Into<Vec<u8>>,
727        B: BroadcastTypeTrait,
728    {
729        let key: String = BroadcastType::get_key(broadcast_type);
730        self.broadcast_map.send(&key, data.into())
731    }
732
733    /// Runs the WebSocket connection, handling incoming requests and outgoing messages.
734    ///
735    /// This asynchronous function continuously monitors for new WebSocket requests
736    /// and incoming broadcast messages, processing them according to the configured hooks.
737    ///
738    /// # Type Parameters
739    ///
740    /// - `B`: The type implementing `BroadcastTypeTrait`.
741    ///
742    /// # Arguments
743    ///
744    /// - `WebSocketConfig<B>` - The WebSocket configuration containing the configuration for this WebSocket instance.
745    ///
746    /// # Panics
747    ///
748    /// Panics if the context in the WebSocket configuration is not set (i.e., it's the default context).
749    /// Panics if the broadcast type in the WebSocket configuration is `BroadcastType::Unknown`.
750    pub async fn run<B: BroadcastTypeTrait>(&self, config: WebSocketConfig<B>) {
751        let ctx: Context = config.get_context().clone();
752        if ctx.to_string() == Context::default().to_string() {
753            panic!("Context must be set");
754        }
755        let buffer_size: usize = config.get_buffer_size();
756        let capacity: Capacity = config.get_capacity();
757        let broadcast_type: BroadcastType<B> = config.get_broadcast_type().clone();
758        let mut receiver: Receiver<Vec<u8>> = match &broadcast_type {
759            BroadcastType::PointToPoint(key1, key2) => self.point_to_point(key1, key2, capacity),
760            BroadcastType::PointToGroup(key) => self.point_to_group(key, capacity),
761            BroadcastType::Unknown => panic!("BroadcastType must be PointToPoint or PointToGroup"),
762        };
763        let key: String = BroadcastType::get_key(broadcast_type);
764        config.get_connected_hook()(&ctx).await;
765        let result_handle = || async {
766            ctx.aborted().await;
767            ctx.closed().await;
768        };
769        loop {
770            tokio::select! {
771                request_res = ctx.ws_from_stream(buffer_size) => {
772                    let mut need_break = false;
773                    if request_res.is_ok() {
774                        config.get_request_hook()(&ctx).await;
775                    } else {
776                        need_break = true;
777                        config.get_closed_hook()(&ctx).await;
778                    }
779                    let body: ResponseBody = ctx.get_response_body().await;
780                    let is_err: bool = self.broadcast_map.send(&key, body).is_err();
781                    config.get_sended_hook()(&ctx).await;
782                    if need_break || is_err {
783                        break;
784                    }
785                },
786                msg_res = receiver.recv() => {
787                    if let Ok(msg) = &msg_res {
788                        let frame_list: Vec<ResponseBody> = WebSocketFrame::create_frame_list(msg);
789                        if ctx.send_body_list_with_data(&frame_list).await.is_ok() {
790                            continue;
791                        }
792                    }
793                    break;
794                }
795            }
796        }
797        result_handle().await;
798    }
799}
hyperlane_plugin_websocket/websocket/impl.rs

hyperlane_plugin_websocket/websocket/
impl.rs
