Refactor incremental render handling and improve error management

- Introduced `IncrementalRenderSink` type to manage streaming updates more effectively.
- Updated `handleIncrementalRenderRequest` to return an optional sink and handle execution context errors gracefully.
- Refactored the `run` function to utilize the new sink for processing updates, enhancing error logging for unexpected chunks.
- Simplified test setup by removing unused sink methods, ensuring tests focus on relevant functionality.

Author: AbanoubGhadban

react_on_rails_pro/packages/node-renderer/src/worker.ts

@@ -25,6 +25,7 @@ import handleGracefulShutdown from './worker/handleGracefulShutdown';
 import {
   handleIncrementalRenderRequest,
   type IncrementalRenderInitialRequest,
+  type IncrementalRenderSink,
 } from './worker/handleIncrementalRenderRequest';
 import { handleIncrementalRenderStream } from './worker/handleIncrementalRenderStream';
 import {
@@ -260,7 +261,7 @@ export default function run(config: Partial<Config>) {
     const { bundleTimestamp } = req.params;
 
     // Stream parser state
-    let renderResult: Awaited<ReturnType<typeof handleIncrementalRenderRequest>> | null = null;
+    let incrementalSink: IncrementalRenderSink | undefined;
 
     try {
       // Handle the incremental render stream
@@ -292,10 +293,12 @@ export default function run(config: Partial<Config>) {
           };
 
           try {
-            renderResult = await handleIncrementalRenderRequest(initial);
+            const { response, sink } = await handleIncrementalRenderRequest(initial);
+            incrementalSink = sink;
+
             return {
-              response: renderResult.response,
-              shouldContinue: true,
+              response,
+              shouldContinue: !!incrementalSink,
             };
           } catch (err) {
             const errorResponse = errorResponseResult(
@@ -313,13 +316,13 @@ export default function run(config: Partial<Config>) {
         },
 
         onUpdateReceived: (obj: unknown) => {
-          // Only process updates if we have a render result
-          if (!renderResult) {
+          if (!incrementalSink) {
+            log.error({ msg: 'Unexpected update chunk received after rendering was aborted', obj });
             return;
           }
 
           try {
-            renderResult.sink.add(obj);
+            incrementalSink.add(obj);
           } catch (err) {
             // Log error but don't stop processing
             log.error({ err, msg: 'Error processing update chunk' });
@@ -331,13 +334,7 @@ export default function run(config: Partial<Config>) {
         },
 
         onRequestEnded: () => {
-          try {
-            if (renderResult) {
-              renderResult.sink.end();
-            }
-          } catch (err) {
-            log.error({ err, msg: 'Error ending render sink' });
-          }
+          // Do nothing
         },
       });
     } catch (err) {

react_on_rails_pro/packages/node-renderer/src/worker/handleIncrementalRenderRequest.ts

@@ -1,15 +1,31 @@
 import type { ResponseResult } from '../shared/utils';
 import { handleRenderRequest } from './handleRenderRequest';
+import log from '../shared/log';
+import { getRequestBundleFilePath } from '../shared/utils';
 
 export type IncrementalRenderSink = {
   /** Called for every subsequent NDJSON object after the first one */
   add: (chunk: unknown) => void;
-  /** Called when the client finishes sending the NDJSON stream */
-  end: () => void;
-  /** Called if the request stream errors or validation fails */
-  abort: (error: unknown) => void;
 };
 
+export type UpdateChunk = {
+  bundleTimestamp: string | number;
+  updateChunk: string;
+};
+
+function assertIsUpdateChunk(value: unknown): asserts value is UpdateChunk {
+  if (
+    typeof value !== 'object' ||
+    value === null ||
+    !('bundleTimestamp' in value) ||
+    !('updateChunk' in value) ||
+    (typeof value.bundleTimestamp !== 'string' && typeof value.bundleTimestamp !== 'number') ||
+    typeof value.updateChunk !== 'string'
+  ) {
+    throw new Error('Invalid incremental render chunk received, missing properties');
+  }
+}
+
 export type IncrementalRenderInitialRequest = {
   renderingRequest: string;
   bundleTimestamp: string | number;
@@ -18,7 +34,7 @@ export type IncrementalRenderInitialRequest = {
 
 export type IncrementalRenderResult = {
   response: ResponseResult;
-  sink: IncrementalRenderSink;
+  sink?: IncrementalRenderSink;
 };
 
 /**
@@ -34,26 +50,34 @@ export async function handleIncrementalRenderRequest(
 
   try {
     // Call handleRenderRequest internally to handle all validation and VM execution
-    const renderResult = await handleRenderRequest({
+    const { response, executionContext } = await handleRenderRequest({
       renderingRequest,
       bundleTimestamp,
       dependencyBundleTimestamps,
       providedNewBundles: undefined,
       assetsToCopy: undefined,
     });
 
-    // Return the result directly with a placeholder sink
+    // If we don't get an execution context, it means there was an early error
+    // (e.g. bundle not found). In this case, the sink will be a no-op.
+    if (!executionContext) {
+      return { response };
+    }
+
+    // Return the result with a sink that uses the execution context
     return {
-      response: renderResult,
+      response,
       sink: {
-        add: () => {
-          /* no-op - will be implemented in next commit */
-        },
-        end: () => {
-          /* no-op - will be implemented in next commit */
-        },
-        abort: () => {
-          /* no-op - will be implemented in next commit */
+        add: (chunk: unknown) => {
+          try {
+            assertIsUpdateChunk(chunk);
+            const bundlePath = getRequestBundleFilePath(chunk.bundleTimestamp);
+            executionContext.runInVM(chunk.updateChunk, bundlePath).catch((err: unknown) => {
+              log.error({ msg: 'Error running incremental render chunk', err, chunk });
+            });
+          } catch (err) {
+            log.error({ msg: 'Invalid incremental render chunk', err, chunk });
+          }
         },
       },
     };
@@ -67,17 +91,6 @@ export async function handleIncrementalRenderRequest(
         headers: { 'Cache-Control': 'no-cache, no-store, max-age=0, must-revalidate' },
         data: errorMessage,
       },
-      sink: {
-        add: () => {
-          /* no-op */
-        },
-        end: () => {
-          /* no-op */
-        },
-        abort: () => {
-          /* no-op */
-        },
-      },
     };
   }
 }

react_on_rails_pro/packages/node-renderer/src/worker/handleRenderRequest.ts

@@ -192,7 +192,7 @@ export async function handleRenderRequest({
   dependencyBundleTimestamps?: string[] | number[];
   providedNewBundles?: ProvidedNewBundle[] | null;
   assetsToCopy?: Asset[] | null;
-}): Promise<ResponseResult> {
+}): Promise<{ response: ResponseResult; executionContext?: ExecutionContext }> {
   try {
     // const bundleFilePathPerTimestamp = getRequestBundleFilePath(bundleTimestamp);
     const allBundleFilePaths = Array.from(
@@ -204,15 +204,20 @@ export async function handleRenderRequest({
 
     if (allBundleFilePaths.length > maxVMPoolSize) {
       return {
-        headers: { 'Cache-Control': 'no-cache, no-store, max-age=0, must-revalidate' },
-        status: 410,
-        data: `Too many bundles uploaded. The maximum allowed is ${maxVMPoolSize}. Please reduce the number of bundles or increase maxVMPoolSize in your configuration.`,
+        response: {
+          headers: { 'Cache-Control': 'no-cache, no-store, max-age=0, must-revalidate' },
+          status: 410,
+          data: `Too many bundles uploaded. The maximum allowed is ${maxVMPoolSize}. Please reduce the number of bundles or increase maxVMPoolSize in your configuration.`,
+        },
       };
     }
 
     try {
       const executionContext = await buildExecutionContext(allBundleFilePaths, /* buildVmsIfNeeded */ false);
-      return await prepareResult(renderingRequest, entryBundleFilePath, executionContext);
+      return {
+        response: await prepareResult(renderingRequest, entryBundleFilePath, executionContext),
+        executionContext,
+      };
     } catch (e) {
       // Ignore VMContextNotFoundError, it means the bundle does not exist.
       // The following code will handle this case.
@@ -225,14 +230,14 @@ export async function handleRenderRequest({
     if (providedNewBundles && providedNewBundles.length > 0) {
       const result = await handleNewBundlesProvided(renderingRequest, providedNewBundles, assetsToCopy);
       if (result) {
-        return result;
+        return { response: result };
       }
     }
 
     // Check if the bundle exists:
     const missingBundleError = await validateBundlesExist(bundleTimestamp, dependencyBundleTimestamps);
     if (missingBundleError) {
-      return missingBundleError;
+      return { response: missingBundleError };
     }
 
     // The bundle exists, but the VM has not yet been created.
@@ -243,7 +248,10 @@ export async function handleRenderRequest({
       workerIdLabel(),
     );
     const executionContext = await buildExecutionContext(allBundleFilePaths, /* buildVmsIfNeeded */ true);
-    return await prepareResult(renderingRequest, entryBundleFilePath, executionContext);
+    return {
+      response: await prepareResult(renderingRequest, entryBundleFilePath, executionContext),
+      executionContext,
+    };
   } catch (error) {
     const msg = formatExceptionMessage(
       renderingRequest,