Refactor VM handling and introduce ExecutionContext

- Replaced the `runInVM` function with a new `ExecutionContext` class to manage VM contexts more effectively.
- Updated the `handleRenderRequest` function to utilize the new `ExecutionContext`, improving the handling of rendering requests.
- Enhanced error management by introducing `VMContextNotFoundError` for better clarity when VM contexts are missing.
- Refactored tests to align with the new execution context structure, ensuring consistent behavior across rendering scenarios.

Author: AbanoubGhadban

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

@@ -27,7 +27,7 @@ import {
 } from '../shared/utils';
 import { getConfig } from '../shared/configBuilder';
 import * as errorReporter from '../shared/errorReporter';
-import { buildVM, hasVMContextForBundle, runInVM } from './vm';
+import { buildExecutionContext, ExecutionContext, VMContextNotFoundError } from './vm';
 
 export type ProvidedNewBundle = {
   timestamp: string | number;
@@ -37,9 +37,10 @@ export type ProvidedNewBundle = {
 async function prepareResult(
   renderingRequest: string,
   bundleFilePathPerTimestamp: string,
+  executionContext: ExecutionContext,
 ): Promise<ResponseResult> {
   try {
-    const result = await runInVM(renderingRequest, bundleFilePathPerTimestamp, cluster);
+    const result = await executionContext.runInVM(renderingRequest, bundleFilePathPerTimestamp, cluster);
 
     let exceptionMessage = null;
     if (!result) {
@@ -209,9 +210,15 @@ export async function handleRenderRequest({
       };
     }
 
-    // If the current VM has the correct bundle and is ready
-    if (allBundleFilePaths.every((bundleFilePath) => hasVMContextForBundle(bundleFilePath))) {
-      return await prepareResult(renderingRequest, entryBundleFilePath);
+    try {
+      const executionContext = await buildExecutionContext(allBundleFilePaths, /* buildVmsIfNeeded */ false);
+      return await prepareResult(renderingRequest, entryBundleFilePath, executionContext);
+    } catch (e) {
+      // Ignore VMContextNotFoundError, it means the bundle does not exist.
+      // The following code will handle this case.
+      if (!(e instanceof VMContextNotFoundError)) {
+        throw e;
+      }
     }
 
     // If gem has posted updated bundle:
@@ -230,10 +237,13 @@ export async function handleRenderRequest({
 
     // The bundle exists, but the VM has not yet been created.
     // Another worker must have written it or it was saved during deployment.
-    log.info('Bundle %s exists. Building VM for worker %s.', entryBundleFilePath, workerIdLabel());
-    await Promise.all(allBundleFilePaths.map((bundleFilePath) => buildVM(bundleFilePath)));
-
-    return await prepareResult(renderingRequest, entryBundleFilePath);
+    log.info(
+      'Bundle %s exists. Building ExecutionContext for worker %s.',
+      entryBundleFilePath,
+      workerIdLabel(),
+    );
+    const executionContext = await buildExecutionContext(allBundleFilePaths, /* buildVmsIfNeeded */ true);
+    return await prepareResult(renderingRequest, entryBundleFilePath, executionContext);
   } catch (error) {
     const msg = formatExceptionMessage(
       renderingRequest,

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

@@ -29,7 +29,7 @@ import * as errorReporter from '../shared/errorReporter';
 const readFileAsync = promisify(fs.readFile);
 const writeFileAsync = promisify(fs.writeFile);
 
-interface VMContext {
+export interface VMContext {
   context: Context;
   sharedConsoleHistory: SharedConsoleHistory;
   lastUsed: number; // Track when this VM was last used
@@ -101,84 +101,14 @@ function manageVMPoolSize() {
   }
 }
 
-/**
- *
- * @param renderingRequest JS Code to execute for SSR
- * @param filePath
- * @param vmCluster
- */
-export async function runInVM(
-  renderingRequest: string,
-  filePath: string,
-  vmCluster?: typeof cluster,
-): Promise<RenderResult> {
-  const { serverBundleCachePath } = getConfig();
-
-  try {
-    // Wait for VM creation if it's in progress
-    if (vmCreationPromises.has(filePath)) {
-      await vmCreationPromises.get(filePath);
-    }
-
-    // Get the correct VM context based on the provided bundle path
-    const vmContext = getVMContext(filePath);
-
-    if (!vmContext) {
-      throw new Error(`No VM context found for bundle ${filePath}`);
-    }
-
-    // Update last used timestamp
-    vmContext.lastUsed = Date.now();
-
-    const { context, sharedConsoleHistory } = vmContext;
-
-    if (log.level === 'debug') {
-      // worker is nullable in the primary process
-      const workerId = vmCluster?.worker?.id;
-      log.debug(`worker ${workerId ? `${workerId} ` : ''}received render request for bundle ${filePath} with code
-${smartTrim(renderingRequest)}`);
-      const debugOutputPathCode = path.join(serverBundleCachePath, 'code.js');
-      log.debug(`Full code executed written to: ${debugOutputPathCode}`);
-      await writeFileAsync(debugOutputPathCode, renderingRequest);
-    }
-
-    let result = sharedConsoleHistory.trackConsoleHistoryInRenderRequest(() => {
-      context.renderingRequest = renderingRequest;
-      try {
-        return vm.runInContext(renderingRequest, context) as RenderCodeResult;
-      } finally {
-        context.renderingRequest = undefined;
-      }
-    });
-
-    if (isReadableStream(result)) {
-      const newStreamAfterHandlingError = handleStreamError(result, (error) => {
-        const msg = formatExceptionMessage(renderingRequest, error, 'Error in a rendering stream');
-        errorReporter.message(msg);
-      });
-      return newStreamAfterHandlingError;
-    }
-    if (typeof result !== 'string') {
-      const objectResult = await result;
-      result = JSON.stringify(objectResult);
-    }
-    if (log.level === 'debug') {
-      log.debug(`result from JS:
-${smartTrim(result)}`);
-      const debugOutputPathResult = path.join(serverBundleCachePath, 'result.json');
-      log.debug(`Wrote result to file: ${debugOutputPathResult}`);
-      await writeFileAsync(debugOutputPathResult, result);
-    }
-
-    return result;
-  } catch (exception) {
-    const exceptionMessage = formatExceptionMessage(renderingRequest, exception);
-    log.debug('Caught exception in rendering request', exceptionMessage);
-    return Promise.resolve({ exceptionMessage });
+export class VMContextNotFoundError extends Error {
+  constructor(bundleFilePath: string) {
+    super(`VMContext not found for bundle: ${bundleFilePath}`);
+    this.name = 'VMContextNotFoundError';
   }
 }
 
-export async function buildVM(filePath: string): Promise<VMContext> {
+async function buildVM(filePath: string): Promise<VMContext> {
   // Return existing promise if VM is already being created
   if (vmCreationPromises.has(filePath)) {
     return vmCreationPromises.get(filePath) as Promise<VMContext>;
@@ -200,12 +130,7 @@ export async function buildVM(filePath: string): Promise<VMContext> {
         additionalContext !== null && additionalContext.constructor === Object;
       const sharedConsoleHistory = new SharedConsoleHistory();
 
-      const runOnOtherBundle = async (bundleTimestamp: string | number, renderingRequest: string) => {
-        const bundlePath = getRequestBundleFilePath(bundleTimestamp);
-        return runInVM(renderingRequest, bundlePath, cluster);
-      };
-
-      const contextObject = { sharedConsoleHistory, runOnOtherBundle };
+      const contextObject = { sharedConsoleHistory };
 
       if (supportModules) {
         // IMPORTANT: When adding anything to this object, update:
@@ -349,6 +274,120 @@ export async function buildVM(filePath: string): Promise<VMContext> {
   return vmCreationPromise;
 }
 
+async function getOrBuildVMContext(bundleFilePath: string, buildVmsIfNeeded: boolean): Promise<VMContext> {
+  const vmContext = getVMContext(bundleFilePath);
+  if (vmContext) {
+    return vmContext;
+  }
+
+  const vmCreationPromise = vmCreationPromises.get(bundleFilePath);
+  if (vmCreationPromise) {
+    return vmCreationPromise;
+  }
+
+  if (buildVmsIfNeeded) {
+    return buildVM(bundleFilePath);
+  }
+
+  throw new VMContextNotFoundError(bundleFilePath);
+}
+
+export type ExecutionContext = {
+  runInVM: (
+    renderingRequest: string,
+    bundleFilePath: string,
+    vmCluster?: typeof cluster,
+  ) => Promise<RenderResult>;
+  getVMContext: (bundleFilePath: string) => VMContext | undefined;
+};
+
+export async function buildExecutionContext(
+  bundlePaths: string[],
+  buildVmsIfNeeded: boolean,
+): Promise<ExecutionContext> {
+  const mapBundleFilePathToVMContext = new Map<string, VMContext>();
+  await Promise.all(
+    bundlePaths.map(async (bundleFilePath) => {
+      const vmContext = await getOrBuildVMContext(bundleFilePath, buildVmsIfNeeded);
+      vmContext.lastUsed = Date.now();
+      mapBundleFilePathToVMContext.set(bundleFilePath, vmContext);
+    }),
+  );
+  const sharedExecutionContext = new Map();
+
+  const runInVM = async (renderingRequest: string, bundleFilePath: string, vmCluster?: typeof cluster) => {
+    try {
+      const { serverBundleCachePath } = getConfig();
+      const vmContext = mapBundleFilePathToVMContext.get(bundleFilePath);
+      if (!vmContext) {
+        throw new VMContextNotFoundError(bundleFilePath);
+      }
+
+      // Update last used timestamp
+      vmContext.lastUsed = Date.now();
+
+      const { context, sharedConsoleHistory } = vmContext;
+
+      if (log.level === 'debug') {
+        // worker is nullable in the primary process
+        const workerId = vmCluster?.worker?.id;
+        log.debug(`worker ${workerId ? `${workerId} ` : ''}received render request for bundle ${bundleFilePath} with code
+  ${smartTrim(renderingRequest)}`);
+        const debugOutputPathCode = path.join(serverBundleCachePath, 'code.js');
+        log.debug(`Full code executed written to: ${debugOutputPathCode}`);
+        await writeFileAsync(debugOutputPathCode, renderingRequest);
+      }
+
+      let result = sharedConsoleHistory.trackConsoleHistoryInRenderRequest(() => {
+        context.renderingRequest = renderingRequest;
+        context.sharedExecutionContext = sharedExecutionContext;
+        context.runOnOtherBundle = (bundleTimestamp: string | number, newRenderingRequest: string) => {
+          const otherBundleFilePath = getRequestBundleFilePath(bundleTimestamp);
+          return runInVM(otherBundleFilePath, newRenderingRequest, vmCluster);
+        };
+
+        try {
+          return vm.runInContext(renderingRequest, context) as RenderCodeResult;
+        } finally {
+          context.renderingRequest = undefined;
+          context.sharedExecutionContext = undefined;
+          context.runOnOtherBundle = undefined;
+        }
+      });
+
+      if (isReadableStream(result)) {
+        const newStreamAfterHandlingError = handleStreamError(result, (error) => {
+          const msg = formatExceptionMessage(renderingRequest, error, 'Error in a rendering stream');
+          errorReporter.message(msg);
+        });
+        return newStreamAfterHandlingError;
+      }
+      if (typeof result !== 'string') {
+        const objectResult = await result;
+        result = JSON.stringify(objectResult);
+      }
+      if (log.level === 'debug') {
+        log.debug(`result from JS:
+  ${smartTrim(result)}`);
+        const debugOutputPathResult = path.join(serverBundleCachePath, 'result.json');
+        log.debug(`Wrote result to file: ${debugOutputPathResult}`);
+        await writeFileAsync(debugOutputPathResult, result);
+      }
+
+      return result;
+    } catch (exception) {
+      const exceptionMessage = formatExceptionMessage(renderingRequest, exception);
+      log.debug('Caught exception in rendering request', exceptionMessage);
+      return Promise.resolve({ exceptionMessage });
+    }
+  };
+
+  return {
+    getVMContext: (bundleFilePath: string) => mapBundleFilePathToVMContext.get(bundleFilePath),
+    runInVM,
+  };
+}
+
 export function resetVM() {
   // Clear all VM contexts
   vmContexts.clear();

react_on_rails_pro/packages/node-renderer/tests/helper.ts

@@ -4,7 +4,7 @@ import path from 'path';
 import fsPromises from 'fs/promises';
 import fs from 'fs';
 import fsExtra from 'fs-extra';
-import { buildVM, resetVM } from '../src/worker/vm';
+import { buildExecutionContext, resetVM } from '../src/worker/vm';
 import { buildConfig } from '../src/shared/configBuilder';
 
 export const mkdirAsync = fsPromises.mkdir;
@@ -59,12 +59,12 @@ export function vmSecondaryBundlePath(testName: string) {
 
 export async function createVmBundle(testName: string) {
   await safeCopyFileAsync(getFixtureBundle(), vmBundlePath(testName));
-  await buildVM(vmBundlePath(testName));
+  await buildExecutionContext([vmBundlePath(testName)], /* buildVmsIfNeeded */ true);
 }
 
 export async function createSecondaryVmBundle(testName: string) {
   await safeCopyFileAsync(getFixtureSecondaryBundle(), vmSecondaryBundlePath(testName));
-  await buildVM(vmSecondaryBundlePath(testName));
+  await buildExecutionContext([vmSecondaryBundlePath(testName)], /* buildVmsIfNeeded */ true);
 }
 
 export function lockfilePath(testName: string) {

react_on_rails_pro/packages/node-renderer/tests/serverRenderRSCReactComponent.test.js

@@ -1,7 +1,7 @@
 import path from 'path';
 import fs from 'fs';
 import { Readable } from 'stream';
-import { buildVM, getVMContext, resetVM } from '../src/worker/vm';
+import { buildExecutionContext, resetVM } from '../src/worker/vm';
 import { getConfig } from '../src/shared/configBuilder';
 
 const SimpleWorkingComponent = () => 'hello';
@@ -62,8 +62,8 @@ describe('serverRenderRSCReactComponent', () => {
   // The serverRenderRSCReactComponent function should only be called when the bundle is compiled with the `react-server` condition.
   // Therefore, we cannot call it directly in the test files. Instead, we run the RSC bundle through the VM and call the method from there.
   const getReactOnRailsRSCObject = async () => {
-    // Use the copied rsc-bundle.js file from temp directory
-    const vmContext = await buildVM(tempRscBundlePath);
+    const executionContext = await buildExecutionContext([tempRscBundlePath], /* buildVmsIfNeeded */ true);
+    const vmContext = executionContext.getVMContext(tempRscBundlePath);
     const { ReactOnRails, React } = vmContext.context;
 
     function SuspensedComponentWithAsyncError() {