Worker: Fix dataclip serialization

.changeset/tame-hats-drum.md

@@ -0,0 +1,5 @@
+---
+'@openfn/ws-worker': patch
+---
+
+Capture more diagnostic detail when the connection to Lightning drops unexpectedly

.changeset/tidy-plums-obey.md

@@ -0,0 +1,5 @@
+---
+'@openfn/ws-worker': patch
+---
+
+Reduce dataclip bloat when sending step results to Lightning

packages/ws-worker/README.md

@@ -51,6 +51,12 @@ Use `-l mock` to connect to a lightning mock server (on the default port).
 
 For a list of supported worker and engine options, see src/start.ts
 
+### Sending output dataclips without double-encoding
+
+By default, the worker JSON-stringifies each step's output dataclip before sending it to Lightning, and Lightning's own transport then re-encodes the whole envelope — this double-encoding bloats large dataclips on the wire. Pass `--no-stringify-state` or set `WORKER_NO_STRINGIFY_STATE` to send the dataclip as a native JSON value instead, avoiding that bloat.
+
+This is only compatible with Lightning 2.19 or later — do not enable it against older Lightning versions.
+
 ## Enforcing memory limits with cgroups
 
 Each run's memory limit is enforced by default through node's max-old-space-size, which only constrains heap size. Native and buffer allocations bypass this limit. This can cause the worker to consume more memory than it is technically allowed, which can in turn cause the whole worker process to be killed by its container (ie, kubernetes).

packages/ws-worker/src/api/execute.ts

@@ -80,6 +80,16 @@ export function execute(
     sentryScope,
   };
 
