Merge pull request #912 from rust-lang/retry-one-at-a-time

Wait for a WebSocket connection to close before reconnecting

Author: shepmaster

ui/frontend/reducers/websocket.ts

@@ -1,7 +1,7 @@
 import { PayloadAction, createSlice } from '@reduxjs/toolkit';
 import z from 'zod';
 
-import { createWebsocketResponseSchema } from '../websocketActions';
+import { createWebsocketResponseSchema, makeWebSocketMeta } from '../websocketActions';
 
 export type State = {
   connected: boolean;
@@ -14,6 +14,11 @@ const initialState: State = {
   featureFlagEnabled: false,
 };
 
+const websocketConnectedPayloadSchema = z.object({
+  iAcceptThisIsAnUnsupportedApi: z.boolean(),
+});
+type websocketConnectedPayload = z.infer<typeof websocketConnectedPayloadSchema>;
+
 const websocketErrorPayloadSchema = z.object({
   error: z.string(),
 });
@@ -23,9 +28,18 @@ const slice = createSlice({
   name: 'websocket',
   initialState,
   reducers: {
-    connected: (state) => {
-      state.connected = true;
-      delete state.error;
+    connected: {
+      reducer: (state, _action: PayloadAction<websocketConnectedPayload>) => {
+        state.connected = true;
+        delete state.error;
+      },
+
+      prepare: () => ({
+        payload: {
+          iAcceptThisIsAnUnsupportedApi: true,
+        },
+        meta: makeWebSocketMeta(),
+      }),
     },
 
     disconnected: (state) => {
@@ -49,6 +63,11 @@ export const {
   featureFlagEnabled: websocketFeatureFlagEnabled,
 } = slice.actions;
 
+export const websocketConnectedSchema = createWebsocketResponseSchema(
+  websocketConnected,
+  websocketConnectedPayloadSchema,
+);
+
 export const websocketErrorSchema = createWebsocketResponseSchema(
   websocketError,
   websocketErrorPayloadSchema,

ui/frontend/websocketMiddleware.ts

@@ -4,29 +4,49 @@ import { z } from 'zod';
 import { wsExecuteResponseSchema } from './reducers/output/execute';
 import {
   websocketConnected,
+  websocketConnectedSchema,
   websocketDisconnected,
   websocketError,
   websocketErrorSchema,
 } from './reducers/websocket';
 
 const WSMessageResponse = z.discriminatedUnion('type', [
+  websocketConnectedSchema,
   websocketErrorSchema,
   wsExecuteResponseSchema,
 ]);
 
-const reportWebSocketError = async (error: string) => {
-  try {
-    await fetch('/nowebsocket', {
-      method: 'post',
-      headers: {
-        'Content-Type': 'application/json',
-      },
-      body: JSON.stringify({ error }),
-    });
-  } catch (reportError) {
-    console.log('Unable to report WebSocket error', error, reportError);
-  }
-};
+const reportWebSocketError = (() => {
+  let lastReport: string | undefined;
+  let lastReportTime = 0;
+
+  return async (error: string) => {
+    // Don't worry about reporting the same thing again.
+    if (lastReport === error) {
+      return;
+    }
+    lastReport = error;
+
+    // Don't worry about spamming the server with reports.
+    const now = Date.now();
+    if (now - lastReportTime < 1000) {
+      return;
+    }
+    lastReportTime = now;
+
+    try {
+      await fetch('/nowebsocket', {
+        method: 'post',
+        headers: {
+          'Content-Type': 'application/json',
+        },
+        body: JSON.stringify({ error }),
+      });
+    } catch (reportError) {
+      console.log('Unable to report WebSocket error', error, reportError);
+    }
+  };
+})();
 
 const openWebSocket = (currentLocation: Location) => {
   try {
@@ -76,18 +96,16 @@ export const websocketMiddleware =
         resetTimeout();
 
         socket.addEventListener('open', () => {
-          store.dispatch(websocketConnected());
-
-          wasConnected = true;
+          if (socket) {
+            socket.send(JSON.stringify(websocketConnected()));
+          }
         });
 
         socket.addEventListener('close', (event) => {
           store.dispatch(websocketDisconnected());
 
           // Reconnect if we've previously connected
           if (wasConnected && !event.wasClean) {
-            wasConnected = false;
-            reconnectAttempt = 0;
             reconnect();
           }
         });
@@ -104,6 +122,12 @@ export const websocketMiddleware =
           try {
             const rawMessage = JSON.parse(event.data);
             const message = WSMessageResponse.parse(rawMessage);
+
+            if (websocketConnected.match(message)) {
+              wasConnected = true;
+              reconnectAttempt = 0;
+            }
+
             store.dispatch(message);
             resetTimeout();
           } catch (e) {
@@ -114,15 +138,10 @@ export const websocketMiddleware =
     };
 
     const reconnect = () => {
-      if (socket && socket.readyState == socket.OPEN) {
-        return;
-      }
-
-      connect();
-
       const delay = backoffMs(reconnectAttempt);
       reconnectAttempt += 1;
-      setTimeout(reconnect, delay);
+
+      window.setTimeout(connect, delay);
     };
 
     connect();

ui/src/server_axum/websocket.rs

@@ -15,6 +15,23 @@ use tokio::{sync::mpsc, task::JoinSet};
 
 type Meta = serde_json::Value;
 
+#[derive(serde::Deserialize)]
+#[serde(tag = "type")]
+enum HandshakeMessage {
+    #[serde(rename = "websocket/connected")]
+    Connected {
+        payload: Connected,
+        #[allow(unused)]
+        meta: Meta,
+    },
+}
+
+#[derive(serde::Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct Connected {
+    i_accept_this_is_an_unsupported_api: bool,
+}
+
 #[derive(serde::Deserialize)]
 #[serde(tag = "type")]
 enum WSMessageRequest {
@@ -107,6 +124,10 @@ pub async fn handle(mut socket: WebSocket) {
     LIVE_WS.inc();
     let start = Instant::now();
 
+    if !connect_handshake(&mut socket).await {
+        return;
+    }
+
     let (tx, mut rx) = mpsc::channel(3);
     let mut tasks = JoinSet::new();
 
@@ -172,6 +193,15 @@ pub async fn handle(mut socket: WebSocket) {
     DURATION_WS.observe(elapsed.as_secs_f64());
 }
 
+async fn connect_handshake(socket: &mut WebSocket) -> bool {
+    let Some(Ok(Message::Text(txt))) = socket.recv().await else { return false };
+    let Ok(HandshakeMessage::Connected { payload, .. }) = serde_json::from_str::<HandshakeMessage>(&txt) else { return false };
+    if !payload.i_accept_this_is_an_unsupported_api {
+        return false;
+    }
+    socket.send(Message::Text(txt)).await.is_ok()
+}
+
 fn error_to_response(error: Error) -> MessageResponse {
     let error = error.to_string();
     // TODO: thread through the Meta from the originating request