1use 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#[derive(Clone, Copy)]
26pub struct SseSetup<T> {
27 session_data: PhantomData<T>,
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
32pub struct FhtmxUiNoSessionData;
33
34impl SseSetup<()> {
35 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 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 #[must_use]
64 pub fn state_data(&self) -> web::Data<SseState<T>> {
65 web::Data::new(SseState::default())
66 }
67
68 pub fn setup_route(&self, path: &str, cfg: &mut web::ServiceConfig) {
70 cfg.route(path, web::get().to(sse_handler::<T>));
71 }
72}
73
74pub struct SseSession<T> {
76 pub data: Option<T>,
78 pub sender: mpsc::Sender<Event>,
80}
81
82pub struct SseState<T> {
84 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 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 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 pub fn remove_session(&self, id: Uuid) -> Option<(Uuid, SseSession<T>)> {
117 self.sessions.remove(&id)
118 }
119
120 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 self.remove_session(id);
126 }
127 Some(())
128 }
129
130 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 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
157pub 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#[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#[derive(Clone, Debug, Serialize, Deserialize)]
189pub struct SseHandlerQuery {
190 pub sse_id: Uuid,
192}