Bunyan log size limiting stream

LogStream.test.ts

import bunyan from "bunyan";
import fs from "fs";
import { LogStream, trimStringToJsonLength } from "LogStream";

describe("LogStream", () => {
  const outputStream = fs.createWriteStream("/dev/null");
  const logger = bunyan.createLogger({
    name: "test",
    stream: new LogStream({ stream: outputStream, maxLength: 1000 }),
  });
  let spiedWrite: jest.SpiedFunction<typeof outputStream.write>;

  beforeEach(() => {
    spiedWrite = jest
      .spyOn(outputStream, "write")
      .mockImplementation((_buffer, encodingOrCb, cb) => {
        (cb ?? (encodingOrCb as unknown as Function))?.();
        return true;
      });
  });

  afterEach(() => {
    jest.restoreAllMocks();
  });

  function getStdoutText() {
    return spiedWrite.mock.calls
      .map(([buf]) =>
        buf instanceof Buffer || buf instanceof Uint8Array
          ? buf.toString("utf8")
          : buf,
      )
      .join("");
  }

  it("should truncate very long messages", () => {
    logger.info("a".repeat(1_000));

    const text = getStdoutText();
    expect(text.length).toBeLessThanOrEqual(1000); // varies based on date, PID etc.
    const json: { msg: string } = JSON.parse(text);
    expect(json.msg.length).toBeLessThan(1_000);
    expect(json.msg.endsWith("<truncated>")).toBe(true);
  });

  it("should truncate very long field with newlines", () => {
    logger.info(
      { reasonableField: "hello", largeField: "hello\n".repeat(1_000) },
      "reasonable message",
    );

    const text = getStdoutText();
    expect(text.length).toBeLessThanOrEqual(1000); // varies based on date, PID etc.
    const json: { msg: string; reasonableField: string; largeField: string } =
      JSON.parse(text);
    expect(json.msg).toEqual("reasonable message");
    expect(json.reasonableField).toEqual("hello");
    expect(json.largeField).toMatch(/^hello\nhello\n.+ <truncated>$/m);
  });

  it("should truncate very large structured log", () => {
    const largeArray = new Array(1024 * 1024).fill("a");

    logger.info({ a: "hello", b: largeArray, c: "world" }, "test");

    const text = getStdoutText();
    expect(text.length).toBeLessThanOrEqual(1000); // varies based on date, PID etc.
    expect(JSON.parse(text)).toEqual({
      name: "test",
      hostname: expect.any(String),
      pid: expect.any(Number),
      level: 30,
      a: "hello",
      b: "<too large>",
      c: "world",
      msg: "test",
      time: expect.any(String),
      v: 0,
    });
  });
});

describe("trimStringToJsonLength", () => {
  it("should trim string to length", () => {
    expect(trimStringToJsonLength("hello", 0)).toEqual("");
    expect(trimStringToJsonLength("hello", 1)).toEqual("");
    expect(trimStringToJsonLength("hello", 2)).toEqual("");
    expect(trimStringToJsonLength("hello", 3)).toEqual("h"); // JSON.parse('"h"').length === 3
    expect(trimStringToJsonLength("hello\nworld\n", 6)).toEqual("hell"); // JSON.parse('"hel"').length === 5
    expect(trimStringToJsonLength("hello\nworld\n", 7)).toEqual("hello");
    expect(trimStringToJsonLength("hello\nworld\n", 8)).toEqual("hello"); // same since we can't fit the newline in 1 extra byte
    expect(trimStringToJsonLength("hello\nworld\n", 9)).toEqual("hello\n");
  });
});

LogStream.ts

import os from "os";
import stream from "stream";

export const bunyanCoreFields = new Set([
  "name",
  "level",
  "time",
  "pid",
  "hostname",
  "v",
  "src",
]);
const truncatedStringSuffix = " <truncated>";

/**
 * A stream that avoids emitting log records that are too large to be practical,
 * and may cause syslog to split the message into multiple records.
 * The solution is to redact large fields until the message fits.
 *
 * @example
 *
 * const logger = bunyan.createLogger({
 *  name: "test",
 *  stream: new LogStream({ stream: process.stdout, maxLength: 5000 }),
 * });
 *
 */
export class LogStream extends stream.Writable {
  readonly stream: stream.Writable;
  readonly maxLength: number;

  constructor(args: { stream: stream.Writable; maxLength: number }) {
    super({
      decodeStrings: false,
    });
    this.stream = args.stream;
    this.maxLength = args.maxLength;
  }

  _write(
    chunk: string,
    _encoding: BufferEncoding,
    cb: (error?: Error | null) => void,
  ): boolean {
    if (chunk.length > this.maxLength) {
      const data: { msg?: string; [k: string]: any } = JSON.parse(chunk);

      // Try to truncate the largest fields first.
      const fieldsLargestToSmallest = Object.entries(data)
        .filter(([key]) => !bunyanCoreFields.has(key))
        .map(([key, value]) => [JSON.stringify(value).length, key, value]);
      fieldsLargestToSmallest.sort(([l1], [l2]) => l2 - l1); // descending

      for (const [flen, fkey, fval] of fieldsLargestToSmallest) {
        switch (typeof fval) {
          case "string":
            const otherFieldsLength = chunk.length - flen;
            const available =
              this.maxLength - otherFieldsLength - truncatedStringSuffix.length;
            data[fkey] =
              trimStringToJsonLength(fval, available) + truncatedStringSuffix;
            break;
          case "object":
            data[fkey] = "<too large>";
            break;
          default:
            continue; // no use triming numbers, booleans
        }

        chunk = JSON.stringify(data) + os.EOL;
        if (chunk.length <= this.maxLength) break; // if it helped, stop.
      }
    }

    return this.stream.write(chunk, cb);
  }
}

/**
 * Trims a string so that its length *represented in JSON* is at most `l`.
 */
export function trimStringToJsonLength(s: string, l: number): string {
  const escaped = JSON.stringify(s).slice(1, -1);
  let trimmed = escaped.slice(0, Math.max(l - 2, 0));
  if (trimmed.endsWith("\\")) trimmed = trimmed.slice(0, -1);
  return JSON.parse(`"${trimmed}"`);
}