Skip to content

Commit bb70c1b

Browse files
phipagphipag
authored
fix(metrics): Clear custom dimensions when flushing. (#2328)
1 parent 7ea89de commit bb70c1b

File tree

2 files changed

+212
-1
lines changed

2 files changed

+212
-1
lines changed

‎powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/internal/EmfMetricsLogger.java‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,27 @@ public void flush() {
164164
}
165165
} else {
166166
emfLogger.flush();
167+
168+
// Clear custom dimensions after flush while preserving default dimensions
169+
clearCustomDimensions();
170+
}
171+
}
172+
173+
private void clearCustomDimensions() {
174+
// Reset all dimensions in the EMF logger
175+
emfLogger.resetDimensions(false);
176+
177+
// Re-apply default dimensions if they exist
178+
if (!defaultDimensions.isEmpty()) {
179+
DimensionSet emfDimensionSet = new DimensionSet();
180+
defaultDimensions.forEach((key, value) -> {
181+
try {
182+
emfDimensionSet.addDimension(key, value);
183+
} catch (Exception e) {
184+
// Ignore dimension errors
185+
}
186+
});
187+
emfLogger.setDimensions(emfDimensionSet);
167188
}
168189
}
169190

@@ -198,7 +219,8 @@ public void flushMetrics(Consumer<Metrics> metricsConsumer) {
198219
metrics.setNamespace(this.namespace);
199220
}
200221
if (!defaultDimensions.isEmpty()) {
201-
metrics.setDefaultDimensions(software.amazon.lambda.powertools.metrics.model.DimensionSet.of(defaultDimensions));
222+
metrics.setDefaultDimensions(
223+
software.amazon.lambda.powertools.metrics.model.DimensionSet.of(defaultDimensions));
202224
}
203225
properties.forEach(metrics::addMetadata);
204226

‎powertools-metrics/src/test/java/software/amazon/lambda/powertools/metrics/internal/EmfMetricsLoggerTest.java‎

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -263,6 +263,65 @@ void shouldAddMetadata() throws Exception {
263263
assertThat(rootNode.get("CustomMetadata").asText()).isEqualTo("MetadataValue");
264264
}
265265

