Engine: try to make calculating payload sizes more efficient (#1155)

* scaffolding: make ensure-payload-size async and add slots for two algorithms

* first swing at async traversal

This actually increases memory overhead

* new traversal algorithm with queue, rather than recuion

* add benchmarking util

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* tidying

* optimisations

Still doesn't get close

* add streaming algo

* update benchmark

* remove traversal algorithm

* tidy

* tidy up

* tidy docs

* remove ignore

* remove old traverse stats

* typing

* restore await

* add a synchronous publish option

important  for process exit events

* revert debug code

---------

Co-authored-by: Claude <noreply@anthropic.com>

Author: josephjclark

packages/engine-multi/memtest.js

@@ -20,7 +20,8 @@
     "@openfn/lexicon": "workspace:^",
     "@openfn/logger": "workspace:*",
     "@openfn/runtime": "workspace:*",
-    "fast-safe-stringify": "^2.1.1"
+    "fast-safe-stringify": "^2.1.1",
+    "json-stream-stringify": "^3.1.6"
   },
   "devDependencies": {
     "@types/node": "^18.19.127",

packages/engine-multi/package.json

@@ -1,3 +1,5 @@
+import { JsonStreamStringify } from 'json-stream-stringify';
+
 // This specifies which keys of an event payload to potentially redact
 // if they are too big
 const KEYS_TO_VERIFY = ['state', 'final_state', 'log'];
@@ -11,18 +13,22 @@ const replacements: Record<string, any> = {
   },
 };
 
