Extract buffer management logic into KeyBuffer class and keep log4j appender minimal.

phipag · phipag · commit 966e5766f42d · 2025-08-29T17:47:01.000+02:00
diff --git a/powertools-logging/powertools-logging-log4j/src/main/java/software/amazon/lambda/powertools/logging/log4j/BufferingAppender.java b/powertools-logging/powertools-logging-log4j/src/main/java/software/amazon/lambda/powertools/logging/log4j/BufferingAppender.java
@@ -15,10 +15,7 @@
 package software.amazon.lambda.powertools.logging.log4j;
 
 import java.io.Serializable;
-import java.util.ArrayDeque;
 import java.util.Deque;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 
 import org.apache.logging.log4j.Level;
 import org.apache.logging.log4j.core.Appender;
@@ -35,24 +32,64 @@
 import org.apache.logging.log4j.core.config.plugins.PluginElement;
 import org.apache.logging.log4j.core.config.plugins.PluginFactory;
 import org.apache.logging.log4j.core.impl.Log4jLogEvent;
-import org.apache.logging.log4j.message.SimpleMessage;
 
 import software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor;
+import software.amazon.lambda.powertools.logging.internal.BufferManager;
+import software.amazon.lambda.powertools.logging.internal.KeyBuffer;
+
+import static software.amazon.lambda.powertools.logging.log4j.BufferingAppenderConstants.NAME;
 
 /**
- * A minimalistic Log4j2 appender that buffers log events based on trace ID
- * and flushes them when error logs are encountered or manually triggered.
+ * A Log4j2 appender that buffers log events by AWS X-Ray trace ID for optimized Lambda logging.
+ * 
+ * <p>This appender is designed specifically for AWS Lambda functions to reduce log ingestion
+ * by buffering lower-level logs and only outputting them when errors occur, preserving
+ * full context for troubleshooting while minimizing routine log volume.
+ * 
+ * <h3>Key Features:</h3>
+ * <ul>
+ *   <li><strong>Trace-based buffering:</strong> Groups logs by AWS X-Ray trace ID</li>
+ *   <li><strong>Selective output:</strong> Only buffers logs at or below configured verbosity level</li>
+ *   <li><strong>Auto-flush on errors:</strong> Automatically outputs buffered logs when ERROR/FATAL events occur</li>
+ *   <li><strong>Memory management:</strong> Prevents memory leaks with configurable buffer size limits</li>
+ *   <li><strong>Overflow protection:</strong> Warns when logs are discarded due to buffer limits</li>
+ * </ul>
+ * 
+ * <h3>Configuration Example:</h3>
+ * <pre>{@code
+ * <BufferingAppender name="BufferedAppender" 
+ *                    bufferAtVerbosity="INFO" 
+ *                    maxBytes="20480" 
+ *                    flushOnErrorLog="true">
+ *   <AppenderRef ref="ConsoleAppender"/>
+ * </BufferingAppender>
+ * }</pre>
+ * 
+ * <h3>Configuration Parameters:</h3>
+ * <ul>
+ *   <li><strong>bufferAtVerbosity:</strong> Log level to buffer (default: DEBUG). Logs at this level and below are buffered</li>
+ *   <li><strong>maxBytes:</strong> Maximum buffer size in bytes per trace ID (default: 20480)</li>
+ *   <li><strong>flushOnErrorLog:</strong> Whether to flush buffer on ERROR/FATAL logs (default: true)</li>
+ * </ul>
+ * 
+ * <h3>Behavior:</h3>
+ * <ul>
+ *   <li>During Lambda INIT phase (no trace ID): logs are output directly</li>
+ *   <li>During Lambda execution (with trace ID): logs are buffered or output based on level</li>
+ *   <li>When buffer overflows: oldest logs are discarded and a warning is logged</li>
+ *   <li>On Lambda completion: remaining buffered logs can be flushed via {@link software.amazon.lambda.powertools.logging.PowertoolsLogging}</li>
+ * </ul>
+ * 
+ * @see software.amazon.lambda.powertools.logging.PowertoolsLogging#flushLogBuffer()
  */
