1use std::{future::Future, net::IpAddr, sync::Arc, time::Duration};
2
3use anyhow::Context as _;
4use chrono::{DateTime, Utc};
5use tokio::{
6 sync::{RwLock, watch},
7 task::JoinHandle,
8};
9use url::Url;
10
11use crate::board::client::{
12 BoardServerClient, BoardServerClientError, SessionCreatedResponse, SharedSessionFileResponse,
13};
14
15#[derive(Clone, Debug, Eq, PartialEq)]
16pub struct BoardSessionContext {
17 pub session_id: String,
18 pub server_ip: IpAddr,
19 pub http_base_url: Url,
20}
21
22#[derive(Clone, Debug)]
23pub struct SharedSessionFile {
24 pub filename: String,
25 pub relative_path: String,
26 pub http_url: Url,
27 pub size: u64,
28 pub uploaded_at: DateTime<Utc>,
29}
30
31#[derive(Debug)]
32pub struct BoardSession {
33 client: BoardServerClient,
34 info: SessionCreatedResponse,
35 lease_expires_at: Arc<RwLock<DateTime<Utc>>>,
36 heartbeat_stop: watch::Sender<bool>,
37 heartbeat_task: Option<JoinHandle<anyhow::Result<()>>>,
38}
39
40impl BoardSession {
41 pub async fn acquire(client: BoardServerClient, board_type: &str) -> anyhow::Result<Self> {
42 Self::acquire_inner(client, board_type, None).await
43 }
44
45 pub async fn acquire_with_board_id(
46 client: BoardServerClient,
47 board_type: &str,
48 board_id: &str,
49 ) -> anyhow::Result<Self> {
50 let board_id = board_id.trim();
51 if board_id.is_empty() {
52 return Err(anyhow!("board_id must not be empty when provided"));
53 }
54 Self::acquire_inner(client, board_type, Some(board_id)).await
55 }
56
57 async fn acquire_inner(
58 client: BoardServerClient,
59 board_type: &str,
60 board_id: Option<&str>,
61 ) -> anyhow::Result<Self> {
62 let info = match board_id {
63 Some(board_id) => {
64 acquire_session_with(
65 board_type,
66 Some(board_id),
67 || client.create_session_with_board_id(board_type, board_id),
68 |duration| tokio::time::sleep(duration),
69 )
70 .await?
71 }
72 None => {
73 acquire_session_with(
74 board_type,
75 None,
76 || client.create_session(board_type),
77 |duration| tokio::time::sleep(duration),
78 )
79 .await?
80 }
81 };
82
83 let lease_expires_at = Arc::new(RwLock::new(info.lease_expires_at));
84 let (heartbeat_stop, heartbeat_rx) = watch::channel(false);
85 let heartbeat_task = Some(tokio::spawn(run_heartbeat_loop(
86 client.clone(),
87 info.session_id.clone(),
88 lease_expires_at.clone(),
89 heartbeat_rx,
90 )));
91
92 Ok(Self {
93 client,
94 info,
95 lease_expires_at,
96 heartbeat_stop,
97 heartbeat_task,
98 })
99 }
100
101 pub fn info(&self) -> &SessionCreatedResponse {
102 &self.info
103 }
104
105 pub async fn current_lease_expires_at(&self) -> DateTime<Utc> {
106 *self.lease_expires_at.read().await
107 }
108
109 pub async fn context(&self) -> anyhow::Result<BoardSessionContext> {
110 let profile = self
111 .client
112 .get_boot_profile(&self.info.session_id)
113 .await
114 .context("failed to get board session network context")?;
115 let server_ip = profile
116 .server_ip
117 .as_deref()
118 .ok_or_else(|| anyhow!("board server did not report a board-reachable IP address"))?
119 .parse::<IpAddr>()
120 .with_context(|| {
121 format!(
122 "board server returned invalid board-reachable IP address `{}`",
123 profile.server_ip.as_deref().unwrap_or_default()
124 )
125 })?;
126 let http_base_url = profile.http_base_url.ok_or_else(|| {
127 anyhow!("board server did not report a board-reachable shared HTTP URL")
128 })?;
129 Ok(BoardSessionContext {
130 session_id: self.info.session_id.clone(),
131 server_ip,
132 http_base_url,
133 })
134 }
135
136 pub async fn upload_shared_file(
137 &self,
138 relative_path: &str,
139 bytes: Vec<u8>,
140 ) -> anyhow::Result<SharedSessionFile> {
141 let response = self
142 .client
143 .upload_session_file(&self.info.session_id, relative_path, bytes)
144 .await
145 .with_context(|| format!("failed to upload shared session file `{relative_path}`"))?;
146 shared_session_file_from_response(relative_path, response)
147 }
148
149 pub async fn release(mut self) -> anyhow::Result<()> {
150 self.stop_heartbeat();
151 if let Some(task) = self.heartbeat_task.take() {
152 match task.await {
153 Ok(Ok(())) => {}
154 Ok(Err(err)) => log::warn!("board heartbeat stopped with error: {err:#}"),
155 Err(err) if err.is_cancelled() => {}
156 Err(err) => log::warn!("board heartbeat task join error: {err}"),
157 }
158 }
159
160 self.client
161 .delete_session(&self.info.session_id)
162 .await
163 .with_context(|| format!("failed to release session `{}`", self.info.session_id))?;
164 Ok(())
165 }
166
167 fn stop_heartbeat(&self) {
168 let _ = self.heartbeat_stop.send(true);
169 }
170}
171
172fn shared_session_file_from_response(
173 requested_relative_path: &str,
174 response: SharedSessionFileResponse,
175) -> anyhow::Result<SharedSessionFile> {
176 let http_url = response.http_url.ok_or_else(|| {
177 anyhow!(
178 "board server did not return an HTTP URL for shared session file `{requested_relative_path}`"
179 )
180 })?;
181 Ok(SharedSessionFile {
182 filename: response.filename,
183 relative_path: requested_relative_path.to_string(),
184 http_url,
185 size: response.size,
186 uploaded_at: response.uploaded_at,
187 })
188}
189
190impl Drop for BoardSession {
191 fn drop(&mut self) {
192 self.stop_heartbeat();
193 if let Some(task) = &self.heartbeat_task {
194 task.abort();
195 }
196 }
197}
198
199async fn run_heartbeat_loop(
200 client: BoardServerClient,
201 session_id: String,
202 lease_expires_at: Arc<RwLock<DateTime<Utc>>>,
203 mut stop_rx: watch::Receiver<bool>,
204) -> anyhow::Result<()> {
205 loop {
206 tokio::select! {
207 _ = tokio::time::sleep(Duration::from_secs(1)) => {}
208 changed = stop_rx.changed() => {
209 if changed.is_err() || *stop_rx.borrow() {
210 break;
211 }
212 continue;
213 }
214 }
215
216 let heartbeat = client
217 .heartbeat(&session_id)
218 .await
219 .with_context(|| format!("heartbeat failed for session `{session_id}`"))?;
220 log::debug!(
221 "Session {} heartbeat extended lease to {}",
222 heartbeat.session_id,
223 heartbeat.lease_expires_at
224 );
225 *lease_expires_at.write().await = heartbeat.lease_expires_at;
226 }
227
228 Ok(())
229}
230
231async fn acquire_session_with<CreateFn, CreateFut, SleepFn, SleepFut>(
232 board_type: &str,
233 board_id: Option<&str>,
234 mut create: CreateFn,
235 mut sleep: SleepFn,
236) -> Result<SessionCreatedResponse, BoardServerClientError>
237where
238 CreateFn: FnMut() -> CreateFut,
239 CreateFut: Future<Output = Result<SessionCreatedResponse, BoardServerClientError>>,
240 SleepFn: FnMut(Duration) -> SleepFut,
241 SleepFut: Future<Output = ()>,
242{
243 let board_id = board_id.map(str::trim);
244 loop {
245 match create().await {
246 Ok(session) => return Ok(session),
247 Err(err)
248 if board_id
249 .map(|board_id| err.is_no_available_board_id(board_id))
250 .unwrap_or(false) =>
251 {
252 println!(
253 "Board `{}` is not available, retrying in 1s...",
254 board_id.unwrap()
255 );
256 sleep(Duration::from_secs(1)).await;
257 }
258 Err(err) if err.is_no_available_board_for(board_type) => {
259 println!("No available board for type `{board_type}`, retrying in 1s...");
260 sleep(Duration::from_secs(1)).await;
261 }
262 Err(err) if err.is_board_type_not_found_for(board_type) => {
263 return Err(err);
264 }
265 Err(err) => return Err(err),
266 }
267 }
268}
269
270#[cfg(test)]
271mod tests {
272 use std::{
273 sync::{Arc, Mutex},
274 time::Duration,
275 };
276
277 use chrono::Utc;
278 use reqwest::StatusCode;
279
280 use super::{acquire_session_with, shared_session_file_from_response};
281 use crate::board::client::{
282 BoardServerClientError, SessionCreatedResponse, SharedSessionFileResponse,
283 };
284
285 fn created_session(id: &str) -> SessionCreatedResponse {
286 SessionCreatedResponse {
287 session_id: id.to_string(),
288 board_id: "demo-01".to_string(),
289 lease_expires_at: Utc::now(),
290 serial_available: true,
291 boot_mode: "uboot".to_string(),
292 ws_url: Some("/api/v1/sessions/demo/serial/ws".to_string()),
293 }
294 }
295
296 fn no_board_error(board_type: &str) -> BoardServerClientError {
297 BoardServerClientError {
298 status: StatusCode::CONFLICT,
299 code: Some("conflict".to_string()),
300 message: format!("no available board for type `{board_type}`"),
301 }
302 }
303
304 fn no_board_id_error(board_id: &str) -> BoardServerClientError {
305 BoardServerClientError {
306 status: StatusCode::CONFLICT,
307 code: Some("conflict".to_string()),
308 message: format!("board `{board_id}` is not available"),
309 }
310 }
311
312 #[test]
313 fn shared_file_keeps_the_requested_relative_path() {
314 let uploaded_at = Utc::now();
315 let response = SharedSessionFileResponse {
316 filename: "probe.sh".to_string(),
317 relative_path: "ostool/sessions/demo/tools/probe.sh".to_string(),
318 tftp_url: None,
319 http_url: Some(
320 "http://192.168.1.2:2999/share/sessions/demo/tools/probe.sh"
321 .parse()
322 .unwrap(),
323 ),
324 size: 5,
325 uploaded_at,
326 };
327
328 let shared = shared_session_file_from_response("tools/probe.sh", response).unwrap();
329
330 assert_eq!(shared.relative_path, "tools/probe.sh");
331 }
332
333 #[tokio::test]
334 async fn acquire_session_retries_until_a_board_is_available() {
335 let responses = Arc::new(Mutex::new(vec![
336 Ok(created_session("demo-session")),
337 Err(no_board_error("rk3568")),
338 Err(no_board_error("rk3568")),
339 ]));
340 let sleeps = Arc::new(Mutex::new(Vec::new()));
341
342 let session = acquire_session_with(
343 "rk3568",
344 None,
345 {
346 let responses = responses.clone();
347 move || {
348 let responses = responses.clone();
349 async move { responses.lock().unwrap().pop().unwrap() }
350 }
351 },
352 {
353 let sleeps = sleeps.clone();
354 move |duration| {
355 let sleeps = sleeps.clone();
356 async move {
357 sleeps.lock().unwrap().push(duration);
358 }
359 }
360 },
361 )
362 .await
363 .unwrap();
364
365 assert_eq!(session.session_id, "demo-session");
366 assert_eq!(sleeps.lock().unwrap().len(), 2);
367 assert!(
368 sleeps
369 .lock()
370 .unwrap()
371 .iter()
372 .all(|duration| *duration == Duration::from_secs(1))
373 );
374 }
375
376 #[tokio::test]
377 async fn acquire_session_retries_until_requested_board_id_is_available() {
378 let responses = Arc::new(Mutex::new(vec![
379 Ok(created_session("demo-session")),
380 Err(no_board_id_error("demo-02")),
381 ]));
382 let sleeps = Arc::new(Mutex::new(Vec::new()));
383
384 let session = acquire_session_with(
385 "rk3568",
386 Some(" demo-02 "),
387 {
388 let responses = responses.clone();
389 move || {
390 let responses = responses.clone();
391 async move { responses.lock().unwrap().pop().unwrap() }
392 }
393 },
394 {
395 let sleeps = sleeps.clone();
396 move |duration| {
397 let sleeps = sleeps.clone();
398 async move {
399 sleeps.lock().unwrap().push(duration);
400 }
401 }
402 },
403 )
404 .await
405 .unwrap();
406
407 assert_eq!(session.session_id, "demo-session");
408 assert_eq!(*sleeps.lock().unwrap(), vec![Duration::from_secs(1)]);
409 }
410
411 #[tokio::test]
412 async fn acquire_session_stops_retrying_on_non_conflict_error() {
413 let error = BoardServerClientError {
414 status: StatusCode::SERVICE_UNAVAILABLE,
415 code: Some("request_failed".to_string()),
416 message: "boom".to_string(),
417 };
418
419 let result = acquire_session_with(
420 "rk3568",
421 None,
422 || {
423 let error = error.clone();
424 async move { Err(error) }
425 },
426 |_| async {},
427 )
428 .await;
429
430 assert_eq!(result.unwrap_err().message, "boom");
431 }
432
433 #[tokio::test]
434 async fn acquire_session_stops_retrying_when_board_type_is_missing() {
435 let error = BoardServerClientError {
436 status: StatusCode::NOT_FOUND,
437 code: Some("not_found".to_string()),
438 message: "board type `rk3568` not found".to_string(),
439 };
440 let sleeps = Arc::new(Mutex::new(Vec::new()));
441
442 let result = acquire_session_with(
443 "rk3568",
444 None,
445 || {
446 let error = error.clone();
447 async move { Err(error) }
448 },
449 {
450 let sleeps = sleeps.clone();
451 move |duration| {
452 let sleeps = sleeps.clone();
453 async move {
454 sleeps.lock().unwrap().push(duration);
455 }
456 }
457 },
458 )
459 .await;
460
461 assert_eq!(result.unwrap_err().message, "board type `rk3568` not found");
462 assert!(sleeps.lock().unwrap().is_empty());
463 }
464}