Skip to content

Commit eb199de

Browse files
committed
Replace body of tokio::select with an enum
1 parent 1aad586 commit eb199de

File tree

1 file changed

+60
-29
lines changed

1 file changed

+60
-29
lines changed

ui/src/server_axum/websocket.rs

Lines changed: 60 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -371,25 +371,55 @@ async fn handle_core(
371371
let mut active_execution_gc_interval = time::interval(Duration::from_secs(30));
372372

373373
loop {
374-
tokio::select! {
375-
request = socket.recv() => {
374+
use Event::*;
375+
376+
enum Event {
377+
Request(Option<Result<Message, axum::Error>>),
378+
Response(Option<Result<MessageResponse, TaggedError>>),
379+
Task(Result<Result<(), TaggedError>, tokio::task::JoinError>),
380+
GarbageCollection,
381+
IdleTimeout,
382+
IdleRequest,
383+
SessionTimeout,
384+
}
385+
386+
let event = tokio::select! {
387+
request = socket.recv() => Request(request),
388+
389+
resp = rx.recv() => Response(resp),
390+
391+
// We don't care if there are no running tasks
392+
Some(task) = manager.join_next() => Task(task),
393+
394+
_ = active_execution_gc_interval.tick() => GarbageCollection,
395+
396+
_ = &mut idle_timeout, if manager.is_empty() => IdleTimeout,
397+
398+
_ = factory.container_requested(), if manager.is_empty() => IdleRequest,
399+
400+
_ = &mut session_timeout => SessionTimeout,
401+
};
402+
403+
match event {
404+
Request(request) => {
376405
metrics::WS_INCOMING.inc();
377406

378407
match request {
379-
None => {
380-
// browser disconnected
381-
break;
382-
}
383-
Some(Ok(Message::Text(txt))) => handle_msg(&txt, &tx, &mut manager, &mut active_executions, &db).await,
384-
Some(Ok(_)) => {
385-
// unknown message type
386-
continue;
408+
// browser disconnected
409+
None => break,
410+
411+
Some(Ok(Message::Text(txt))) => {
412+
handle_msg(&txt, &tx, &mut manager, &mut active_executions, &db).await
387413
}
414+
415+
// unknown message type
416+
Some(Ok(_)) => continue,
417+
388418
Some(Err(e)) => super::record_websocket_error(e.to_string()),
389419
}
390-
},
420+
}
391421

392-
resp = rx.recv() => {
422+
Response(resp) => {
393423
let resp = resp.expect("The rx should never close as we have a tx");
394424

395425
let success = resp.is_ok();
@@ -403,10 +433,9 @@ async fn handle_core(
403433

404434
let success = if success { "true" } else { "false" };
405435
metrics::WS_OUTGOING.with_label_values(&[success]).inc();
406-
},
436+
}
407437

408-
// We don't care if there are no running tasks
409-
Some(task) = manager.join_next() => {
438+
Task(task) => {
410439
// The last task has completed which means we are a
411440
// candidate for idling in a little while.
412441
if manager.is_empty() {
@@ -415,17 +444,21 @@ async fn handle_core(
415444

416445
let (error, meta) = match task {
417446
Ok(Ok(())) => continue,
447+
418448
Ok(Err(error)) => error,
449+
419450
Err(error) => {
420451
// The task was cancelled; no need to report
421-
let Ok(panic) = error.try_into_panic() else { continue };
452+
let Ok(panic) = error.try_into_panic() else {
453+
continue;
454+
};
422455

423456
let text = match panic.downcast::<String>() {
424457
Ok(text) => *text,
425458
Err(panic) => match panic.downcast::<&str>() {
426459
Ok(text) => text.to_string(),
427460
_ => "An unknown panic occurred".into(),
428-
}
461+
},
429462
};
430463
(WebSocketTaskPanicSnafu { text }.build(), None)
431464
}
@@ -435,32 +468,30 @@ async fn handle_core(
435468
// We can't send a response
436469
break;
437470
}
438-
},
471+
}
439472

440-
_ = active_execution_gc_interval.tick() => {
473+
GarbageCollection => {
441474
active_executions = mem::take(&mut active_executions)
442475
.into_iter()
443476
.filter(|(_id, (_, tx))| tx.as_ref().is_some_and(|tx| !tx.is_closed()))
444477
.collect();
445-
},
478+
}
446479

447-
_ = &mut idle_timeout, if manager.is_empty() => {
480+
IdleTimeout => {
448481
if handle_idle(&mut manager, &tx).await.is_break() {
449-
break
482+
break;
450483
}
451-
},
484+
}
452485

453-
_ = factory.container_requested(), if manager.is_empty() => {
486+
IdleRequest => {
454487
info!("Container requested to idle");
455488

456489
if handle_idle(&mut manager, &tx).await.is_break() {
457-
break
490+
break;
458491
}
459-
},
460-
461-
_ = &mut session_timeout => {
462-
break;
463492
}
493+
494+
SessionTimeout => break,
464495
}
465496
}
466497

0 commit comments

Comments
 (0)