feat: support lazy environment activation in NodeEnvManager (#2841)

Author: idoros

packages/runtime-node/src/micro-rpc.ts

@@ -30,6 +30,11 @@ export function rpcCall<T>(worker: Worker, type: string, timeout = 10000): Promi
     return result;
 }
 
+export function rpcPost(worker: Worker, type: string, value?: unknown) {
+    const outgoingMessage = { type, id: getNextMessageId(), value };
+    worker.postMessage(outgoingMessage);
+}
+
 export function bindRpcListener<T>(type: string, customFetcher: () => Promise<T> | T) {
     const handler = async (message: unknown) => {
         if (isValidRpcMessage(message) && message.type === type) {

packages/runtime-node/src/node-env-manager.ts

@@ -19,6 +19,7 @@ export interface RunningNodeEnvironment {
     id: string;
     dispose(): Promise<void>;
     getMetrics(): Promise<PerformanceMetrics>;
+    activate?(): Promise<void>;
 }
 
 export interface NodeEnvConfig extends Pick<AnyEnvironment, 'env' | 'endpointType'> {
@@ -55,6 +56,7 @@ export class NodeEnvManager implements IDisposable {
         }: ILaunchHttpServerOptions & {
             connectionHandlers?: ConnectionHandlers;
         } = {},
+        lazy = false,
     ) {
         process.env.ENGINE_FLOW_V2_DIST_URL = this.importMeta.url;
         const disposeMetricsListener = bindMetricsListener(() => this.collectMetricsFromAllOpenEnvironments());
@@ -95,6 +97,10 @@ export class NodeEnvManager implements IDisposable {
         }
         await this.runFeatureEnvironments(verbose, runtimeOptions, forwardingCom);
 
+        if (!lazy) {
+            await this.activateEnvs();
+        }
+
         app.get('/health', (_req, res) => {
             res.status(200).end();
         });
@@ -120,6 +126,15 @@ export class NodeEnvManager implements IDisposable {
         return { port };
     }
 
+    async activateEnvs() {
+        const activatedEnvs: Promise<void>[] = [];
+        for (const env of this.openEnvironments.values()) {
+            if (!env.activate) continue;
+            activatedEnvs.push(env.activate());
+        }
+        await Promise.all(activatedEnvs);
+    }
+
     async closeAll() {
         await Promise.all([...this.openEnvironments.values()].map((env) => this.closeEnv(env)));
     }
@@ -186,7 +201,7 @@ export class NodeEnvManager implements IDisposable {
                 workerURL: this.createEnvironmentFileUrl(envName),
                 runtimeOptions: runtimeOptions,
             });
-            await envWithInit.initialize();
+            envWithInit.preLoad();
             runningEnv = envWithInit;
         }
 

packages/runtime-node/src/worker-thread-initializer2.ts

@@ -3,11 +3,13 @@ import { Worker } from '@dazl/isomorphic-worker/worker';
 import { type UniversalWorkerOptions } from '@dazl/isomorphic-worker/types';
 import { createDisposables } from '@dazl/patterns';
 import { getMetricsFromWorker } from './metrics-utils.js';
-import { rpcCall } from './micro-rpc.js';
+import { rpcCall, rpcPost } from './micro-rpc.js';
 import type { RunningNodeEnvironment } from './node-env-manager.js';
 
 export interface WorkerThreadInitializer2 extends RunningNodeEnvironment {
     initialize: () => Promise<void>;
+    preLoad: () => void;
+    activate: () => Promise<void>;
 }
 
 export interface WorkerThreadInitializerOptions2 {
@@ -37,7 +39,8 @@ export function workerThreadInitializer2({
     const instanceId = communication.getEnvironmentInstanceId(env.env, env.endpointType);
     const envIsReady = communication.envReady(instanceId);
     let worker: Worker;
-    const initialize = async () => {
+
+    const preLoad = () => {
         const envRuntimeOptions = new Map(runtimeOptions?.entries());
         envRuntimeOptions.set('environment_id', instanceId);
 
@@ -73,13 +76,21 @@ export function workerThreadInitializer2({
             communication.clearEnvironment(instanceId);
             communication.removeMessageHandler(host);
         });
+    };
 
+    const activate = async () => {
+        rpcPost(worker, 'activate');
         await envIsReady;
     };
 
     return {
         id: instanceId,
-        initialize,
+        initialize: async () => {
+            preLoad();
+            await activate();
+        },
+        preLoad,
+        activate,
         dispose: () => disposables.dispose(),
         getMetrics: () => getMetricsFromWorker(worker),
     };

packages/runtime-node/test-kit/entrypoints/a.node.ts

@@ -38,6 +38,11 @@ export function runEnv({
 
 if (workerData) {
     const unbindMetricsListener = bindMetricsListener();
+    let running: ReturnType<typeof runEnv>;
+    const unbindActivateListener = bindRpcListener('activate', () => {
+        unbindActivateListener();
+        running = runEnv();
+    });
     const unbindTerminationListener = bindRpcListener('terminate', async () => {
         if (verbose) {
             console.log(`[${env.env}]: Termination Requested. Waiting for engine.`);
@@ -55,8 +60,6 @@ if (workerData) {
             return;
         }
     });
-    // used by dispose to wait for engine to be ready
-    const running = runEnv();
 } else {
     console.log('running engine in test mode');
 }

packages/runtime-node/test-kit/entrypoints/b.node.ts

@@ -38,6 +38,11 @@ export function runEnv({
 
 if (workerData) {
     const unbindMetricsListener = bindMetricsListener();
+    let running: ReturnType<typeof runEnv>;
+    const unbindActivateListener = bindRpcListener('activate', () => {
+        unbindActivateListener();
+        running = runEnv();
+    });
     const unbindTerminationListener = bindRpcListener('terminate', async () => {
         if (verbose) {
             console.log(`[${env.env}]: Termination Requested. Waiting for engine.`);
@@ -55,8 +60,6 @@ if (workerData) {
             return;
         }
     });
-    // used by dispose to wait for engine to be ready
-    const running = runEnv();
 } else {
     console.log('running engine in test mode');
 }

packages/runtime-node/test/node-env.manager.unit.ts

@@ -1,5 +1,6 @@
 import { expect } from 'chai';
 import sinon from 'sinon';
+import { timeout } from 'promise-assist';
 import { BaseHost, COM, Communication, WsClientHost } from '@dazl/engine-core';
 import {
     IConnectionHandler,
@@ -155,6 +156,45 @@ describe('NodeEnvManager', () => {
         });
     });
 
+    describe('NodeEnvManager with lazy activation', () => {
+        it('should not respond before activateEnvs and work after it', async () => {
+            const featureEnvironmentsMapping: NodeEnvsFeatureMapping = {
+                featureToEnvironments: {
+                    'test-feature': [aEnv.env, bEnv.env],
+                },
+                availableEnvironments: {
+                    a: {
+                        env: aEnv.env,
+                        endpointType: 'single',
+                        envType: 'node',
+                    },
+                    b: {
+                        env: bEnv.env,
+                        endpointType: 'single',
+                        envType: 'node',
+                    },
+                },
+            };
+
+            const manager = disposeAfterTest(new NodeEnvManager(meta, featureEnvironmentsMapping));
+            const { port } = await manager.autoLaunch(new Map([['feature', 'test-feature']]), {}, true);
+            const communication = disposeAfterTest(getClientCom(port));
+            const api = communication.apiProxy<EchoService>({ id: aEnv.env }, { id: 'test-feature.echoAService' });
+
+            // env should not respond before activation
+            const error = await timeout(api.echo(), 500).catch((e: Error) => e);
+            expect(error, 'uninitialized env')
+                .to.be.instanceOf(Error)
+                .with.property('message')
+                .that.includes('timed out');
+
+            // activate and verify env responds
+            await manager.activateEnvs();
+
+            expect(await api.echo()).to.equal('a');
+        });
+    });
+
     describe('NodeEnvManager with connection handlers', () => {
         it('should call connection, reconnection and disconnection handlers', async () => {
             const connectionHandler = sinon.spy();