Replace body of tokio::select with an enum

Author: shepmaster

ui/src/server_axum/websocket.rs

@@ -371,25 +371,55 @@ async fn handle_core(
     let mut active_execution_gc_interval = time::interval(Duration::from_secs(30));
 
     loop {
-        tokio::select! {
-            request = socket.recv() => {
+        use Event::*;
+
+        enum Event {
+            Request(Option<Result<Message, axum::Error>>),
+            Response(Option<Result<MessageResponse, TaggedError>>),
+            Task(Result<Result<(), TaggedError>, tokio::task::JoinError>),
+            GarbageCollection,
+            IdleTimeout,
+            IdleRequest,
+            SessionTimeout,
+        }
+
+        let event = tokio::select! {
+            request = socket.recv() => Request(request),
+
+            resp = rx.recv() => Response(resp),
+
+            // We don't care if there are no running tasks
+            Some(task) = manager.join_next() => Task(task),
+
+            _ = active_execution_gc_interval.tick() => GarbageCollection,
+
+            _ = &mut idle_timeout, if manager.is_empty() => IdleTimeout,
+
+            _ = factory.container_requested(), if manager.is_empty() => IdleRequest,
+
+            _ = &mut session_timeout => SessionTimeout,
+        };
+
+        match event {
+            Request(request) => {
                 metrics::WS_INCOMING.inc();
 
                 match request {
-                    None => {
-                        // browser disconnected
-                        break;
-                    }
-                    Some(Ok(Message::Text(txt))) => handle_msg(&txt, &tx, &mut manager, &mut active_executions, &db).await,
-                    Some(Ok(_)) => {
-                        // unknown message type
-                        continue;
+                    // browser disconnected
+                    None => break,
+
+                    Some(Ok(Message::Text(txt))) => {
+                        handle_msg(&txt, &tx, &mut manager, &mut active_executions, &db).await
                     }
+
+                    // unknown message type
+                    Some(Ok(_)) => continue,
+
                     Some(Err(e)) => super::record_websocket_error(e.to_string()),
                 }
-            },
+            }
 
-            resp = rx.recv() => {
+            Response(resp) => {
                 let resp = resp.expect("The rx should never close as we have a tx");
 
                 let success = resp.is_ok();
@@ -403,10 +433,9 @@ async fn handle_core(
 
                 let success = if success { "true" } else { "false" };
                 metrics::WS_OUTGOING.with_label_values(&[success]).inc();
-            },
+            }
 
-            // We don't care if there are no running tasks
-            Some(task) = manager.join_next() => {
+            Task(task) => {
                 // The last task has completed which means we are a
                 // candidate for idling in a little while.
                 if manager.is_empty() {
@@ -415,17 +444,21 @@ async fn handle_core(
 
                 let (error, meta) = match task {
                     Ok(Ok(())) => continue,
+
                     Ok(Err(error)) => error,
+
                     Err(error) => {
                         // The task was cancelled; no need to report
-                        let Ok(panic) = error.try_into_panic() else { continue };
+                        let Ok(panic) = error.try_into_panic() else {
+                            continue;
+                        };
 
                         let text = match panic.downcast::<String>() {
                             Ok(text) => *text,
                             Err(panic) => match panic.downcast::<&str>() {
                                 Ok(text) => text.to_string(),
                                 _ => "An unknown panic occurred".into(),
-                            }
+                            },
                         };
                         (WebSocketTaskPanicSnafu { text }.build(), None)
                     }
@@ -435,32 +468,30 @@ async fn handle_core(
                     // We can't send a response
                     break;
                 }
-            },
+            }
 
-            _ = active_execution_gc_interval.tick() => {
+            GarbageCollection => {
                 active_executions = mem::take(&mut active_executions)
                     .into_iter()
                     .filter(|(_id, (_, tx))| tx.as_ref().is_some_and(|tx| !tx.is_closed()))
                     .collect();
-            },
+            }
 
-            _ = &mut idle_timeout, if manager.is_empty() => {
+            IdleTimeout => {
                 if handle_idle(&mut manager, &tx).await.is_break() {
-                    break
+                    break;
                 }
-            },
+            }
 
-            _ = factory.container_requested(), if manager.is_empty() => {
+            IdleRequest => {
                 info!("Container requested to idle");
 
                 if handle_idle(&mut manager, &tx).await.is_break() {
-                    break
+                    break;
                 }
-            },
-
-            _ = &mut session_timeout => {
-                break;
             }
+
+            SessionTimeout => break,
         }
     }