feat(logger): Add log buffer and flush method (#3617)

ConnorKirk · web-flow · commit 6968ca87f55e · 2025-02-19T13:34:32.000+01:00
diff --git a/packages/logger/src/Logger.ts b/packages/logger/src/Logger.ts
@@ -13,6 +13,7 @@ import { LogJsonIndent, LogLevelThreshold, ReservedKeys } from './constants.js';
 import type { LogFormatter } from './formatter/LogFormatter.js';
 import type { LogItem } from './formatter/LogItem.js';
 import { PowertoolsLogFormatter } from './formatter/PowertoolsLogFormatter.js';
+import { CircularMap } from './logBuffer.js';
 import type { ConfigServiceInterface } from './types/ConfigServiceInterface.js';
 import type {
   ConstructorOptions,
@@ -142,7 +143,10 @@ class Logger extends Utility implements LoggerInterface {
    * Sometimes we need to log warnings before the logger is fully initialized, however we can't log them
    * immediately because the logger is not ready yet. This buffer stores those logs until the logger is ready.
    */
-  #buffer: [number, Parameters<Logger['createAndPopulateLogItem']>][] = [];
+  readonly #initBuffer: [
+    number,
+    Parameters<Logger['createAndPopulateLogItem']>,
+  ][] = [];
   /**
    * Flag used to determine if the logger is initialized.
    */
@@ -165,6 +169,29 @@ class Logger extends Utility implements LoggerInterface {
    */
   #jsonReplacerFn?: CustomJsonReplacerFn;
 
+  /**
+   * Represents whether the buffering functionality is enabled in the logger
+   */
+  protected isBufferEnabled = false;
+
+  /**
+   * Log level threshold for the buffer
+   * Logs with a level lower than this threshold will be buffered
+   */
+  protected bufferLogThreshold: number = LogLevelThreshold.DEBUG;
+  /**
+   * Max size of the buffer. Additions to the buffer beyond this size will
+   * cause older logs to be evicted from the buffer
+   */
+  readonly #maxBufferBytesSize = 1024;
+
+  /**
+   * Contains buffered logs, grouped by _X_AMZN_TRACE_ID, each group with a max size of `maxBufferBytesSize`
+   */
+  readonly #buffer: CircularMap<string> = new CircularMap({
+    maxBytesSize: this.#maxBufferBytesSize,
+  });
+
   /**
    * Log level used by the current instance of Logger.
    *
@@ -182,11 +209,11 @@ class Logger extends Utility implements LoggerInterface {
     // all logs are buffered until the logger is initialized
     this.setOptions(rest);
     this.#isInitialized = true;
-    for (const [level, log] of this.#buffer) {
+    for (const [level, log] of this.#initBuffer) {
       // we call the method directly and create the log item just in time
       this.printLog(level, this.createAndPopulateLogItem(...log));
     }
-    this.#buffer = [];
+    this.#initBuffer = [];
   }
 
   /**
@@ -919,7 +946,7 @@ class Logger extends Utility implements LoggerInterface {
   }
 
   /**
-   * Print a given log with given log level.
+   * Print or buffer a given log with given log level.
    *
    * @param logLevel - The log level threshold
    * @param input - The log message
@@ -937,7 +964,33 @@ class Logger extends Utility implements LoggerInterface {
           this.createAndPopulateLogItem(logLevel, input, extraInput)
         );
       } else {
-        this.#buffer.push([logLevel, [logLevel, input, extraInput]]);
+        this.#initBuffer.push([logLevel, [logLevel, input, extraInput]]);
+      }
+      return;
+    }
+
+    const traceId = this.envVarsService.getXrayTraceId();
+    if (traceId !== undefined && this.shouldBufferLog(traceId, logLevel)) {
+      try {
+        this.bufferLogItem(
+          traceId,
+          this.createAndPopulateLogItem(logLevel, input, extraInput),
+          logLevel
+        );
+      } catch (error) {
+        this.printLog(
+          LogLevelThreshold.WARN,
+          this.createAndPopulateLogItem(
+            LogLevelThreshold.WARN,
+            `Unable to buffer log: ${(error as Error).message}`,
+            [error as Error]
+          )
+        );
+
+        this.printLog(
+          logLevel,
+          this.createAndPopulateLogItem(logLevel, input, extraInput)
+        );
       }
     }
   }
@@ -1169,6 +1222,68 @@ class Logger extends Utility implements LoggerInterface {
     });
     persistentKeys && this.appendPersistentKeys(persistentKeys);
   }
+
+  /**
+   * Add a log to the buffer
+   * @param xrayTraceId - _X_AMZN_TRACE_ID of the request
+   * @param log - Log to be buffered
+   * @param logLevel - level of log to be buffered
+   */
+  protected bufferLogItem(
+    xrayTraceId: string,
+    log: LogItem,
+    logLevel: number
+  ): void {
+    log.prepareForPrint();
+
+    const stringified = JSON.stringify(
+      log.getAttributes(),
+      this.getJsonReplacer(),
+      this.logIndentation
+    );
+
+    this.#buffer.setItem(xrayTraceId, stringified, logLevel);
+  }
+
+  /**
+   * Flushes all items of the respective _X_AMZN_TRACE_ID within
+   * the buffer.
+   */
+  protected flushBuffer(): void {
+    const traceId = this.envVarsService.getXrayTraceId();
+    if (traceId === undefined) {
+      return;
+    }
+
+    const buffer = this.#buffer.get(traceId) || [];
+
+    for (const item of buffer) {
+      const consoleMethod =
+        item.logLevel === LogLevelThreshold.CRITICAL
+          ? 'error'
+          : (this.getLogLevelNameFromNumber(
+              item.logLevel
+            ).toLowerCase() as keyof Omit<LogFunction, 'critical'>);
+      this.console[consoleMethod](item.value);
+    }
+
+    this.#buffer.delete(traceId);
+  }
+  /**
+   * Tests if the log meets the criteria to be buffered
+   * @param traceId - _X_AMZN_TRACE_ID of the request
+   * @param logLevel - The  level of the log being considered
+   */
+  protected shouldBufferLog(
+    traceId: string | undefined,
+    logLevel: number
+  ): boolean {
+    return (
+      this.isBufferEnabled &&
+      traceId !== undefined &&
+      logLevel <= this.bufferLogThreshold
+    );
+  }
 }
 
 export { Logger };
