[Flight] Use more robust web socket implementation in fixture (#34338)

The `WebSocketStream` implementation seems to be a bit unreliable. We've
seen `Cannot close a ERRORED writable stream` errors when expanding the
logged deep object, for example. And when reducing the fixture to a
minimal app, we even get `Connection closed` errors, because the web
socket connection is closed before all debug chunks are sent.

We can improve the reliability of the web socket connection by using a
normal `WebSocket` instance on the client, along with manually creating
a `WritableStream` and a `ReadableStream` for processing the messages.

As an additional benefit, the debug channel now also works in Firefox
and Safari.

On the server, we're simplifying the integration with the Express server
a bit by utilizing the `server` property for `WebSocket.Server`, instead
of the `noServer` property with the manual upgrade handling.

fixtures/flight/server/region.js

@@ -74,13 +74,7 @@ function getDebugChannel(req) {
   return activeDebugChannels.get(requestId);
 }
 
-async function renderApp(
-  res,
-  returnValue,
-  formState,
-  noCache,
-  promiseForDebugChannel
-) {
+async function renderApp(res, returnValue, formState, noCache, debugChannel) {
   const {renderToPipeableStream} = await import(
     'react-server-dom-webpack/server'
   );
@@ -132,7 +126,7 @@ async function renderApp(
   // For client-invoked server actions we refresh the tree and return a return value.
   const payload = {root, returnValue, formState};
   const {pipe} = renderToPipeableStream(payload, moduleMap, {
-    debugChannel: await promiseForDebugChannel,
+    debugChannel,
     filterStackFrame,
   });
   pipe(res);
@@ -385,23 +379,20 @@ app.on('error', function (error) {
 if (process.env.NODE_ENV === 'development') {
   // Open a websocket server for Debug information
   const WebSocket = require('ws');
-  const webSocketServer = new WebSocket.Server({noServer: true});
 
-  httpServer.on('upgrade', (request, socket, head) => {
-    const DEBUG_CHANNEL_PATH = '/debug-channel?';
-    if (request.url.startsWith(DEBUG_CHANNEL_PATH)) {
-      const requestId = request.url.slice(DEBUG_CHANNEL_PATH.length);
-      const promiseForWs = new Promise(resolve => {
-        webSocketServer.handleUpgrade(request, socket, head, ws => {
-          ws.on('close', () => {
-            activeDebugChannels.delete(requestId);
-          });
-          resolve(ws);
-        });
-      });
-      activeDebugChannels.set(requestId, promiseForWs);
-    } else {
-      socket.destroy();
-    }
+  const webSocketServer = new WebSocket.Server({
+    server: httpServer,
+    path: '/debug-channel',
+  });
+
+  webSocketServer.on('connection', (ws, req) => {
+    const url = new URL(req.url, `http://${req.headers.host}`);
+    const requestId = url.searchParams.get('id');
+
+    activeDebugChannels.set(requestId, ws);
+
+    ws.on('close', (code, reason) => {
+      activeDebugChannels.delete(requestId);
+    });
   });
 }

fixtures/flight/src/index.js

@@ -14,18 +14,52 @@ function findSourceMapURL(fileName) {
   );
 }
 
+async function createWebSocketStream(url) {
+  const ws = new WebSocket(url);
+  ws.binaryType = 'arraybuffer';
+
+  await new Promise((resolve, reject) => {
+    ws.addEventListener('open', resolve, {once: true});
+    ws.addEventListener('error', reject, {once: true});
+  });
+
+  const writable = new WritableStream({
+    write(chunk) {
+      ws.send(chunk);
+    },
+    close() {
+      ws.close();
+    },
+    abort(reason) {
+      ws.close(1000, reason && String(reason));
+    },
+  });
+
+  const readable = new ReadableStream({
+    start(controller) {
+      ws.addEventListener('message', event => {
+        controller.enqueue(event.data);
+      });
+      ws.addEventListener('close', () => {
+        controller.close();
+      });
+      ws.addEventListener('error', err => {
+        controller.error(err);
+      });
+    },
+  });
+
+  return {readable, writable};
+}
+
 let updateRoot;
 async function callServer(id, args) {
   let response;
-  if (
-    process.env.NODE_ENV === 'development' &&
-    typeof WebSocketStream === 'function'
-  ) {
+  if (process.env.NODE_ENV === 'development') {
     const requestId = crypto.randomUUID();
-    const wss = new WebSocketStream(
-      'ws://localhost:3001/debug-channel?' + requestId
+    const debugChannel = await createWebSocketStream(
+      `ws://localhost:3001/debug-channel?id=${requestId}`
     );
-    const debugChannel = await wss.opened;
     response = createFromFetch(
       fetch('/', {
         method: 'POST',
@@ -74,15 +108,11 @@ function Shell({data}) {
 
 async function hydrateApp() {
   let response;
-  if (
-    process.env.NODE_ENV === 'development' &&
-    typeof WebSocketStream === 'function'
-  ) {
+  if (process.env.NODE_ENV === 'development') {
     const requestId = crypto.randomUUID();
-    const wss = new WebSocketStream(
-      'ws://localhost:3001/debug-channel?' + requestId
+    const debugChannel = await createWebSocketStream(
+      `ws://localhost:3001/debug-channel?id=${requestId}`
     );
-    const debugChannel = await wss.opened;
     response = createFromFetch(
       fetch('/', {
         headers: {