266+
@Test
267+
void shouldClearMetadataAfterFlush() throws Exception {
268+
// Given - Add metadata and flush first time
269+
metrics.addMetadata("RequestId", "req-123");
270+
metrics.addMetadata("UserAgent", "test-agent");
271+
metrics.addMetric("FirstMetric", 1.0);
272+
metrics.flush();
273+
274+
// Capture first flush output and reset for second flush
275+
String firstFlushOutput = outputStreamCaptor.toString().trim();
276+
outputStreamCaptor.reset();
277+
278+
// When - Add another metric and flush again using the SAME metrics instance
279+
metrics.addMetric("SecondMetric", 2.0);
280+
metrics.flush();
281+
282+
// Then - Verify first flush had metadata
283+
JsonNode firstRootNode = objectMapper.readTree(firstFlushOutput);
284+
assertThat(firstRootNode.has("RequestId")).isTrue();
285+
assertThat(firstRootNode.get("RequestId").asText()).isEqualTo("req-123");
286+
assertThat(firstRootNode.has("UserAgent")).isTrue();
287+
assertThat(firstRootNode.get("UserAgent").asText()).isEqualTo("test-agent");
288+
assertThat(firstRootNode.has("FirstMetric")).isTrue();
289+
290+
// Verify second flush does NOT have metadata from first flush
291+
// The EMF library automatically clears metadata after flush
292+
String secondFlushOutput = outputStreamCaptor.toString().trim();
293+
JsonNode secondRootNode = objectMapper.readTree(secondFlushOutput);
294+
295+
// Metadata should be cleared after first flush by the EMF library
296+
assertThat(secondRootNode.has("RequestId")).isFalse();
297+
assertThat(secondRootNode.has("UserAgent")).isFalse();
298+
assertThat(secondRootNode.has("SecondMetric")).isTrue();
299+
}
300+
301+
@Test
302+
void shouldInheritMetadataInFlushMetricsMethod() throws Exception {
303+
// Given - Add metadata to the main metrics instance
304+
metrics.addMetadata("PersistentMetadata", "should-inherit");
305+
metrics.addMetadata("GlobalContext", "main-instance");
306+
307+
// When - Use flushMetrics to create a separate metrics context
308+
metrics.flushMetrics(separateMetrics -> {
309+
separateMetrics.addMetric("SeparateMetric", 1.0);
310+
// Don't add any metadata to the separate instance
311+
});
312+
313+
// Then - The separate metrics context SHOULD inherit metadata from main instance
314+
String flushMetricsOutput = outputStreamCaptor.toString().trim();
315+
JsonNode rootNode = objectMapper.readTree(flushMetricsOutput);
316+
317+
// The separate metrics should have inherited metadata (this is expected behavior)
318+
assertThat(rootNode.has("PersistentMetadata")).isTrue();
319+
assertThat(rootNode.get("PersistentMetadata").asText()).isEqualTo("should-inherit");
320+
assertThat(rootNode.has("GlobalContext")).isTrue();
321+
assertThat(rootNode.get("GlobalContext").asText()).isEqualTo("main-instance");
322+
assertThat(rootNode.has("SeparateMetric")).isTrue();
323+
}
324+
266325
@Test
267326
void shouldSetDefaultDimensions() throws Exception {
268327
// Given
@@ -547,4 +606,134 @@ void shouldNotFlushSingleMetricWhenDisabled() {
547606
String emfOutput = outputStreamCaptor.toString().trim();
548607
assertThat(emfOutput).isEmpty();
549608
}
609+
610+
@Test
611+
void shouldClearCustomDimensionsAfterFlush() throws Exception {
612+
// Given - Set up default dimensions that should persist
613+
DimensionSet defaultDimensions = DimensionSet.of("Service", "TestService", "Environment", "Test");
614+
metrics.setDefaultDimensions(defaultDimensions);
615+
616+
// First invocation - add custom dimensions and flush
617+
DimensionSet customDimensions = DimensionSet.of("EXAMPLE_KEY", "EXAMPLE_VALUE");
618+
metrics.addDimension(customDimensions);
619+
metrics.addMetric("SERL", 1.0);
620+
metrics.flush();
621+
622+
// Capture first flush output
623+
String firstFlushOutput = outputStreamCaptor.toString().trim();
624+
outputStreamCaptor.reset(); // Clear for second flush
625+
626+
// Second invocation - should NOT have custom dimensions from first invocation
627+
metrics.addMetric("Expected", 1.0);
628+
metrics.flush();
629+
630+
// Then - Verify first flush had both default and custom dimensions
631+
JsonNode firstRootNode = objectMapper.readTree(firstFlushOutput);
632+
assertThat(firstRootNode.has("Service")).isTrue();
633+
assertThat(firstRootNode.get("Service").asText()).isEqualTo("TestService");
634+
assertThat(firstRootNode.has("Environment")).isTrue();
635+
assertThat(firstRootNode.get("Environment").asText()).isEqualTo("Test");
636+
assertThat(firstRootNode.has("EXAMPLE_KEY")).isTrue();
637+
assertThat(firstRootNode.get("EXAMPLE_KEY").asText()).isEqualTo("EXAMPLE_VALUE");
638+
assertThat(firstRootNode.has("SERL")).isTrue();
639+
640+
// Verify second flush has ONLY default dimensions (custom dimensions should be cleared)
641+
String secondFlushOutput = outputStreamCaptor.toString().trim();
642+
JsonNode secondRootNode = objectMapper.readTree(secondFlushOutput);
643+
644+
// Default dimensions should still be present
645+
assertThat(secondRootNode.has("Service")).isTrue();
646+
assertThat(secondRootNode.get("Service").asText()).isEqualTo("TestService");
647+
assertThat(secondRootNode.has("Environment")).isTrue();
648+
assertThat(secondRootNode.get("Environment").asText()).isEqualTo("Test");
649+
650+
// Custom dimensions should be cleared (this is the failing assertion that demonstrates the bug)
651+
assertThat(secondRootNode.has("EXAMPLE_KEY")).isFalse();
652+
assertThat(secondRootNode.has("Expected")).isTrue();
653+
654+
// Verify dimensions in CloudWatchMetrics section
655+
JsonNode secondDimensions = secondRootNode.get("_aws").get("CloudWatchMetrics").get(0).get("Dimensions").get(0);
656+
boolean hasExampleKey = false;
657+
boolean hasService = false;
658+
boolean hasEnvironment = false;
659+
660+
for (JsonNode dimension : secondDimensions) {
661+
String dimName = dimension.asText();
662+
if ("EXAMPLE_KEY".equals(dimName)) {
663+
hasExampleKey = true;
664+
} else if ("Service".equals(dimName)) {
665+
hasService = true;
666+
} else if ("Environment".equals(dimName)) {
667+
hasEnvironment = true;
668+
}
669+
}
670+
671+
// Default dimensions should be in CloudWatchMetrics
672+
assertThat(hasService).isTrue();
673+
assertThat(hasEnvironment).isTrue();
674+
// Custom dimension should NOT be in CloudWatchMetrics (this should fail initially)
675+
assertThat(hasExampleKey).isFalse();
676+
}
677+
678+
@Test
679+
void shouldHandleEmptyCustomDimensionsGracefully() throws Exception {
680+
// Given - Only default dimensions, no custom dimensions
681+
metrics.setDefaultDimensions(DimensionSet.of("Service", "TestService"));
682+
683+
// When - Flush without adding custom dimensions
684+
metrics.addMetric("TestMetric", 1.0);
685+
metrics.flush();
686+
outputStreamCaptor.reset();
687+
688+
// Second flush
689+
metrics.addMetric("TestMetric2", 2.0);
690+
metrics.flush();
691+
692+
// Then - Should work normally with only default dimensions
693+
String output = outputStreamCaptor.toString().trim();
694+
JsonNode rootNode = objectMapper.readTree(output);
695+
696+
assertThat(rootNode.has("Service")).isTrue();
697+
assertThat(rootNode.get("Service").asText()).isEqualTo("TestService");
698+
assertThat(rootNode.has("TestMetric2")).isTrue();
699+
}
700+
701+
@Test
702+
void shouldClearCustomDimensionsWhenNoDefaultDimensionsSet() throws Exception {
703+
// Given - No default dimensions set
704+
metrics.clearDefaultDimensions();
705+
706+
// When - Add custom dimensions and flush
707+
metrics.addDimension("CustomDim", "CustomValue");
708+
metrics.addMetric("Metric1", 1.0);
709+
metrics.flush();
710+
outputStreamCaptor.reset();
711+
712+
// Second flush without custom dimensions
713+
metrics.addMetric("Metric2", 2.0);
714+
metrics.flush();
715+
716+
// Then - Custom dimensions should be cleared
717+
String output = outputStreamCaptor.toString().trim();
718+
JsonNode rootNode = objectMapper.readTree(output);
719+
720+
assertThat(rootNode.has("CustomDim")).isFalse();
721+
assertThat(rootNode.has("Metric2")).isTrue();
722+
723+
// Verify no custom dimensions in CloudWatchMetrics section
724+
JsonNode dimensionsArray = rootNode.get("_aws").get("CloudWatchMetrics").get(0).get("Dimensions");
725+
boolean hasCustomDim = false;
726+
if (dimensionsArray != null && dimensionsArray.size() > 0) {
727+
JsonNode dimensions = dimensionsArray.get(0);
728+
if (dimensions != null) {
729+
for (JsonNode dimension : dimensions) {
730+
if ("CustomDim".equals(dimension.asText())) {
731+
hasCustomDim = true;
732+
break;
733+
}
734+
}
735+
}
736+
}
737+
assertThat(hasCustomDim).isFalse();
738+
}
550739
}

0 commit comments

Comments
 (0)