-export const verify = (value: any, limit_mb: number = 10) => {
+export const verify = async (
+  value: any,
+  limit_mb: number = 10,
+  algo: 'stringify' | 'stream' = 'stringify'
+) => {
   if (value && !isNaN(limit_mb)) {
-    let size_mb = 0;
-    try {
-      const str = typeof value === 'string' ? value : JSON.stringify(value);
-      const size_bytes = Buffer.byteLength(str, 'utf8');
-      size_mb = size_bytes / 1024 / 1024;
-    } catch (e) {
-      // do nothing
+    const limitBytes = limit_mb * 1024 * 1024;
+
+    let sizeBytes: number;
+    if (algo === 'stream') {
+      sizeBytes = await calculateSizeStream(value, limitBytes);
+    } else {
+      sizeBytes = calculateSizeStringify(value);
     }
 
-    if (size_mb > limit_mb) {
+    if (sizeBytes > limitBytes) {
       const e = new Error();
       // @ts-ignore
       e.name = 'PAYLOAD_TOO_LARGE';
@@ -32,12 +38,39 @@ export const verify = (value: any, limit_mb: number = 10) => {
   }
 };
 
-export default (payload: any, limit_mb: number = 10) => {
+export const calculateSizeStringify = (value: any): number => {
+  const str = typeof value === 'string' ? value : JSON.stringify(value);
+  const size_bytes = Buffer.byteLength(str, 'utf8');
+  return size_bytes;
+};
+
+export const calculateSizeStream = async (
+  value: any,
+  limit?: number
+): Promise<number> => {
+  let size_bytes = 0;
+
+  const stream = new JsonStreamStringify(value);
+
+  for await (const chunk of stream) {
+    // Each chunk is a string token from the JSON output
+    size_bytes += Buffer.byteLength(chunk, 'utf8');
+
+    if (limit !== undefined && size_bytes > limit) {
+      break;
+    }
+  }
+  stream.destroy();
+
+  return size_bytes;
+};
+
+export default async (payload: any, limit_mb: number = 10) => {
   const newPayload = { ...payload };
 
   for (const key of KEYS_TO_VERIFY) {
     try {
-      verify(payload[key], limit_mb);
+      await verify(payload[key], limit_mb);
     } catch (e) {
       Object.assign(newPayload[key], replacements[key] ?? replacements.default);
       newPayload.redacted = true;

packages/engine-multi/src/util/ensure-payload-size.ts

@@ -47,23 +47,36 @@ type Options = {
   payloadLimits?: PayloadLimits;
 };
 
-export const publish = (
+export const publish = async (
   type: string,
   payload: Omit<Event, 'threadId' | 'type'>
 ) => {
   // Validate the size of every outgoing message
   // Redact any payloads that are too large
-  // Use event-specific limit if available, otherwise use default
   const limit =
     payloadLimits?.[type as keyof PayloadLimits] ?? payloadLimits?.default;
-  const safePayload = ensurePayloadSize(payload, limit);
-  const x = {
+  const safePayload = await ensurePayloadSize(payload, limit);
+
+  parentPort!.postMessage({
     type,
     threadId,
     processId,
     ...safePayload,
-  };
-  parentPort!.postMessage(x);
+    ...payload,
+  });
+};
+
+// publishes without payload validation
+export const publishSync = async (
+  type: string,
+  payload: Omit<Event, 'threadId' | 'type'>
+) => {
+  parentPort!.postMessage({
+    type,
+    threadId,
+    processId,
+    ...payload,
+  });
 };
 
 const run = (task: string, args: any[], options: Options = {}) => {
@@ -90,7 +103,7 @@ const run = (task: string, args: any[], options: Options = {}) => {
 };
 
 process.on('exit', (code) => {
-  publish(ENGINE_REJECT_TASK, {
+  publishSync(ENGINE_REJECT_TASK, {
     error: {
       name: 'ExitError',
       message: `Worker thread exited with code: ${code}`,

packages/engine-multi/src/worker/thread/runtime.ts

@@ -4,7 +4,7 @@
 import { getHeapStatistics } from 'node:v8';
 import { createMockLogger } from '@openfn/logger';
 import { randomUUID } from 'node:crypto';
-import createAPI from '../src/api';
+import createAPI from '../../src/api';
 
 let completedCount = 0;
 const logger = createMockLogger();

packages/engine-multi/test/memtest.ts …/engine-multi/test/benchmarks/memtest.tspackages/engine-multi/test/memtest.ts renamed to packages/engine-multi/test/benchmarks/memtest.ts packages/engine-multi/test/memtest.ts renamed to packages/engine-multi/test/benchmarks/memtest.ts

@@ -0,0 +1,27 @@
+/**
+ * To measure memory usage:
+ * 
+ * /usr/bin/time -v pnpm tsx test/payload.ts
+ */
+
+import { verify } from '../../src/util/ensure-payload-size';
+
+// this is 2.63mb of data
+// (I'm using using docs.json from adaptors)
+// @ts-ignore path may not exist
+import data from '../tmp/docs.json' with { type: 'json' };
+
+// this is 16mb of data
+// import data from './tmp/mtg-StandardAtomic.json' with { type: 'json' };
+
+const algo = process.argv[2] ?? 'stringify'
+
+console.log('Running with algo', algo)
+
+try {
+  // @ts-ignore
+  await verify(data, 0.1, algo as any)
+} catch(e) {
+  console.log(e)
+  // do nothing
+}

packages/engine-multi/test/benchmarks/payload.ts

@@ -0,0 +1,107 @@
+#!/usr/bin/env bun
+// @ts-ignore
+import { $ } from 'bun';
+
+const count = parseInt(process.argv[2] ?? '10');
+
+async function run(algo: string) {
+  // Run the command and capture both stdout and stderr
+  const proc =
+    await $`/usr/bin/time -v pnpm tsx test/payload.ts ${algo}`.quiet();
+
+  // time -v outputs to stderr
+  const output = proc.stderr.toString();
+
+  // Extract user time (in seconds)
+  const userTimeMatch = output.match(/User time \(seconds\): ([\d.]+)/);
+  const userTime = userTimeMatch ? parseFloat(userTimeMatch[1]) : null;
+
+  // Extract maximum resident set size (in KB)
+  const maxRssMatch = output.match(
+    /Maximum resident set size \(kbytes\): (\d+)/
+  );
+  const maxRss = maxRssMatch ? parseInt(maxRssMatch[1]) : null;
+
+  // Create JSON output
+  return {
+    time: userTime,
+    mem: maxRss,
+  };
+}
+
+async function bench(algo: string) {
+  console.log(`\n==RUNNING BENCHMARK FOR ${algo.toUpperCase()}==\n`);
+
+  // one once to warm any caching - the very first run is often slowe
+  await run(algo);
+
+  // Run 10 times and collect results
+  const results = [];
+  for (let i = 0; i < count; i++) {
+    const result = await run(algo);
+    results.push(result);
+    console.log(`Run ${i + 1}: used ${result.mem}kb in ${result.time}ms`);
+  }
+
+  // Calculate statistics
+  const times = results.map((r) => r.time || 0);
+  const mems = results.map((r) => r.mem || 0);
+
+  const avgTime = times.reduce((sum, t) => sum + t, 0) / times.length;
+  const minTime = Math.min(...times);
+  const maxTime = Math.max(...times);
+
+  const avgMem = mems.reduce((sum, m) => sum + m, 0) / mems.length;
+  const minMem = Math.min(...mems);
+  const maxMem = Math.max(...mems);
+
+  const stats = {
+    time: { min: minTime, max: maxTime, avg: avgTime },
+    mem: { min: minMem, max: maxMem, avg: avgMem },
+  };
+
+  console.log(`\n${algo} results:`);
+  console.log(
+    `  Time (ms): min=${minTime.toFixed(2)}, max=${maxTime.toFixed(
+      2
+    )}, avg=${avgTime.toFixed(2)}`
+  );
+  console.log(
+    `  Memory (MB): min=${(minMem / 1024).toFixed(2)}, max=${(
+      maxMem / 1024
+    ).toFixed(2)}, avg=${(avgMem / 1024).toFixed(2)}`
+  );
+  console.log();
+
+  return stats;
+}
+
+// @ts-ignore
+const stringifyStats = await bench('stringify');
+// @ts-ignore
+const streamStats = await bench('stream');
+
+// Calculate percentage differences (traverse vs stringify baseline)
+function calcDiff(baseline: number, comparison: number) {
+  return (((comparison - baseline) / baseline) * 100).toFixed(2);
+}
+
+console.log('\n=== COMPARISON (stream vs stringify baseline) ===\n');
+
+console.log('Time Performance:');
+console.log(
+  `  Min: ${calcDiff(stringifyStats.time.min, streamStats.time.min)}%`
+);
+console.log(
+  `  Max: ${calcDiff(stringifyStats.time.max, streamStats.time.max)}%`
+);
+console.log(
+  `  Avg: ${calcDiff(stringifyStats.time.avg, streamStats.time.avg)}%`
+);
+
+console.log('\nMemory Usage:');
+console.log(`  Min: ${calcDiff(stringifyStats.mem.min, streamStats.mem.min)}%`);
+console.log(`  Max: ${calcDiff(stringifyStats.mem.max, streamStats.mem.max)}%`);
+console.log(`  Avg: ${calcDiff(stringifyStats.mem.avg, streamStats.mem.avg)}%`);
+
+console.log('\n(Negative values mean stream is better)');

packages/engine-multi/test/benchmarks/serialization.ts

@@ -518,7 +518,7 @@ test.serial('redact final state if it exceeds the payload limit', (t) => {
 
     const expression = `
 export default [(state) => {
-  state.data = new Array(1024 * 1024).fill('a')
+  state.data = new Array(1024 * 512).fill('a').join('')
   return state;
 }]`;
 
@@ -528,7 +528,7 @@ export default [(state) => {
       },
     ]);
     const options = {
-      payloadLimitMb: 0.5,
+      payloadLimitMb: 0.1,
     };
 
     api