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

hyperlane_plugin_websocket/websocket/
impl.rs
