Merge pull request #839 from balena-io/add-cluster-mode-test

Add ability to run tests with cluster mode

test/lib/pine-in-process.ts

@@ -1,20 +1,74 @@
 import { exit } from 'process';
+import cluster from 'node:cluster';
 import type { PineTestOptions } from './pine-init';
 import { init } from './pine-init';
 import { tasks } from '../../src/server-glue/module';
 import { PINE_TEST_SIGNALS } from './common';
+import { type Serializable } from 'child_process';
+
+const createWorker = (
+	readyWorkers: Set<number>,
+	processArgs: PineTestOptions,
+) => {
+	const worker = cluster.fork(process.env);
+	worker.on('message', (msg) => {
+		if ('init' in msg && msg.init === 'ready') {
+			readyWorkers.add(worker.id);
+			if (readyWorkers.size === processArgs.clusterInstances && process.send) {
+				process.send({ init: 'success' });
+			}
+		}
+	});
+};
 
 export async function forkInit() {
+	const processArgs: PineTestOptions = JSON.parse(process.argv[2]);
+
+	if (cluster.isPrimary) {
+		const readyWorkers = new Set<number>();
+		process.on('message', async (message: Serializable) => {
+			console.log('Received message in primary process', message);
+			for (const id of readyWorkers.keys()) {
+				cluster.workers?.[id]?.send(message);
+			}
+		});
+		if (processArgs.clusterInstances && processArgs.clusterInstances > 1) {
+			for (let i = 0; i < processArgs.clusterInstances; i++) {
+				createWorker(readyWorkers, processArgs);
+			}
+			cluster.on('exit', (worker) => {
+				// While pine is initializing on empty db a worker might die
+				// as several instances try at the same time to create tables etc
+				// This is not a problem as the worker just tries to recreate until
+				// everything syncs up
+				console.info(`Worker ${worker.process.pid} died`);
+				createWorker(readyWorkers, processArgs);
+			});
+		}
+	}
+
+	if (!cluster.isPrimary || processArgs.clusterInstances === 1) {
+		await runApp(processArgs);
+	}
+
+	if (
+		cluster.isPrimary &&
+		process.send &&
+		(!processArgs.clusterInstances || processArgs.clusterInstances === 1)
+	) {
+		// If single instance or no clustering, send success directly
+		process.send({ init: 'success' });
+	}
+}
+
+async function runApp(processArgs: PineTestOptions) {
+	console.error('Running app in', processArgs);
 	try {
-		// this file is forked - so here we have to evaluate process argv again.
-		// fork only has argv as string[], we use JSON string as arguments
-		const processArgs: PineTestOptions = JSON.parse(process.argv[2]);
 		const { default: initConfig } = await import(processArgs.configPath);
 		console.info(`listenPort: ${processArgs.listenPort}`);
 		const app = await init(
 			initConfig.default,
 			processArgs.listenPort,
-			processArgs.deleteDb,
 			processArgs.withLoginRoute,
 		);
 
@@ -39,7 +93,7 @@ export async function forkInit() {
 		}
 
 		if (process.send) {
-			process.send({ init: 'success' });
+			process.send({ init: 'ready' });
 		}
 
 		process.on('message', async (message) => {

test/lib/pine-init.ts

@@ -11,12 +11,12 @@ export type PineTestOptions = {
 	withLoginRoute?: boolean;
 	deleteDb: boolean;
 	listenPort: number;
+	clusterInstances?: number;
 };
 
 export async function init(
 	initConfig: pine.ConfigLoader.Config,
 	initPort: number,
-	deleteDb = false,
 	withLoginRoute = false,
 ) {
 	const app = express();
@@ -53,7 +53,6 @@ export async function init(
 	});
 
 	try {
-		await cleanInit(deleteDb);
 		const loader = await pine.init(app, initConfig);
 		if (withLoginRoute) {
 			await pine.mountLoginRouter(loader, app);
@@ -69,27 +68,3 @@ export async function init(
 		exit(1);
 	}
 }
-
-async function cleanInit(deleteDb = false) {
-	if (!deleteDb) {
-		return;
-	}
-
-	try {
-		const initDbOptions = {
-			engine:
-				process.env.DATABASE_URL?.slice(
-					0,
-					process.env.DATABASE_URL?.indexOf(':'),
-				) ?? 'postgres',
-			params: process.env.DATABASE_URL ?? 'localhost',
-		};
-		const initDb = pine.dbModule.connect(initDbOptions);
-		await initDb.executeSql(
-			'DROP SCHEMA "public" CASCADE; CREATE SCHEMA "public";',
-		);
-		console.info(`Postgres database dropped`);
-	} catch (e) {
-		console.error(`Error during dropping postgres database: ${e}`);
-	}
-}

test/lib/test-init.ts

@@ -1,6 +1,7 @@
 import type { ChildProcess } from 'child_process';
 import { fork } from 'child_process';
 import { boolVar } from '@balena/env-parsing';
+import * as pine from '../../src/server-glue/module';
 import type { PineTestOptions } from './pine-init';
 import type { OptionalField } from '../../src/sbvr-api/common-types';
 export const listenPortDefault = 1337;
@@ -18,7 +19,11 @@ export async function testInit(
 			taskHandlersPath: options.taskHandlersPath,
 			routesPath: options.routesPath,
 			withLoginRoute: options.withLoginRoute,
+			clusterInstances: options.clusterInstances ?? 1,
 		};
+		if (processArgs.deleteDb) {
+			await cleanDb();
+		}
 		const testServer = fork(
 			__dirname + '/pine-in-process.ts',
 			[JSON.stringify(processArgs)],
@@ -35,7 +40,7 @@ export async function testInit(
 		await new Promise((resolve, reject) => {
 			testServer.on('message', (msg: { init: string }) => {
 				console.info(`init pine in separate process`);
-				if ('init' in msg) {
+				if ('init' in msg && msg.init === 'success') {
 					resolve(msg.init);
 				}
 			});
@@ -56,3 +61,23 @@ export async function testInit(
 export function testDeInit(testServer: ChildProcess) {
 	testServer?.kill();
 }
+
+async function cleanDb() {
+	try {
+		const initDbOptions = {
+			engine:
+				process.env.DATABASE_URL?.slice(
+					0,
+					process.env.DATABASE_URL?.indexOf(':'),
+				) ?? 'postgres',
+			params: process.env.DATABASE_URL ?? 'localhost',
+		};
+		const initDb = pine.dbModule.connect(initDbOptions);
+		await initDb.executeSql(
+			'DROP SCHEMA "public" CASCADE; CREATE SCHEMA "public";',
+		);
+		console.info(`Postgres database dropped`);
+	} catch (e) {
+		console.error(`Error during dropping postgres database: ${e}`);
+	}
+}