+  // Log pheonix channel errors to sentry
+  channel.onError((...args: any) => {
+    sentryScope.addBreadcrumb({
+      category: 'channel',
+      message: 'Channel error',
+      level: 'warning',
+      data: { state: channel.state, args },
+    });
+  });
+
   Sentry.withIsolationScope(sentryScope, async () => {
     Sentry.addBreadcrumb({
       category: 'run',

packages/ws-worker/src/channels/run.ts

@@ -63,7 +63,13 @@ const joinRunChannel = (
     channel.onError((...args: any) => {
       // Error occurred on the channel
       // (the socket will try to reconnect with backoff)
-      logger.debug(`Critical error in channel ${channelName}`, args);
+      // Note we don't report to sentry here - the socket error handler does that
+      logger.error(
+        `Critical error in channel ${channelName}`,
+        args,
+        'state:',
+        channel.state
+      );
     });
   });
 };

packages/ws-worker/src/channels/worker-queue.ts

@@ -103,9 +103,15 @@ const connectToWorkerQueue = (
 
     // On close, the socket will try and reconnect itself
     // Forever, so far as I can tell
-    socket.onClose((_e: any) => {
-      logger.debug('queue socket closed');
-      events.emit('disconnect');
+    socket.onClose((e: any) => {
+      logger.warn(
+        `queue socket closed: code=${e?.code} reason=${e?.reason} clean=${e?.wasClean}`
+      );
+      events.emit('disconnect', {
+        code: e?.code,
+        reason: e?.reason,
+        wasClean: e?.wasClean,
+      });
     });
 
     // if we fail to connect, the socket will try to reconnect

packages/ws-worker/src/errors.ts

@@ -35,3 +35,24 @@ export class LightningTimeoutError extends Error {
     super(`[${event}] timeout`);
   }
 }
+
+export type SocketCloseDetails = {
+  code?: number;
+  reason?: string;
+  wasClean?: boolean;
+};
+
+export class LightningSocketClosedError extends Error {
+  name = 'LightningSocketClosedError';
+  code?: number;
+  reason?: string;
+  wasClean?: boolean;
+  constructor({ code, reason, wasClean }: SocketCloseDetails = {}) {
+    super(
+      `Lightning socket closed: code=${code} reason=${reason ?? 'unknown'}`
+    );
+    this.code = code;
+    this.reason = reason;
+    this.wasClean = wasClean;
+  }
+}

packages/ws-worker/src/events/step-complete.ts

@@ -98,10 +98,14 @@ export default async function onStepComplete(
     ]);
   } else {
     evt.output_dataclip_id = dataclipId;
+    // Write the dataclip if it's not too big
     if (!options || options.outputDataclips !== false) {
-      const payload = stringify(outputState);
-      // Write the dataclip if it's not too big
-      evt.output_dataclip = payload;
+      // For back compatibility, stringify the the state object before sending
+      // Note that this causes payloads to bloat
+      // In a major version soon, we should remove the option and never stringify
+      evt.output_dataclip = options?.noStringifyState
+        ? outputState
+        : stringify(outputState);
     }
   }
 

packages/ws-worker/src/mock/sockets.ts

@@ -4,6 +4,8 @@ type EventHandler = (evt?: any) => void;
 export const mockChannel = (
   callbacks: Record<string, EventHandler> = {}
 ): any => {
+  const closeCallbacks: EventHandler[] = [];
+  const errorCallbacks: EventHandler[] = [];
   const c = {
     on: (event: string, fn: EventHandler) => {
       // TODO support multiple callbacks
@@ -71,8 +73,21 @@ export const mockChannel = (
       return receive;
     },
     leave: () => {},
-    onClose: () => {},
-    onError: () => {},
+    // Real phoenix channels support multiple onClose/onError bindings (each
+    // call pushes onto an array), which is now relied on in production -
+    // run.ts and execute.ts both bind onError on the same channel. So this
+    // collects every registered callback rather than keeping only the last
+    onClose: (fn: EventHandler) => {
+      closeCallbacks.push(fn);
+    },
+    onError: (fn: EventHandler) => {
+      errorCallbacks.push(fn);
+    },
+    // test helpers: fire every registered callback, as the real socket would
+    _triggerClose: (...args: any[]) =>
+      closeCallbacks.forEach((fn) => fn(...args)),
+    _triggerError: (...args: any[]) =>
+      errorCallbacks.forEach((fn) => fn(...args)),
   };
   return c;
 };

packages/ws-worker/src/server.ts

@@ -31,6 +31,7 @@ import { convertRun } from './util';
 import parseWorkloops from './util/parse-workloops';
 import getDefaultWorkloopConfig from './util/get-default-workloop-config';
 import { matchesIgnoredError } from './util/ignored-errors';
+import { LightningSocketClosedError, SocketCloseDetails } from './errors';
 
 const exec = promisify(_exec);
 
@@ -61,6 +62,7 @@ export type ServerOptions = {
   claimTimeoutSeconds?: number;
   payloadLimitMb?: number; // max memory limit for socket payload (ie, step:complete, log)
   logPayloadLimitMb?: number; // max memory limit for log payloads specifically
+  noStringifyState?: boolean; // send output dataclips as native JSON instead of a pre-stringified string. Requires lightning support
   collectionsVersion?: string;
   collectionsUrl?: string;
   monorepoDir?: string;
@@ -137,18 +139,32 @@ function connect(app: ServerApp, logger: Logger, options: ServerOptions = {}) {
   };
 
   // We were disconnected from the queue
-  const onDisconnect = () => {
+  const onDisconnect = (details: SocketCloseDetails = {}) => {
     for (const w of app.workloops) {
       if (!w.isStopped()) {
         w.stop('Socket disconnected unexpectedly');
       }
     }
     if (!app.destroyed) {
-      logger.info('Connection to lightning lost');
+      logger.info(
+        `Connection to lightning lost (code=${details.code} reason=${details.reason} clean=${details.wasClean})`
+      );
       logger.info(
         'Worker will automatically reconnect when lightning is back online'
       );
-      // So far as I know, the socket will try and reconnect in the background forever
+      Sentry.captureException(
+        new LightningSocketClosedError(details),
+        (scope) => {
+          scope.setFingerprint([
+            'LightningSocketClosedError',
+            String(details.code),
+            details.reason ?? '',
+          ]);
+          scope.setTag('close_code', String(details.code));
+          scope.setExtras(details);
+          return scope;
+        }
+      );
     }
   };
 
@@ -353,6 +369,7 @@ function createServer(engine: RuntimeEngine, options: ServerOptions = {}) {
           options.logPayloadLimitMb = app.options.logPayloadLimitMb;
         }
 
+        options.noStringifyState = app.options.noStringifyState;
         options.timeoutRetryCount = app.options.timeoutRetryCount;
         options.timeoutRetryDelay =
           app.options.timeoutRetryDelayMs ?? app.options.socketTimeoutSeconds;

packages/ws-worker/src/start.ts

@@ -59,6 +59,7 @@ function engineReady(engine: any) {
     maxWorkflows: effectiveCapacity,
     workloopConfigs,
     payloadLimitMb: args.payloadMemory,
+    noStringifyState: args.noStringifyState,
     logPayloadLimitMb: args.logPayloadMemory ?? 1, // Default to 1MB
     collectionsVersion: args.collectionsVersion,
     collectionsUrl: args.collectionsUrl,

packages/ws-worker/src/util/cli.ts

@@ -33,6 +33,7 @@ type Args = {
   messageTimeoutSeconds?: number;
   mock?: boolean;
   monorepoDir?: string;
+  noStringifyState?: boolean;
   payloadMemory?: number;
   port?: number;
   profile?: boolean;
@@ -96,6 +97,7 @@ export default function parseArgs(argv: string[]): Args {
     WORKER_MAX_RUN_MEMORY_MB,
     WORKER_MAX_STATE_MEMORY_MB,
     WORKER_MESSAGE_TIMEOUT_SECONDS,
+    WORKER_NO_STRINGIFY_STATE,
     WORKER_PORT,
     WORKER_PROFILE_POLL_INTERVAL_MS,
     WORKER_PROFILE,
@@ -233,6 +235,11 @@ export default function parseArgs(argv: string[]): Args {
         'Maximum memory allocated to a single run, in mb. Env: WORKER_MAX_PAYLOAD_MB',
       type: 'number',
     })
+    .option('stringify-state', {
+      description:
+        'Pass --no-stringify-state or set WORKER_NO_STRINGIFY_STATE to optimize stateful payloads sent to lightning. Not back compatible with lightning versions older than 2.19.',
+      type: 'boolean',
+    })
     .option('cgroup', {
       alias: ['enable-cgroup-enforcement', 'cgroups'],
       description:
@@ -298,7 +305,7 @@ export default function parseArgs(argv: string[]): Args {
       'production start configuration with 1 fast lane workloop (capacity 1) and a second workloop with capacity 4'
     );
 
-  const args = parser.parse() as Args;
+  const args = parser.parse() as Args & { stringifyState?: boolean };
 
   const resolvedWorkloops = setArg(args.workloops, WORKER_WORKLOOPS) as
     | string
@@ -348,6 +355,13 @@ export default function parseArgs(argv: string[]): Args {
         ? parseInt(WORKER_MAX_STATE_MEMORY_MB, 10)
         : undefined),
     payloadMemory: setArg(args.payloadMemory, WORKER_MAX_PAYLOAD_MB, 10),
+    // args.stringifyState is positively framed (see the --stringify-state
+    // option above); everything downstream of parseArgs uses the negatively
+    // framed noStringifyState, matching WORKER_NO_STRINGIFY_STATE
+    noStringifyState:
+      args.stringifyState !== undefined
+        ? !args.stringifyState
+        : setArg(undefined, WORKER_NO_STRINGIFY_STATE, false),
     logPayloadMemory: setArg(
       args.logPayloadMemory,
       WORKER_MAX_LOG_PAYLOAD_MB,

packages/ws-worker/src/util/convert-lightning-plan.ts

@@ -33,6 +33,9 @@ export type WorkerRunOptions = ExecuteOptions & {
   outputDataclips?: boolean;
   payloadLimitMb?: number;
   logPayloadLimitMb?: number;
+  // Send the output dataclip as a native JSON value instead of a
+  // pre-stringified string. Defaults to false (old behaviour)
+  noStringifyState?: boolean;
   jobLogLevel?: LogLevel;
   timeoutRetryCount?: number;
   timeoutRetryDelay?: number;

packages/ws-worker/test/api/execute.test.ts

@@ -402,6 +402,32 @@ test('execute should return a context object', async (t) => {
   });
 });
 
+test('execute should breadcrumb a channel error onto the run scope', async (t) => {
+  const channel = mockChannel(mockEventHandlers);
+  const engine = await createMockRTE();
+  const logger = createMockLogger();
+
+  const plan = {
+    id: 'a',
+    workflow: {
+      steps: [
+        {
+          expression: 'fn(() => ({ done: true }))',
+        },
+      ],
+    },
+  } as ExecutionPlan;
+
+  const context = execute(channel, engine, logger, plan, {}, {}, () => {});
+
+  channel._triggerError('boom');
+
+  const breadcrumbs = context.sentryScope!.getScopeData().breadcrumbs;
+  const found = breadcrumbs.find((b: any) => b.message === 'Channel error');
+  t.truthy(found);
+  t.is(found!.category, 'channel');
+});
+
 // TODO this is more of an engine test really, but worth having I suppose
 test('execute should lazy-load a credential', async (t) => {
   const logger = createMockLogger();

packages/ws-worker/test/channels/run.test.ts

@@ -43,3 +43,21 @@ test('should fail to join an run channel with an invalid token', async (t) => {
     t.pass();
   }
 });
+
+test('should log an error including channel state when the channel errors', async (t) => {
+  const logger = createMockLogger();
+  const channel = mockChannel({
+    join: () => ({ status: 'ok' }),
+    [GET_PLAN]: () => runs['run-1'],
+  });
+  const socket = new MockSocket('www', { 'run:a': channel });
+
+  await joinRunChannel(socket, 'x.y.z', 'a', logger);
+
+  channel.state = 'errored';
+  channel._triggerError('boom');
+
+  const log = logger._find('error', /Critical error in channel run:a/);
+  t.truthy(log);
+  t.regex(log!.message as string, /errored/);
+});

packages/ws-worker/test/channels/worker-queue.test.ts

@@ -24,6 +24,32 @@ test('should connect', (t) => {
   });
 });
 
+test('should emit disconnect with the close code, reason and wasClean', (t) => {
+  return new Promise((done) => {
+    connectToWorkerQueue('www', 'a', 'secret', logger, {
+      SocketConstructor: MockSocket as any,
+    })
+      .on('connect', ({ socket }) => {
+        // Real phoenix sockets invoke onClose with a CloseEvent-like object -
+        // MockSocket doesn't drive this itself, so trigger it directly
+        // @ts-ignore accessing test-only internals
+        socket.callbacks.onClose({
+          code: 1009,
+          reason: 'message too big',
+          wasClean: false,
+        });
+      })
+      .on('disconnect', (details) => {
+        t.deepEqual(details, {
+          code: 1009,
+          reason: 'message too big',
+          wasClean: false,
+        });
+        done();
+      });
+  });
+});
+
 test('should connect with an auth token', async (t) => {
   return new Promise((done) => {
     const workerId = 'x';