1use tokio::sync::{mpsc, oneshot, watch};
2use tokio::task::JoinHandle;
3
4pub type RealtimeEventStream = mpsc::Receiver<super::ImEvent>;
5
6pub struct RealtimeSession {
7 events: Option<RealtimeEventStream>,
8 status: watch::Receiver<super::RealtimeStatus>,
9 shutdown: super::ShutdownSignal,
10 exit: Option<oneshot::Receiver<super::RealtimeExit>>,
11 worker: Option<JoinHandle<()>>,
12}
13
14impl RealtimeSession {
15 pub(crate) fn new(
16 events: RealtimeEventStream,
17 status: watch::Receiver<super::RealtimeStatus>,
18 shutdown: super::ShutdownSignal,
19 exit: oneshot::Receiver<super::RealtimeExit>,
20 worker: Option<JoinHandle<()>>,
21 ) -> Self {
22 Self {
23 events: Some(events),
24 status,
25 shutdown,
26 exit: Some(exit),
27 worker,
28 }
29 }
30
31 pub fn subscribe(&mut self) -> crate::ImResult<RealtimeEventStream> {
32 self.events.take().ok_or_else(|| {
33 crate::ImError::invalid_input(
34 Some("session".to_owned()),
35 "realtime event stream is already attached",
36 )
37 })
38 }
39
40 pub fn status(&self) -> super::RealtimeStatus {
41 self.status.borrow().clone()
42 }
43
44 pub fn status_updates(&self) -> watch::Receiver<super::RealtimeStatus> {
45 self.status.clone()
46 }
47
48 pub async fn stop(&self) -> crate::ImResult<()> {
49 self.shutdown.request();
50 Ok(())
51 }
52
53 pub async fn join(&mut self) -> crate::ImResult<super::RealtimeExit> {
54 let Some(exit) = self.exit.take() else {
55 return Err(crate::ImError::invalid_input(
56 Some("session".to_owned()),
57 "realtime session exit has already been joined",
58 ));
59 };
60 let exit = match exit.await {
61 Ok(exit) => exit,
62 Err(_) => {
63 return Err(join_worker_error(self.worker.take()).await);
64 }
65 };
66 if let Some(worker) = self.worker.take() {
67 worker.await.map_err(|err| crate::ImError::Internal {
68 message: if err.is_panic() {
69 "realtime session task panicked".to_owned()
70 } else if err.is_cancelled() {
71 "realtime session task was cancelled".to_owned()
72 } else {
73 format!("realtime session task failed: {err}")
74 },
75 })?;
76 }
77 Ok(exit)
78 }
79}
80
81impl Drop for RealtimeSession {
82 fn drop(&mut self) {
83 self.shutdown.request();
84 }
85}
86
87async fn join_worker_error(worker: Option<JoinHandle<()>>) -> crate::ImError {
88 if let Some(worker) = worker {
89 match worker.await {
90 Ok(()) => crate::ImError::Internal {
91 message: "realtime session task exited without an exit result".to_owned(),
92 },
93 Err(err) if err.is_panic() => crate::ImError::Internal {
94 message: "realtime session task panicked".to_owned(),
95 },
96 Err(err) if err.is_cancelled() => crate::ImError::Internal {
97 message: "realtime session task was cancelled".to_owned(),
98 },
99 Err(err) => crate::ImError::Internal {
100 message: format!("realtime session task failed: {err}"),
101 },
102 }
103 } else {
104 crate::ImError::Internal {
105 message: "realtime session task exited without an exit result".to_owned(),
106 }
107 }
108}
109
110pub(crate) fn initial_realtime_status(
111 options: &super::RealtimeOptions,
112 state: super::RealtimeConnectionState,
113 last_error: Option<String>,
114) -> super::RealtimeStatus {
115 super::RealtimeStatus {
116 connected: state == super::RealtimeConnectionState::Connected,
117 state,
118 subscriptions: options.subscriptions.clone(),
119 last_error,
120 }
121}
122
123#[cfg(test)]
124mod tests {
125 use tokio::sync::{mpsc, oneshot, watch};
126
127 use super::*;
128
129 #[tokio::test]
130 async fn realtime_session_allows_single_event_stream_and_stop_request() {
131 let options = super::super::RealtimeOptions::default();
132 let (sender, receiver) = mpsc::channel(options.event_buffer);
133 let (status_sender, status_receiver) = watch::channel(initial_realtime_status(
134 &options,
135 super::super::RealtimeConnectionState::Disconnected,
136 None,
137 ));
138 let (exit_sender, exit_receiver) = oneshot::channel();
139 let shutdown = super::super::ShutdownSignal::pending();
140 let mut session = RealtimeSession::new(
141 receiver,
142 status_receiver,
143 shutdown.clone(),
144 exit_receiver,
145 None,
146 );
147
148 let mut events = session.subscribe().unwrap();
149 assert!(session.subscribe().is_err());
150 sender
151 .send(super::super::ImEvent::ConnectionStateChanged(
152 super::super::ConnectionStateChanged {
153 state: super::super::RealtimeConnectionState::Connected,
154 reason: None,
155 },
156 ))
157 .await
158 .unwrap();
159 let _ = status_sender.send(super::super::RealtimeStatus {
160 connected: true,
161 state: super::super::RealtimeConnectionState::Connected,
162 subscriptions: options.subscriptions,
163 last_error: None,
164 });
165
166 assert!(matches!(
167 events.recv().await,
168 Some(super::super::ImEvent::ConnectionStateChanged(
169 super::super::ConnectionStateChanged {
170 state: super::super::RealtimeConnectionState::Connected,
171 ..
172 }
173 ))
174 ));
175 assert!(session.status().connected);
176 session.stop().await.unwrap();
177 assert!(shutdown.is_requested());
178 exit_sender
179 .send(super::super::RealtimeExit {
180 reason: super::super::RealtimeExitReason::ShutdownRequested,
181 reconnect_attempts: 0,
182 warnings: Vec::new(),
183 })
184 .unwrap();
185 let exit = session.join().await.unwrap();
186 assert_eq!(
187 exit.reason,
188 super::super::RealtimeExitReason::ShutdownRequested
189 );
190 }
191
192 #[tokio::test]
193 async fn realtime_session_drop_requests_worker_shutdown_without_join() {
194 let options = super::super::RealtimeOptions::default();
195 let (_sender, receiver) = mpsc::channel(options.event_buffer);
196 let (_status_sender, status_receiver) = watch::channel(initial_realtime_status(
197 &options,
198 super::super::RealtimeConnectionState::Disconnected,
199 None,
200 ));
201 let (exit_sender, exit_receiver) = oneshot::channel();
202 drop(exit_sender);
203 let shutdown = super::super::ShutdownSignal::pending();
204 let worker_shutdown = shutdown.clone();
205 let (observed_sender, observed_receiver) = oneshot::channel();
206 let worker = tokio::spawn(async move {
207 while !worker_shutdown.is_requested() {
208 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
209 }
210 let _ = observed_sender.send(());
211 });
212
213 let session = RealtimeSession::new(
214 receiver,
215 status_receiver,
216 shutdown,
217 exit_receiver,
218 Some(worker),
219 );
220 drop(session);
221
222 tokio::time::timeout(std::time::Duration::from_millis(250), observed_receiver)
223 .await
224 .expect("drop should request shutdown before timeout")
225 .expect("worker should observe shutdown");
226 }
227
228 #[tokio::test]
229 async fn realtime_session_join_reports_worker_panic() {
230 let options = super::super::RealtimeOptions::default();
231 let (_sender, receiver) = mpsc::channel(options.event_buffer);
232 let (_status_sender, status_receiver) = watch::channel(initial_realtime_status(
233 &options,
234 super::super::RealtimeConnectionState::Disconnected,
235 None,
236 ));
237 let (exit_sender, exit_receiver) = oneshot::channel();
238 drop(exit_sender);
239 let shutdown = super::super::ShutdownSignal::pending();
240 let worker = tokio::spawn(async {
241 panic!("realtime worker panic test");
242 });
243 let mut session = RealtimeSession::new(
244 receiver,
245 status_receiver,
246 shutdown,
247 exit_receiver,
248 Some(worker),
249 );
250
251 let error = session.join().await.unwrap_err();
252
253 assert!(
254 matches!(error, crate::ImError::Internal { message } if message == "realtime session task panicked")
255 );
256 }
257}