-@Plugin(name = "BufferingAppender", category = Core.CATEGORY_NAME, elementType = Appender.ELEMENT_TYPE)
-public class BufferingAppender extends AbstractAppender {
+@Plugin(name = NAME, category = Core.CATEGORY_NAME, elementType = Appender.ELEMENT_TYPE)
+public class BufferingAppender extends AbstractAppender implements BufferManager {
 
     private final AppenderRef[] appenderRefs;
     private final Configuration configuration;
     private final Level bufferAtVerbosity;
-    private final int maxBytes;
     private final boolean flushOnErrorLog;
-    private final Map<String, Deque<LogEvent>> bufferCache = new ConcurrentHashMap<>();
-    private final ThreadLocal<Boolean> bufferOverflowTriggered = new ThreadLocal<>();
+    private final KeyBuffer<String, LogEvent> buffer;
 
     protected BufferingAppender(String name, Filter filter, Layout<? extends Serializable> layout,
             AppenderRef[] appenderRefs, Configuration configuration, Level bufferAtVerbosity, int maxBytes,
@@ -61,18 +98,18 @@ protected BufferingAppender(String name, Filter filter, Layout<? extends Seriali
         this.appenderRefs = appenderRefs;
         this.configuration = configuration;
         this.bufferAtVerbosity = bufferAtVerbosity;
-        this.maxBytes = maxBytes;
         this.flushOnErrorLog = flushOnErrorLog;
+        this.buffer = new KeyBuffer<>(maxBytes, event -> event.getMessage().getFormattedMessage().length());
     }
 
     @Override
     public void append(LogEvent event) {
         if (appenderRefs == null || appenderRefs.length == 0) {
             return;
         }
+
         LambdaHandlerProcessor.getXrayTraceId().ifPresentOrElse(
                 traceId -> {
-                    // Check if we should buffer this log level
                     if (shouldBuffer(event.getLevel())) {
                         bufferEvent(traceId, event);
                     } else {
@@ -102,62 +139,22 @@ private boolean shouldBuffer(Level level) {
     }
 
     private void bufferEvent(String traceId, LogEvent event) {
-        // Create immutable copy to prevent mutation
         LogEvent immutableEvent = Log4jLogEvent.createMemento(event);
-
-        // Check if single event is larger than buffer - discard if so
-        int eventSize = immutableEvent.getMessage().getFormattedMessage().length();
-        if (eventSize > maxBytes) {
-            if (Boolean.TRUE != bufferOverflowTriggered.get()) {
-                bufferOverflowTriggered.set(true);
-            }
-            return;
-        }
-
-        bufferCache.computeIfAbsent(traceId, k -> new ArrayDeque<>()).add(immutableEvent);
-
-        // Simple size check - remove oldest if over limit
-        Deque<LogEvent> buffer = bufferCache.get(traceId);
-        while (getBufferSize(buffer) > maxBytes && !buffer.isEmpty()) {
-            if (Boolean.TRUE != bufferOverflowTriggered.get()) {
-                bufferOverflowTriggered.set(true);
-            }
-            buffer.removeFirst();
-        }
-    }
-
-    private int getBufferSize(Deque<LogEvent> buffer) {
-        return buffer.stream()
-                .mapToInt(event -> event.getMessage().getFormattedMessage().length())
-                .sum();
+        buffer.add(traceId, immutableEvent);
     }
 
     public void clearBuffer() {
-        LambdaHandlerProcessor.getXrayTraceId().ifPresent(bufferCache::remove);
+        LambdaHandlerProcessor.getXrayTraceId().ifPresent(buffer::clear);
     }
 
     public void flushBuffer() {
         LambdaHandlerProcessor.getXrayTraceId().ifPresent(this::flushBuffer);
     }
 
     private void flushBuffer(String traceId) {
-        // Emit buffer overflow warning if it occurred
-        if (Boolean.TRUE == bufferOverflowTriggered.get()) {
-            // Create LogEvent directly since Log4j status logger may not reach target appenders
-            LogEvent warningEvent = Log4jLogEvent.newBuilder()
-                    .setLoggerName("BufferingAppender")
-                    .setLevel(Level.WARN)
-                    .setMessage(new SimpleMessage(
-                            "Some logs are not displayed because they were evicted from the buffer. Increase buffer size to store more logs in the buffer."))
-                    .setTimeMillis(System.currentTimeMillis())
-                    .build();
-            callAppenders(warningEvent);
-            bufferOverflowTriggered.remove();
-        }
-
-        Deque<LogEvent> buffer = bufferCache.remove(traceId);
-        if (buffer != null) {
-            buffer.forEach(this::callAppenders);
+        Deque<LogEvent> events = buffer.removeAll(traceId);
+        if (events != null) {
+            events.forEach(this::callAppenders);
         }
     }
 
diff --git a/powertools-logging/powertools-logging-log4j/src/main/java/software/amazon/lambda/powertools/logging/log4j/BufferingAppenderConstants.java b/powertools-logging/powertools-logging-log4j/src/main/java/software/amazon/lambda/powertools/logging/log4j/BufferingAppenderConstants.java
@@ -0,0 +1,30 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates.
+ * Licensed under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package software.amazon.lambda.powertools.logging.log4j;
+
+/**
+ * Constants for BufferingAppender configuration and references.
+ */
+public final class BufferingAppenderConstants {
+
+    /**
+     * The name used for BufferingAppender in Log4j2 configuration and references.
+     */
+    public static final String NAME = "BufferingAppender";
+
+    private BufferingAppenderConstants() {
+        // Utility class
+    }
+}
diff --git a/powertools-logging/powertools-logging-log4j/src/main/java/software/amazon/lambda/powertools/logging/log4j/internal/Log4jLoggingManager.java b/powertools-logging/powertools-logging-log4j/src/main/java/software/amazon/lambda/powertools/logging/log4j/internal/Log4jLoggingManager.java
@@ -24,6 +24,8 @@
 import software.amazon.lambda.powertools.logging.internal.LoggingManager;
 import software.amazon.lambda.powertools.logging.log4j.BufferingAppender;
 
+import static software.amazon.lambda.powertools.logging.log4j.BufferingAppenderConstants.NAME;
+
 /**
  * LoggingManager for Log4j2 that provides log level management and buffer operations.
  * Implements both {@link LoggingManager} and {@link BufferManager} interfaces.
@@ -57,7 +59,7 @@ public org.slf4j.event.Level getLogLevel(Logger logger) {
     public void flushBuffer() {
         LoggerContext ctx = (LoggerContext) LogManager.getContext(false);
         BufferingAppender bufferingAppender = (BufferingAppender) ctx.getConfiguration()
-                .getAppender("BufferedAppender");
+                .getAppender(NAME);
         if (bufferingAppender != null) {
             bufferingAppender.flushBuffer();
         }
@@ -70,7 +72,7 @@ public void flushBuffer() {
     public void clearBuffer() {
         LoggerContext ctx = (LoggerContext) LogManager.getContext(false);
         BufferingAppender bufferingAppender = (BufferingAppender) ctx.getConfiguration()
-                .getAppender("BufferedAppender");
+                .getAppender(NAME);
         if (bufferingAppender != null) {
             bufferingAppender.clearBuffer();
         }
diff --git a/powertools-logging/src/main/java/software/amazon/lambda/powertools/logging/internal/KeyBuffer.java b/powertools-logging/src/main/java/software/amazon/lambda/powertools/logging/internal/KeyBuffer.java
@@ -0,0 +1,114 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates.
+ * Licensed under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package software.amazon.lambda.powertools.logging.internal;
+
+import java.util.ArrayDeque;
+import java.util.Deque;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Generic buffer data structure for storing events by key with size-based eviction.
+ * 
+ * <p>This buffer maintains separate event queues for each key, with configurable size limits
+ * to prevent memory exhaustion. When buffers exceed their size limit, older events are
+ * automatically evicted to make room for newer ones.
+ * 
+ * <h3>Key Features:</h3>
+ * <ul>
+ *   <li><strong>Per-key buffering:</strong> Each key maintains its own independent buffer</li>
+ *   <li><strong>Size-based eviction:</strong> Oldest events are removed when buffer size exceeds limit</li>
+ *   <li><strong>Overflow protection:</strong> Events larger than buffer size are rejected entirely</li>
+ *   <li><strong>Thread-safe:</strong> Supports concurrent access across different keys</li>
+ *   <li><strong>Overflow tracking:</strong> Logs warnings when events are evicted or rejected</li>
+ * </ul>
+ * 
+ * <h3>Eviction Behavior:</h3>
+ * <ul>
+ *   <li><strong>Buffer overflow:</strong> When adding an event would exceed maxBytes, oldest events are evicted first</li>
+ *   <li><strong>Large events:</strong> Events larger than maxBytes are rejected without evicting existing events</li>
+ *   <li><strong>FIFO eviction:</strong> Events are removed in first-in-first-out order during overflow</li>
+ *   <li><strong>Overflow warnings:</strong> Automatic logging when events are evicted or rejected</li>
+ * </ul>
+ * 
+ * <h3>Thread Safety:</h3>
+ * <p>This class is thread-safe for concurrent operations. Different keys can be accessed
+ * simultaneously, and operations on the same key are synchronized to prevent data corruption.
+ * 
+ * @param <K> the type of key used for buffering (e.g., String for trace IDs)
+ * @param <T> the type of events to buffer (must be compatible with the size calculator)
+ */
+public class KeyBuffer<K, T> {
+
+    private static final Logger logger = LoggerFactory.getLogger(KeyBuffer.class);
+
+    private final Map<K, Deque<T>> keyBufferCache = new ConcurrentHashMap<>();
+    private final Map<K, Boolean> overflowTriggered = new ConcurrentHashMap<>();
+    private final int maxBytes;
+    private final Function<T, Integer> sizeCalculator;
+
+    public KeyBuffer(int maxBytes, Function<T, Integer> sizeCalculator) {
+        this.maxBytes = maxBytes;
+        this.sizeCalculator = sizeCalculator;
+    }
+
+    public void add(K key, T event) {
+        int eventSize = sizeCalculator.apply(event);
+        if (eventSize > maxBytes) {
+            overflowTriggered.put(key, true);
+            return;
+        }
+
+        Deque<T> buffer = keyBufferCache.computeIfAbsent(key, k -> new ArrayDeque<>());
+        synchronized (buffer) {
+            buffer.add(event);
+            while (getBufferSize(buffer) > maxBytes && !buffer.isEmpty()) {
+                overflowTriggered.put(key, true);
+                buffer.removeFirst();
+            }
+        }
+    }
+
+    public Deque<T> removeAll(K key) {
+        logOverflowWarningIfNeeded(key);
+        Deque<T> buffer = keyBufferCache.remove(key);
+        if (buffer != null) {
+            synchronized (buffer) {
+                return new ArrayDeque<>(buffer);
+            }
+        }
+        return buffer;
+    }
+
+    public void clear(K key) {
+        keyBufferCache.remove(key);
+        overflowTriggered.remove(key);
+    }
+
+    private void logOverflowWarningIfNeeded(K key) {
+        if (Boolean.TRUE.equals(overflowTriggered.remove(key))) {
+            logger.warn(
+                    "Some logs are not displayed because they were evicted from the buffer. Increase buffer size to store more logs in the buffer.");
+        }
+    }
+
+    private int getBufferSize(Deque<T> buffer) {
+        return buffer.stream().mapToInt(sizeCalculator::apply).sum();
+    }
+}
diff --git a/powertools-logging/src/test/java/software/amazon/lambda/powertools/logging/internal/KeyBufferTest.java b/powertools-logging/src/test/java/software/amazon/lambda/powertools/logging/internal/KeyBufferTest.java
