attributes for handler loggers

Author: nicktrn

apps/coordinator/src/index.ts

@@ -3,6 +3,7 @@ import { Server } from "socket.io";
 import {
   CoordinatorToPlatformMessages,
   CoordinatorToProdWorkerMessages,
+  omit,
   PlatformToCoordinatorMessages,
   ProdWorkerSocketData,
   ProdWorkerToCoordinatorMessages,
@@ -127,17 +128,28 @@ class TaskCoordinator {
       clientMessages: CoordinatorToPlatformMessages,
       serverMessages: PlatformToCoordinatorMessages,
       authToken: PLATFORM_SECRET,
+      logHandlerPayloads: false,
       handlers: {
         RESUME_AFTER_DEPENDENCY: async (message) => {
+          const log = logger.child({
+            eventName: "RESUME_AFTER_DEPENDENCY",
+            ...omit(message, "completions", "executions"),
+            completions: message.completions.map((c) => c.id),
+            executions: message.executions.length,
+          });
+
+          log.debug("Handling RESUME_AFTER_DEPENDENCY");
+
           const taskSocket = await this.#getAttemptSocket(message.attemptFriendlyId);
 
           if (!taskSocket) {
-            logger.log("Socket for attempt not found", {
-              attemptFriendlyId: message.attemptFriendlyId,
-            });
+            log.debug("Socket for attempt not found");
             return;
           }
 
+          log.addFields({ socketData: taskSocket.data });
+          log.log("Found task socket for RESUME_AFTER_DEPENDENCY");
+
           await chaosMonkey.call();
 
           // In case the task resumed faster than we could checkpoint
@@ -146,12 +158,17 @@ class TaskCoordinator {
           taskSocket.emit("RESUME_AFTER_DEPENDENCY", message);
         },
         RESUME_AFTER_DEPENDENCY_WITH_ACK: async (message) => {
+          const log = logger.child({
+            eventName: "RESUME_AFTER_DEPENDENCY_WITH_ACK",
+            ...omit(message, "completions", "executions"),
+          });
+
+          log.debug("Handling RESUME_AFTER_DEPENDENCY_WITH_ACK");
+
           const taskSocket = await this.#getAttemptSocket(message.attemptFriendlyId);
 
           if (!taskSocket) {
-            logger.log("Socket for attempt not found", {
-              attemptFriendlyId: message.attemptFriendlyId,
-            });
+            log.debug("Socket for attempt not found");
             return {
               success: false,
               error: {
@@ -161,11 +178,12 @@ class TaskCoordinator {
             };
           }
 
+          log.addFields({ socketData: taskSocket.data });
+          log.log("Found task socket for RESUME_AFTER_DEPENDENCY_WITH_ACK");
+
           //if this is set, we want to kill the process because it will be resumed with the checkpoint from the queue
           if (taskSocket.data.requiresCheckpointResumeWithMessage) {
-            logger.log("RESUME_AFTER_DEPENDENCY_WITH_ACK: Checkpoint is set so going to nack", {
-              socketData: taskSocket.data,
-            });
+            log.log("RESUME_AFTER_DEPENDENCY_WITH_ACK: Checkpoint is set so going to nack");
 
             return {
               success: false,
@@ -189,6 +207,13 @@ class TaskCoordinator {
           };
         },
         RESUME_AFTER_DURATION: async (message) => {
+          const log = logger.child({
+            eventName: "RESUME_AFTER_DURATION",
+            ...message,
+          });
+
+          log.debug("Handling RESUME_AFTER_DURATION");
+
           const taskSocket = await this.#getAttemptSocket(message.attemptFriendlyId);
 
           if (!taskSocket) {
@@ -198,11 +223,19 @@ class TaskCoordinator {
             return;
           }
 
+          log.addFields({ socketData: taskSocket.data });
+          log.log("Found task socket for RESUME_AFTER_DURATION");
+
           await chaosMonkey.call();
 
           taskSocket.emit("RESUME_AFTER_DURATION", message);
         },
         REQUEST_ATTEMPT_CANCELLATION: async (message) => {
+          const log = logger.child({
+            eventName: "REQUEST_ATTEMPT_CANCELLATION",
+            ...message,
+          });
+
           const taskSocket = await this.#getAttemptSocket(message.attemptFriendlyId);
 
           if (!taskSocket) {
@@ -212,9 +245,17 @@ class TaskCoordinator {
             return;
           }
 
+          log.addFields({ socketData: taskSocket.data });
+          log.log("Found task socket for REQUEST_ATTEMPT_CANCELLATION");
+
           taskSocket.emit("REQUEST_ATTEMPT_CANCELLATION", message);
         },
         REQUEST_RUN_CANCELLATION: async (message) => {
+          const log = logger.child({
+            eventName: "REQUEST_RUN_CANCELLATION",
+            ...message,
+          });
+
           const taskSocket = await this.#getRunSocket(message.runId);
 
           if (!taskSocket) {
@@ -224,6 +265,9 @@ class TaskCoordinator {
             return;
           }
 
+          log.addFields({ socketData: taskSocket.data });
+          log.log("Found task socket for REQUEST_RUN_CANCELLATION");
+
           this.#cancelCheckpoint(message.runId);
 
           if (message.delayInMs) {
@@ -239,6 +283,11 @@ class TaskCoordinator {
           }
         },
         READY_FOR_RETRY: async (message) => {
+          const log = logger.child({
+            eventName: "READY_FOR_RETRY",
+            ...message,
+          });
+
           const taskSocket = await this.#getRunSocket(message.runId);
 
           if (!taskSocket) {
@@ -248,11 +297,19 @@ class TaskCoordinator {
             return;
           }
 
+          log.addFields({ socketData: taskSocket.data });
+          log.log("Found task socket for READY_FOR_RETRY");
+
           await chaosMonkey.call();
 
           taskSocket.emit("READY_FOR_RETRY", message);
         },
         DYNAMIC_CONFIG: async (message) => {
+          const log = logger.child({
+            eventName: "DYNAMIC_CONFIG",
+            ...message,
+          });
+
           this.#delayThresholdInMs = message.checkpointThresholdInMs;
 
           // The first time we receive a dynamic config, the worker namespace will be created

packages/core/src/v3/utils/structuredLogger.ts

@@ -60,6 +60,13 @@ export class SimpleStructuredLogger implements StructuredLogger {
     this.#structuredLog(console.debug, message, "debug", ...args);
   }
 
+  addFields(fields: Record<string, unknown>) {
+    this.fields = {
+      ...this.fields,
+      ...fields,
+    };
+  }
+
   #structuredLog(
     loggerFunction: (message: string, ...args: any[]) => void,
     message: string,

packages/core/src/v3/zodNamespace.ts

@@ -41,6 +41,7 @@ interface ZodNamespaceOptions<
   handlers?: ZodSocketMessageHandlers<TClientMessages>;
   authToken?: string;
   logger?: StructuredLogger;
+  logHandlerPayloads?: boolean;
   preAuth?: (
     socket: ZodNamespaceSocket<TClientMessages, TServerMessages, TServerSideEvents, TSocketData>,
     next: (err?: ExtendedError) => void,
@@ -96,6 +97,7 @@ export class ZodNamespace<
     this.#handler = new ZodSocketMessageHandler({
       schema: opts.clientMessages,
       handlers: opts.handlers,
+      logPayloads: opts.logHandlerPayloads,
     });
 
     this.io = opts.io;

packages/core/src/v3/zodSocket.ts

@@ -68,6 +68,7 @@ export type ZodSocketMessageHandlerOptions<TMessageCatalog extends ZodSocketMess
     schema: TMessageCatalog;
     handlers?: ZodSocketMessageHandlers<TMessageCatalog>;
     logger?: StructuredLogger;
+    logPayloads?: boolean;
   };
 
 type MessageFromSocketSchema<
@@ -92,12 +93,14 @@ export class ZodSocketMessageHandler<TRPCCatalog extends ZodSocketMessageCatalog
   #schema: TRPCCatalog;
   #handlers: ZodSocketMessageHandlers<TRPCCatalog> | undefined;
   #logger: StructuredLogger;
+  #logPayloads: boolean;
 
   constructor(options: ZodSocketMessageHandlerOptions<TRPCCatalog>) {
     this.#schema = options.schema;
     this.#handlers = options.handlers;
     this.#logger =
       options.logger ?? new SimpleStructuredLogger("socket-message-handler", LogLevel.info);
+    this.#logPayloads = options.logPayloads ?? !!process.env.LOG_SOCKET_HANDLER_PAYLOADS ?? false;
   }
 
   public async handleMessage(message: unknown) {
@@ -120,7 +123,7 @@ export class ZodSocketMessageHandler<TRPCCatalog extends ZodSocketMessageCatalog
     const handler = this.#handlers[type];
 
     if (!handler) {
-      console.error(`No handler for message type: ${String(type)}`);
+      this.#logger.error("No handler for message type", { type, payload });
       return;
     }
 
@@ -164,7 +167,7 @@ export class ZodSocketMessageHandler<TRPCCatalog extends ZodSocketMessageCatalog
     const parsedPayload = schema.safeParse(messageWithVersion);
 
     if (!parsedPayload.success) {
-      console.error("Failed to parse message payload", {
+      this.#logger.error("Failed to parse message payload", {
         message,
         payload: messageWithVersion,
       });
@@ -194,8 +197,9 @@ export class ZodSocketMessageHandler<TRPCCatalog extends ZodSocketMessageCatalog
 
     for (const eventName of Object.keys(this.#handlers)) {
       emitter.on(eventName, async (message: any, callback?: any): Promise<void> => {
-        log.info(`handling ${eventName}`, {
-          payload: message,
+        log.info(`Incoming event ${eventName}`, {
+          eventName,
+          ...(this.#logPayloads ? { eventMessage: message } : {}),
           hasCallback: !!callback,
         });
 
@@ -344,6 +348,7 @@ interface ZodSocketConnectionOptions<
   handlers?: ZodSocketMessageHandlers<TServerMessages>;
   authToken?: string;
   ioOptions?: Partial<ManagerOptions & SocketOptions>;
+  logHandlerPayloads?: boolean;
   onConnection?: (
     socket: ZodSocket<TServerMessages, TClientMessages>,
     handler: ZodSocketMessageHandler<TServerMessages>,
@@ -399,6 +404,7 @@ export class ZodSocketConnection<
     this.#handler = new ZodSocketMessageHandler({
       schema: opts.serverMessages,
       handlers: opts.handlers,
+      logPayloads: opts.logHandlerPayloads,
     });
     this.#handler.registerHandlers(this.socket, this.#logger);