diff --git a/packages/logger/tests/unit/logBuffer.test.ts b/packages/logger/tests/unit/logBuffer.test.ts
@@ -1,146 +1,110 @@
-import { describe, expect, it, vi } from 'vitest';
-import { CircularMap, SizedItem, SizedSet } from '../../src/logBuffer.js';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { Logger } from '../../src/Logger.js';
+import { LogLevelThreshold } from '../../src/constants.js';
+
+class TestLogger extends Logger {
+  public enableBuffering() {
+    this.isBufferEnabled = true;
+  }
+  public disableBuffering() {
+    this.isBufferEnabled = false;
+  }
+
+  public flushBufferWrapper(): void {
+    this.flushBuffer();
+  }
+
+  public overrideBufferLogItem(): void {
+    this.bufferLogItem = vi.fn().mockImplementation(() => {
+      throw new Error('bufferLogItem error');
+    });
+  }
+
+  public setbufferLevelThreshold(level: number): void {
+    this.bufferLogThreshold = level;
+  }
+}
 
-describe('SizedItem', () => {
-  it('calculates the byteSize based on string value', () => {
+describe('bufferLog', () => {
+  it('outputs a warning when there is an error buffering the log', () => {
     // Prepare
-    const logEntry = 'hello world';
+    process.env.POWERTOOLS_DEV = 'true';
+    const logger = new TestLogger();
+    logger.enableBuffering();
+    logger.overrideBufferLogItem();
 
     // Act
-    const item = new SizedItem(logEntry, 1);
+    logger.debug('This is a debug');
 
     // Assess
-    const expectedByteSize = Buffer.byteLength(logEntry);
-    expect(item.byteSize).toBe(expectedByteSize);
-  });
-
-  it('throws an error if value is not a string', () => {
-    // Prepare
-    const invalidValue = { message: 'not a string' };
-
-    // Act & Assess
-    expect(
-      () => new SizedItem(invalidValue as unknown as string, 1)
-    ).toThrowError('Value should be a string');
+    expect(console.debug).toBeCalledTimes(1);
+    expect(console.warn).toBeCalledTimes(1);
   });
 });
 
-describe('SizedSet', () => {
-  it('adds an item and updates currentBytesSize correctly', () => {
-    // Prepare
-    const set = new SizedSet<string>();
-    const item = new SizedItem('value', 1);
-
-    // Act
-    set.add(item);
+describe('flushBuffer', () => {
+  const ENVIRONMENT_VARIABLES = process.env;
 
-    // Assess
-    expect(set.currentBytesSize).toBe(item.byteSize);
-    expect(set.has(item)).toBe(true);
+  beforeEach(() => {
+    process.env = {
+      ...ENVIRONMENT_VARIABLES,
+      POWERTOOLS_LOGGER_LOG_EVENT: 'true',
+      POWERTOOLS_DEV: 'true',
+    };
+    vi.clearAllMocks();
   });
 
-  it('deletes an item and updates currentBytesSize correctly', () => {
+  it('outputs buffered logs', () => {
     // Prepare
-    const set = new SizedSet<string>();
-    const item = new SizedItem('value', 1);
-    set.add(item);
-    const initialSize = set.currentBytesSize;
+    const logger = new TestLogger({ logLevel: 'SILENT' });
+    logger.enableBuffering();
+    logger.setbufferLevelThreshold(LogLevelThreshold.CRITICAL);
 
     // Act
-    const result = set.delete(item);
+    logger.debug('This is a debug');
+    logger.warn('This is a warning');
+    logger.critical('this is a critical');
 
     // Assess
-    expect(result).toBe(true);
-    expect(set.currentBytesSize).toBe(initialSize - item.byteSize);
-    expect(set.has(item)).toBe(false);
-  });
-
-  it('clears all items and resets currentBytesSize to 0', () => {
-    // Prepare
-    const set = new SizedSet<string>();
-    set.add(new SizedItem('b', 1));
-    set.add(new SizedItem('d', 1));
+    expect(console.warn).toHaveBeenCalledTimes(0);
+    expect(console.error).toHaveBeenCalledTimes(0);
 
     // Act
-    set.clear();
+    logger.flushBufferWrapper();
 
     // Assess
-    expect(set.currentBytesSize).toBe(0);
-    expect(set.size).toBe(0);
+    expect(console.warn).toHaveBeenCalledTimes(1);
+    expect(console.error).toHaveBeenCalledTimes(1);
   });
 
-  it('removes the first inserted item with shift', () => {
+  it('handles an empty buffer', () => {
     // Prepare
-    const set = new SizedSet<string>();
-    const item1 = new SizedItem('first', 1);
-    const item2 = new SizedItem('second', 1);
-    set.add(item1);
-    set.add(item2);
+    const logger = new TestLogger();
+    logger.enableBuffering();
 
     // Act
-    const shiftedItem = set.shift();
-
-    // Assess
-    expect(shiftedItem).toEqual(item1);
-    expect(set.has(item1)).toBe(false);
-    expect(set.currentBytesSize).toBe(item2.byteSize);
+    logger.flushBufferWrapper();
   });
-});
 
-describe('CircularMap', () => {
-  it('adds items to a new buffer for a given key', () => {
+  it('does not output buffered logs when trace id is not set', () => {
     // Prepare
-    const maxBytes = 200;
-    const circularMap = new CircularMap<string>({
-      maxBytesSize: maxBytes,
-    });
+    process.env._X_AMZN_TRACE_ID = undefined;
+    const logger = new TestLogger({});
+    logger.enableBuffering();
 
     // Act
-    circularMap.setItem('trace-1', 'first log', 1);
+    logger.debug('This is a debug');
+    logger.warn('this is a warning');
 
     // Assess
-    const buffer = circularMap.get('trace-1');
-    expect(buffer).toBeDefined();
-    if (buffer) {
-      expect(buffer.currentBytesSize).toBeGreaterThan(0);
-      expect(buffer.size).toBe(1);
-    }
-  });
-
-  it('throws an error when an item exceeds maxBytesSize', () => {
-    // Prepare
-    const maxBytes = 10;
-    const circularMap = new CircularMap<string>({
-      maxBytesSize: maxBytes,
-    });
-
-    // Act & Assess
-    expect(() => {
-      circularMap.setItem('trace-1', 'a very long message', 1);
-    }).toThrowError('Item too big');
-  });
-
-  it('evicts items when the buffer overflows and call the overflow callback', () => {
-    // Prepare
-    const options = {
-      maxBytesSize: 15,
-      onBufferOverflow: vi.fn(),
-    };
-    const circularMap = new CircularMap<string>(options);
-    const smallEntry = '12345';
-
-    const entryByteSize = Buffer.byteLength(smallEntry);
-    const entriesCount = Math.ceil(options.maxBytesSize / entryByteSize);
+    expect(console.debug).toHaveBeenCalledTimes(0);
+    expect(console.warn).toHaveBeenCalledTimes(1);
 
     // Act
-    for (let i = 0; i < entriesCount; i++) {
-      circularMap.setItem('trace-1', smallEntry, 1);
-    }
+    logger.flushBufferWrapper();
 
     // Assess
-    expect(options.onBufferOverflow).toHaveBeenCalledTimes(1);
-    expect(circularMap.get('trace-1')?.currentBytesSize).toBeLessThan(
-      options.maxBytesSize
-    );
+    expect(console.debug).toHaveBeenCalledTimes(0);
+    expect(console.warn).toHaveBeenCalledTimes(1);
   });
 });
diff --git a/packages/logger/tests/unit/logBufferStructures.test.ts b/packages/logger/tests/unit/logBufferStructures.test.ts
