[Flight] Add Debug Channel option for stateful connection to the backend in DEV (#33627)

This adds plumbing for opening a stream from the Flight Client to the
Flight Server so it can ask for more data on-demand. In this mode, the
Flight Server keeps the connection open as long as the client is still
alive and there's more objects to load. It retains any depth limited
objects so that they can be asked for later. In this first PR it just
releases the object when it's discovered on the server and doesn't
actually lazy load it yet. That's coming in a follow up.

This strategy is built on the model that each request has its own
channel for this. Instead of some global registry. That ensures that
referential identity is preserved within a Request and the Request can
refer to previously written objects by reference.

The fixture implements a WebSocket per request but it doesn't have to be
done that way. It can be multiplexed through an existing WebSocket for
example. The current protocol is just a Readable(Stream) on the server
and WritableStream on the client. It could even be sent through a HTTP
request body if browsers implemented full duplex (which they don't).

This PR only implements the direction of messages from Client to Server.
However, I also plan on adding Debug Channel in the other direction to
allow debug info (optionally) be sent from Server to Client through this
channel instead of through the main RSC request. So the `debugChannel`
option will be able to take writable or readable or both.

---------

Co-authored-by: Hendrik Liebau <mail@hendrik-liebau.de>

Author: sebmarkbage

fixtures/flight/server/global.js

@@ -104,6 +104,9 @@ async function renderApp(req, res, next) {
   if (req.headers['cache-control']) {
     proxiedHeaders['Cache-Control'] = req.get('cache-control');
   }
+  if (req.get('rsc-request-id')) {
+    proxiedHeaders['rsc-request-id'] = req.get('rsc-request-id');
+  }
 
   const requestsPrerender = req.path === '/prerender';
 

fixtures/flight/server/region.js

@@ -50,7 +50,27 @@ const {readFile} = require('fs').promises;
 
 const React = require('react');
 
-async function renderApp(res, returnValue, formState, noCache) {
+const activeDebugChannels =
+  process.env.NODE_ENV === 'development' ? new Map() : null;
+
+function getDebugChannel(req) {
+  if (process.env.NODE_ENV !== 'development') {
+    return undefined;
+  }
+  const requestId = req.get('rsc-request-id');
+  if (!requestId) {
+    return undefined;
+  }
+  return activeDebugChannels.get(requestId);
+}
+
+async function renderApp(
+  res,
+  returnValue,
+  formState,
+  noCache,
+  promiseForDebugChannel
+) {
   const {renderToPipeableStream} = await import(
     'react-server-dom-webpack/server'
   );
@@ -101,7 +121,9 @@ async function renderApp(res, returnValue, formState, noCache) {
   );
   // For client-invoked server actions we refresh the tree and return a return value.
   const payload = {root, returnValue, formState};
-  const {pipe} = renderToPipeableStream(payload, moduleMap);
+  const {pipe} = renderToPipeableStream(payload, moduleMap, {
+    debugChannel: await promiseForDebugChannel,
+  });
   pipe(res);
 }
 
@@ -166,7 +188,7 @@ app.get('/', async function (req, res) {
   if ('prerender' in req.query) {
     await prerenderApp(res, null, null, noCache);
   } else {
-    await renderApp(res, null, null, noCache);
+    await renderApp(res, null, null, noCache, getDebugChannel(req));
   }
 });
 
@@ -204,7 +226,7 @@ app.post('/', bodyParser.text(), async function (req, res) {
       // We handle the error on the client
     }
     // Refresh the client and return the value
-    renderApp(res, result, null, noCache);
+    renderApp(res, result, null, noCache, getDebugChannel(req));
   } else {
     // This is the progressive enhancement case
     const UndiciRequest = require('undici').Request;
@@ -220,11 +242,11 @@ app.post('/', bodyParser.text(), async function (req, res) {
       // Wait for any mutations
       const result = await action();
       const formState = decodeFormState(result, formData);
-      renderApp(res, null, formState, noCache);
+      renderApp(res, null, formState, noCache, undefined);
     } catch (x) {
       const {setServerState} = await import('../src/ServerState.js');
       setServerState('Error: ' + x.message);
-      renderApp(res, null, null, noCache);
+      renderApp(res, null, null, noCache, undefined);
     }
   }
 });
@@ -324,7 +346,7 @@ if (process.env.NODE_ENV === 'development') {
   });
 }
 
-app.listen(3001, () => {
+const httpServer = app.listen(3001, () => {
   console.log('Regional Flight Server listening on port 3001...');
 });
 
@@ -346,3 +368,27 @@ app.on('error', function (error) {
       throw 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();
+    }
+  });
+}

fixtures/flight/src/App.js

@@ -123,7 +123,6 @@ async function ServerComponent({noCache}) {
 export default async function App({prerender, noCache}) {
   const res = await fetch('http://localhost:3001/todos');
   const todos = await res.json();
-  console.log(res);
 
   const dedupedChild = <ServerComponent noCache={noCache} />;
   const message = getServerState();

fixtures/flight/src/index.js

@@ -42,17 +42,43 @@ function Shell({data}) {
 }
 
 async function hydrateApp() {
-  const {root, returnValue, formState} = await createFromFetch(
-    fetch('/', {
-      headers: {
-        Accept: 'text/x-component',
-      },
-    }),
-    {
-      callServer,
-      findSourceMapURL,
-    }
-  );
+  let response;
+  if (
+    process.env.NODE_ENV === 'development' &&
+    typeof WebSocketStream === 'function'
+  ) {
+    const requestId = crypto.randomUUID();
+    const wss = new WebSocketStream(
+      'ws://localhost:3001/debug-channel?' + requestId
+    );
+    const debugChannel = await wss.opened;
+    response = createFromFetch(
+      fetch('/', {
+        headers: {
+          Accept: 'text/x-component',
+          'rsc-request-id': requestId,
+        },
+      }),
+      {
+        callServer,
+        debugChannel,
+        findSourceMapURL,
+      }
+    );
+  } else {
+    response = createFromFetch(
+      fetch('/', {
+        headers: {
+          Accept: 'text/x-component',
+        },
+      }),
+      {
+        callServer,
+        findSourceMapURL,
+      }
+    );
+  }
+  const {root, returnValue, formState} = await response;
 
   ReactDOM.hydrateRoot(
     document,

packages/react-client/src/ReactFlightClient.js

@@ -328,6 +328,8 @@ export type FindSourceMapURLCallback = (
   environmentName: string,
 ) => null | string;
 
+export type DebugChannelCallback = (message: string) => void;
+
 export type Response = {
   _bundlerConfig: ServerConsumerModuleMap,
   _serverReferenceConfig: null | ServerManifest,
@@ -351,6 +353,7 @@ export type Response = {
   _debugRootStack?: null | Error, // DEV-only
   _debugRootTask?: null | ConsoleTask, // DEV-only
   _debugFindSourceMapURL?: void | FindSourceMapURLCallback, // DEV-only
+  _debugChannel?: void | DebugChannelCallback, // DEV-only
   _replayConsole: boolean, // DEV-only
   _rootEnvironmentName: string, // DEV-only, the requested environment name.
 };
@@ -687,6 +690,15 @@ export function reportGlobalError(response: Response, error: Error): void {
       triggerErrorOnChunk(chunk, error);
     }
   });
+  if (__DEV__) {
+    const debugChannel = response._debugChannel;
+    if (debugChannel !== undefined) {
+      // If we don't have any more ways of reading data, we don't have to send any
+      // more neither. So we close the writable side.
+      debugChannel('');
+      response._debugChannel = undefined;
+    }
+  }
   if (enableProfilerTimer && enableComponentPerformanceTrack) {
     markAllTracksInOrder();
     flushComponentPerformance(
@@ -1667,6 +1679,14 @@ function parseModelString(
       }
       case 'Y': {
         if (__DEV__) {
+          if (value.length > 2) {
+            const debugChannel = response._debugChannel;
+            if (debugChannel) {
+              const ref = value.slice(2);
+              debugChannel('R:' + ref); // Release this reference immediately
+            }
+          }
+
           // In DEV mode we encode omitted objects in logs as a getter that throws
           // so that when you try to access it on the client, you know why that
           // happened.
@@ -1730,9 +1750,10 @@ function ResponseInstance(
   encodeFormAction: void | EncodeFormActionCallback,
   nonce: void | string,
   temporaryReferences: void | TemporaryReferenceSet,
-  findSourceMapURL: void | FindSourceMapURLCallback,
-  replayConsole: boolean,
-  environmentName: void | string,
+  findSourceMapURL: void | FindSourceMapURLCallback, // DEV-only
+  replayConsole: boolean, // DEV-only
+  environmentName: void | string, // DEV-only
+  debugChannel: void | DebugChannelCallback, // DEV-only
 ) {
   const chunks: Map<number, SomeChunk<any>> = new Map();
   this._bundlerConfig = bundlerConfig;
@@ -1787,6 +1808,7 @@ function ResponseInstance(
       );
     }
     this._debugFindSourceMapURL = findSourceMapURL;
+    this._debugChannel = debugChannel;
     this._replayConsole = replayConsole;
     this._rootEnvironmentName = rootEnv;
   }
@@ -1802,9 +1824,10 @@ export function createResponse(
   encodeFormAction: void | EncodeFormActionCallback,
   nonce: void | string,
   temporaryReferences: void | TemporaryReferenceSet,
-  findSourceMapURL: void | FindSourceMapURLCallback,
-  replayConsole: boolean,
-  environmentName: void | string,
+  findSourceMapURL: void | FindSourceMapURLCallback, // DEV-only
+  replayConsole: boolean, // DEV-only
+  environmentName: void | string, // DEV-only
+  debugChannel: void | DebugChannelCallback, // DEV-only
 ): Response {
   // $FlowFixMe[invalid-constructor]: the shapes are exact here but Flow doesn't like constructors
   return new ResponseInstance(
@@ -1818,6 +1841,7 @@ export function createResponse(
     findSourceMapURL,
     replayConsole,
     environmentName,
+    debugChannel,
   );
 }