Skip to main content

fhtmx_actix/
sse.rs

1//! # Server sent events
2
3use actix_web::{Responder, web};
4use actix_web_lab::sse::{Data, Event};
5use dashmap::DashMap;
6use serde::{Deserialize, Serialize};
7use std::{marker::PhantomData, sync::Arc, time::Duration};
8use tokio::sync::mpsc;
9use uuid::Uuid;
10
11/// Setups Server sent event state and routes
12///
13/// # Example
14///
15/// ```rust,ignore
16/// use actix_web::App;
17/// use fhtmx_actix::sse::SseSetup;
18///
19/// let sse_setup = SseSetup::new();
20/// let sse_data = sse_setup.state_data();
21/// App::new()
22///     .configure(|cfg| sse_setup.setup_route("/sse", cfg))
23///     .app_data(sse_data);
24/// ```
25#[derive(Clone, Copy)]
26pub struct SseSetup<T> {
27    session_data: PhantomData<T>,
28}
29
30/// Setup sse without session data
31#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
32pub struct FhtmxUiNoSessionData;
33
34impl SseSetup<()> {
35    /// Creates a new setup that carries session data of type `T`.
36    pub fn new_with_data<T>() -> SseSetup<T> {
37        SseSetup {
38            session_data: PhantomData,
39        }
40    }
41}
42
43impl Default for SseSetup<FhtmxUiNoSessionData> {
44    fn default() -> Self {
45        Self::new()
46    }
47}
48
49impl SseSetup<FhtmxUiNoSessionData> {
50    /// Creates a new setup without session data.
51    pub fn new() -> Self {
52        SseSetup {
53            session_data: PhantomData,
54        }
55    }
56}
57
58impl<T> SseSetup<T>
59where
60    T: Send + Sync + 'static,
61{
62    /// Gets a `SseState` instance for you to add it to your app
63    #[must_use]
64    pub fn state_data(&self) -> web::Data<SseState<T>> {
65        web::Data::new(SseState::default())
66    }
67
68    /// Setups the sse route
69    pub fn setup_route(&self, path: &str, cfg: &mut web::ServiceConfig) {
70        cfg.route(path, web::get().to(sse_handler::<T>));
71    }
72}
73
74/// An active SSE session.
75pub struct SseSession<T> {
76    /// Optional per-session data.
77    pub data: Option<T>,
78    /// Channel sender for pushing events.
79    pub sender: mpsc::Sender<Event>,
80}
81
82/// Shared SSE state holding all active sessions.
83pub struct SseState<T> {
84    /// Map of session id to session.
85    pub sessions: Arc<DashMap<Uuid, SseSession<T>>>,
86}
87
88impl<T> Default for SseState<T> {
89    fn default() -> Self {
90        Self {
91            sessions: Arc::new(DashMap::new()),
92        }
93    }
94}
95
96impl<T: Clone> SseState<T> {
97    /// Clones the session data for `id`.
98    pub fn get_session_data(&self, id: Uuid) -> Option<T> {
99        self.sessions.get(&id).and_then(|x| x.data.clone())
100    }
101}
102
103impl<T> SseState<T> {
104    /// Registers a new session.
105    pub fn add_session(
106        &self,
107        id: Uuid,
108        data: Option<T>,
109        sender: mpsc::Sender<Event>,
110    ) -> Option<SseSession<T>> {
111        let session = SseSession { data, sender };
112        self.sessions.insert(id, session)
113    }
114
115    /// Removes and returns a session.
116    pub fn remove_session(&self, id: Uuid) -> Option<(Uuid, SseSession<T>)> {
117        self.sessions.remove(&id)
118    }
119
120    /// Sends a message to session id
121    pub fn send_message(&self, id: Uuid, data: Data) -> Option<()> {
122        let sender = self.sessions.get(&id)?.sender.clone();
123        if sender.try_send(Event::Data(data)).is_err() {
124            // Channel is closed so we remove the session
125            self.remove_session(id);
126        }
127        Some(())
128    }
129
130    /// Broadcast a message to all sessions and returns the number of sent messages
131    pub fn broadcast(&self, data: Data) -> usize {
132        let senders = self
133            .sessions
134            .iter()
135            .map(|o| o.value().sender.clone())
136            .collect::<Vec<_>>();
137        sse_broadcast(senders, data)
138    }
139
140    /// Broadcast a message to all sessions but one id and returns the number of sent messages
141    pub fn broadcast_all_but(&self, id: Uuid, data: Data) -> usize {
142        let senders = self
143            .sessions
144            .iter()
145            .filter_map(|o| {
146                if *o.key() == id {
147                    None
148                } else {
149                    Some(o.value().sender.clone())
150                }
151            })
152            .collect::<Vec<_>>();
153        sse_broadcast(senders, data)
154    }
155}
156
157/// Broadcasts an event to all senders, returning the number of successful sends.
158pub fn sse_broadcast(senders: Vec<mpsc::Sender<Event>>, data: Data) -> usize {
159    senders
160        .into_iter()
161        .filter_map(|o| o.try_send(Event::Data(data.clone())).ok())
162        .count()
163}
164
165/// Route to handle web sockets
166#[tracing::instrument(skip_all)]
167pub async fn sse_handler<T: Send + Sync + 'static>(
168    state: web::Data<SseState<T>>,
169) -> impl Responder {
170    let (tx, rx) = mpsc::channel(8);
171    let id = Uuid::new_v4();
172    state.add_session(id, None, tx.clone());
173
174    let _ = tx
175        .send(Event::Data(Data::new(id.to_string()).event("sse_id")))
176        .await;
177
178    let sessions = state.sessions.clone();
179    tokio::spawn(async move {
180        tx.closed().await;
181        sessions.remove(&id);
182    });
183
184    actix_web_lab::sse::Sse::from_infallible_receiver(rx).with_keep_alive(Duration::from_secs(3))
185}
186
187/// Query parameter for identifying an SSE session.
188#[derive(Clone, Debug, Serialize, Deserialize)]
189pub struct SseHandlerQuery {
190    /// The SSE session id.
191    pub sse_id: Uuid,
192}