Skip to content

Commit 6dde0a5

Browse files
phipagphipag
committed
Add buffer cleaning to avoid memory leak and flushBufferOnUncaughtError option to Logging aspect.
1 parent 59011e5 commit 6dde0a5

File tree

4 files changed

+168
-0
lines changed

4 files changed

+168
-0
lines changed

‎powertools-logging/src/main/java/software/amazon/lambda/powertools/logging/Logging.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,4 +102,10 @@
102102
* Set this attribute to true if you want all custom keys to be deleted on each request.
103103
*/
104104
boolean clearState() default false;
105+
106+
/**
107+
* Set to true if you want to flush the log buffer when an uncaught exception occurs.
108+
* This ensures that buffered logs are output when errors happen.
109+
*/
110+
boolean flushBufferOnUncaughtError() default true;
105111
}

‎powertools-logging/src/main/java/software/amazon/lambda/powertools/logging/internal/LambdaLoggingAspect.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,15 @@ public Object around(ProceedingJoinPoint pjp,
160160
// Call Function Handler
161161
lambdaFunctionResponse = pjp.proceed(proceedArgs);
162162
} catch (Throwable t) {
163+
if (LOGGING_MANAGER instanceof BufferManager) {
164+
if (logging.flushBufferOnUncaughtError()) {
165+
((BufferManager) LOGGING_MANAGER).flushBuffer();
166+
} else {
167+
// Clear buffer before error logging to prevent unintended flush. If flushOnErrorLog is enabled on
168+
// the appender the next line would otherwise cause an unintended flush by the appender directly.
169+
((BufferManager) LOGGING_MANAGER).clearBuffer();
170+
}
171+
}
163172
if (logging.logError() || POWERTOOLS_LOG_ERROR) {
164173
// logging the exception with additional context
165174
logger(pjp).error(MarkerFactory.getMarker("FATAL"), "Exception in Lambda Handler", t);
@@ -169,6 +178,10 @@ public Object around(ProceedingJoinPoint pjp,
169178
if (logging.clearState()) {
170179
MDC.clear();
171180
}
181+
// Clear buffer after each handler invocation
182+
if (LOGGING_MANAGER instanceof BufferManager) {
183+
((BufferManager) LOGGING_MANAGER).clearBuffer();
184+
}
172185
coldStartDone();
173186
}
174187

‎powertools-logging/src/test/java/software/amazon/lambda/powertools/logging/handlers/PowertoolsLogErrorNoFlush.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
package software.amazon.lambda.powertools.logging.handlers;
2+
3+
import com.amazonaws.services.lambda.runtime.Context;
4+
import com.amazonaws.services.lambda.runtime.RequestHandler;
5+
6+
import software.amazon.lambda.powertools.logging.Logging;
7+
8+
public class PowertoolsLogErrorNoFlush implements RequestHandler<String, String> {
9+
10+
@Override
11+
@Logging(logError = true, flushBufferOnUncaughtError = false)
12+
public String handleRequest(String input, Context context) {
13+
throw new RuntimeException("This is an error without buffer flush");
14+
}
15+
}

‎powertools-logging/src/test/java/software/amazon/lambda/powertools/logging/internal/LambdaLoggingAspectTest.java‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@
7878
import software.amazon.lambda.powertools.logging.handlers.PowertoolsLogEnabled;
7979
import software.amazon.lambda.powertools.logging.handlers.PowertoolsLogEnabledForStream;
8080
import software.amazon.lambda.powertools.logging.handlers.PowertoolsLogError;
81+
import software.amazon.lambda.powertools.logging.handlers.PowertoolsLogErrorNoFlush;
8182
import software.amazon.lambda.powertools.logging.handlers.PowertoolsLogEvent;
8283
import software.amazon.lambda.powertools.logging.handlers.PowertoolsLogEventBridgeCorrelationId;
8384
import software.amazon.lambda.powertools.logging.handlers.PowertoolsLogEventEnvVar;
@@ -108,6 +109,13 @@ void setUp() throws IllegalAccessException, NoSuchMethodException, InvocationTar
108109
writeStaticField(LoggingConstants.class, "POWERTOOLS_LOG_LEVEL", null, true);
109110
writeStaticField(LoggingConstants.class, "POWERTOOLS_LOG_EVENT", false, true);
110111
writeStaticField(LoggingConstants.class, "POWERTOOLS_SAMPLING_RATE", null, true);
112+
113+
// Reset buffer state for clean test isolation
114+
LoggingManager loggingManager = LoggingManagerRegistry.getLoggingManager();
115+
if (loggingManager instanceof TestLoggingManager) {
116+
((TestLoggingManager) loggingManager).resetBufferState();
117+
}
118+
111119
try {
112120
FileChannel.open(Paths.get("target/logfile.json"), StandardOpenOption.WRITE).truncate(0).close();
113121
} catch (NoSuchFileException e) {
@@ -717,6 +725,132 @@ void shouldLogCorrelationIdOnAppSyncEvent() throws IOException {
717725
.containsEntry("correlation_id", eventId);
718726
}
719727

728+
@Test
729+
void shouldClearBufferAfterSuccessfulHandlerExecution() {
730+
// WHEN
731+
requestHandler.handleRequest(new Object(), context);
732+
733+
// THEN
734+
LoggingManager loggingManager = LoggingManagerRegistry.getLoggingManager();
735+
if (loggingManager instanceof TestLoggingManager) {
736+
assertThat(((TestLoggingManager) loggingManager).isBufferCleared()).isTrue();
737+
}
738+
}
739+
740+
@Test
741+
void shouldClearBufferAfterSuccessfulStreamHandlerExecution() throws IOException {
742+
// WHEN
743+
requestStreamHandler.handleRequest(new ByteArrayInputStream(new byte[] {}), new ByteArrayOutputStream(),
744+
context);
745+
746+
// THEN
747+
LoggingManager loggingManager = LoggingManagerRegistry.getLoggingManager();
748+
if (loggingManager instanceof TestLoggingManager) {
749+
assertThat(((TestLoggingManager) loggingManager).isBufferCleared()).isTrue();
750+
}
751+
}
752+
753+
@Test
754+
void shouldClearBufferAfterHandlerExceptionWithLogError() {
755+
// GIVEN
756+
requestHandler = new PowertoolsLogError();
757+
758+
// WHEN
759+
try {
760+
requestHandler.handleRequest("input", context);
761+
} catch (Exception e) {
762+
// ignore
763+
}
764+
765+
// THEN
766+
LoggingManager loggingManager = LoggingManagerRegistry.getLoggingManager();
767+
if (loggingManager instanceof TestLoggingManager) {
768+
assertThat(((TestLoggingManager) loggingManager).isBufferCleared()).isTrue();
769+
}
770+
}
771+
772+
@Test
773+
void shouldClearBufferAfterHandlerExceptionWithEnvVarLogError() {
774+
try {
775+
// GIVEN
776+
LoggingConstants.POWERTOOLS_LOG_ERROR = true;
777+
requestHandler = new PowertoolsLogEnabled(true);
778+
779+
// WHEN
780+
try {
781+
requestHandler.handleRequest("input", context);
782+
} catch (Exception e) {
783+
// ignore
784+
}
785+
786+
// THEN
787+
LoggingManager loggingManager = LoggingManagerRegistry.getLoggingManager();
788+
if (loggingManager instanceof TestLoggingManager) {
789+
assertThat(((TestLoggingManager) loggingManager).isBufferCleared()).isTrue();
790+
}
791+
} finally {
792+
LoggingConstants.POWERTOOLS_LOG_ERROR = false;
793+
}
794+
}
795+
796+
@Test
797+
void shouldFlushBufferOnUncaughtErrorWhenEnabled() {
798+
// GIVEN
799+
requestHandler = new PowertoolsLogError();
800+
801+
// WHEN
802+
try {
803+
requestHandler.handleRequest("input", context);
804+
} catch (Exception e) {
805+
// ignore
806+
}
807+
808+
// THEN
809+
LoggingManager loggingManager = LoggingManagerRegistry.getLoggingManager();
810+
if (loggingManager instanceof TestLoggingManager) {
811+
assertThat(((TestLoggingManager) loggingManager).isBufferFlushed()).isTrue();
812+
}
813+
}
814+
815+
@Test
816+
void shouldNotFlushBufferOnUncaughtErrorWhenDisabled() {
817+
// GIVEN
818+
PowertoolsLogErrorNoFlush handler = new PowertoolsLogErrorNoFlush();
819+
820+
// WHEN
821+
try {
822+
handler.handleRequest("input", context);
823+
} catch (Exception e) {
824+
// ignore
825+
}
826+
827+
// THEN
828+
LoggingManager loggingManager = LoggingManagerRegistry.getLoggingManager();
829+
if (loggingManager instanceof TestLoggingManager) {
830+
assertThat(((TestLoggingManager) loggingManager).isBufferFlushed()).isFalse();
831+
}
832+
}
833+
834+
@Test
835+
void shouldClearBufferBeforeErrorLoggingWhenFlushBufferOnUncaughtErrorDisabled() {
836+
// GIVEN
837+
PowertoolsLogErrorNoFlush handler = new PowertoolsLogErrorNoFlush();
838+
839+
// WHEN
840+
try {
841+
handler.handleRequest("input", context);
842+
} catch (Exception e) {
843+
// ignore
844+
}
845+
846+
// THEN - Buffer should be cleared and not flushed
847+
LoggingManager loggingManager = LoggingManagerRegistry.getLoggingManager();
848+
if (loggingManager instanceof TestLoggingManager) {
849+
assertThat(((TestLoggingManager) loggingManager).isBufferCleared()).isTrue();
850+
assertThat(((TestLoggingManager) loggingManager).isBufferFlushed()).isFalse();
851+
}
852+
}
853+
720854
private void resetLogLevel(Level level)
721855
throws NoSuchMethodException, IllegalAccessException, InvocationTargetException {
722856
Method setLogLevels = LambdaLoggingAspect.class.getDeclaredMethod("setLogLevels", Level.class);

0 commit comments

Comments
 (0)