Skip to content

Commit ed6a69b

Browse files
phipagphipag
committed
Add proxy to manage thread local metrics provider instances to achieve zero-lock thread-safety.
1 parent 4e83d8a commit ed6a69b

File tree

9 files changed

+525
-27
lines changed

9 files changed

+525
-27
lines changed

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

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,11 @@
1616

1717
import org.crac.Core;
1818
import org.crac.Resource;
19+
1920
import software.amazon.lambda.powertools.common.internal.ClassPreLoader;
2021
import software.amazon.lambda.powertools.common.internal.LambdaConstants;
2122
import software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor;
23+
import software.amazon.lambda.powertools.metrics.internal.ThreadLocalMetricsProxy;
2224
import software.amazon.lambda.powertools.metrics.model.DimensionSet;
2325
import software.amazon.lambda.powertools.metrics.provider.EmfMetricsProvider;
2426
import software.amazon.lambda.powertools.metrics.provider.MetricsProvider;
@@ -28,7 +30,7 @@
2830
*/
2931
public final class MetricsFactory implements Resource {
3032
private static MetricsProvider provider = new EmfMetricsProvider();
31-
private static Metrics metrics;
33+
private static ThreadLocalMetricsProxy metricsProxy;
3234

3335
// Dummy instance to register MetricsFactory with CRaC
3436
private static final MetricsFactory INSTANCE = new MetricsFactory();
@@ -44,23 +46,23 @@ public final class MetricsFactory implements Resource {
4446
* @return the singleton Metrics instance
4547
*/
4648
public static synchronized Metrics getMetricsInstance() {
47-
if (metrics == null) {
48-
metrics = provider.getMetricsInstance();
49+
if (metricsProxy == null) {
50+
metricsProxy = new ThreadLocalMetricsProxy(provider);
4951

5052
// Apply default configuration from environment variables
5153
String envNamespace = System.getenv("POWERTOOLS_METRICS_NAMESPACE");
5254
if (envNamespace != null) {
53-
metrics.setNamespace(envNamespace);
55+
metricsProxy.setNamespace(envNamespace);
5456
}
5557

5658
// Only set Service dimension if it's not the default undefined value
5759
String serviceName = LambdaHandlerProcessor.serviceName();
5860
if (!LambdaConstants.SERVICE_UNDEFINED.equals(serviceName)) {
59-
metrics.setDefaultDimensions(DimensionSet.of("Service", serviceName));
61+
metricsProxy.setDefaultDimensions(DimensionSet.of("Service", serviceName));
6062
}
6163
}
6264

63-
return metrics;
65+
return metricsProxy;
6466
}
6567

6668
/**
@@ -73,8 +75,8 @@ public static synchronized void setMetricsProvider(MetricsProvider metricsProvid
7375
throw new IllegalArgumentException("Metrics provider cannot be null");
7476
}
7577
provider = metricsProvider;
76-
// Reset the metrics instance so it will be recreated with the new provider
77-
metrics = null;
78+
// Reset the metrics proxy so it will be recreated with the new provider
79+
metricsProxy = null;
7880
}
7981

8082
@Override

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

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
/*
2+
* Copyright 2023 Amazon.com, Inc. or its affiliates.
3+
* Licensed under the Apache License, Version 2.0 (the
4+
* "License"); you may not use this file except in compliance
5+
* with the License. You may obtain a copy of the License at
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
* Unless required by applicable law or agreed to in writing, software
8+
* distributed under the License is distributed on an "AS IS" BASIS,
9+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
* See the License for the specific language governing permissions and
11+
* limitations under the License.
12+
*
13+
*/
14+
15+
package software.amazon.lambda.powertools.metrics.internal;
16+
17+
import java.time.Instant;
18+
import java.util.HashMap;
19+
import java.util.Optional;
20+
import java.util.concurrent.atomic.AtomicBoolean;
21+
import java.util.concurrent.atomic.AtomicReference;
22+
import java.util.function.Consumer;
23+
24+
import com.amazonaws.services.lambda.runtime.Context;
25+
26+
import software.amazon.lambda.powertools.metrics.Metrics;
27+
import software.amazon.lambda.powertools.metrics.model.DimensionSet;
28+
import software.amazon.lambda.powertools.metrics.model.MetricResolution;
29+
import software.amazon.lambda.powertools.metrics.model.MetricUnit;
30+
import software.amazon.lambda.powertools.metrics.provider.MetricsProvider;
31+
32+
public class ThreadLocalMetricsProxy implements Metrics {
33+
private final InheritableThreadLocal<Metrics> threadLocalMetrics = new InheritableThreadLocal<>();
34+
private final MetricsProvider provider;
35+
private final AtomicReference<String> initialNamespace = new AtomicReference<>();
36+
private final AtomicReference<DimensionSet> initialDefaultDimensions = new AtomicReference<>();
37+
private final AtomicBoolean initialRaiseOnEmptyMetrics = new AtomicBoolean(false);
38+
39+
public ThreadLocalMetricsProxy(MetricsProvider provider) {
40+
this.provider = provider;
41+
}
42+
43+
private Metrics getOrCreateThreadLocalMetrics() {
44+
Metrics metrics = threadLocalMetrics.get();
45+
if (metrics == null) {
46+
metrics = provider.getMetricsInstance();
47+
String namespace = initialNamespace.get();
48+
if (namespace != null) {
49+
metrics.setNamespace(namespace);
50+
}
51+
DimensionSet dimensions = initialDefaultDimensions.get();
52+
if (dimensions != null) {
53+
metrics.setDefaultDimensions(dimensions);
54+
}
55+
metrics.setRaiseOnEmptyMetrics(initialRaiseOnEmptyMetrics.get());
56+
threadLocalMetrics.set(metrics);
57+
}
58+
return metrics;
59+
}
60+
61+
// Configuration methods - called by MetricsFactory and MetricsBuilder
62+
// These methods DO NOT eagerly create thread-local instances because they are typically called
63+
// outside the Lambda handler (e.g., during class initialization) potentially on a different thread.
64+
// We delay instance creation until the first operation that needs the metrics backend (e.g., addMetric).
65+
// See {@link software.amazon.lambda.powertools.metrics.MetricsFactory#getMetricsInstance()}
66+
// and {@link software.amazon.lambda.powertools.metrics.MetricsBuilder#build()}
67+
68+
@Override
69+
public void setNamespace(String namespace) {
70+
this.initialNamespace.set(namespace);
71+
Optional.ofNullable(threadLocalMetrics.get()).ifPresent(m -> m.setNamespace(namespace));
72+
}
73+
74+
@Override
75+
public void setDefaultDimensions(DimensionSet dimensionSet) {
76+
if (dimensionSet == null) {
77+
throw new IllegalArgumentException("DimensionSet cannot be null");
78+
}
79+
this.initialDefaultDimensions.set(dimensionSet);
80+
Optional.ofNullable(threadLocalMetrics.get()).ifPresent(m -> m.setDefaultDimensions(dimensionSet));
81+
}
82+
83+
@Override
84+
public void setRaiseOnEmptyMetrics(boolean raiseOnEmptyMetrics) {
85+
this.initialRaiseOnEmptyMetrics.set(raiseOnEmptyMetrics);
86+
Optional.ofNullable(threadLocalMetrics.get()).ifPresent(m -> m.setRaiseOnEmptyMetrics(raiseOnEmptyMetrics));
87+
}
88+
89+
@Override
90+
public DimensionSet getDefaultDimensions() {
91+
Metrics metrics = threadLocalMetrics.get();
92+
if (metrics != null) {
93+
return metrics.getDefaultDimensions();
94+
}
95+
DimensionSet dimensions = initialDefaultDimensions.get();
96+
return dimensions != null ? dimensions : DimensionSet.of(new HashMap<>());
97+
}
98+
99+
// Metrics operations - these eagerly create thread-local instances
100+
101+
@Override
102+
public void addMetric(String key, double value, MetricUnit unit, MetricResolution resolution) {
103+
getOrCreateThreadLocalMetrics().addMetric(key, value, unit, resolution);
104+
}
105+
106+
@Override
107+
public void addDimension(DimensionSet dimensionSet) {
108+
getOrCreateThreadLocalMetrics().addDimension(dimensionSet);
109+
}
110+
111+
@Override
112+
public void setTimestamp(Instant timestamp) {
113+
getOrCreateThreadLocalMetrics().setTimestamp(timestamp);
114+
}
115+
116+
@Override
117+
public void addMetadata(String key, Object value) {
118+
getOrCreateThreadLocalMetrics().addMetadata(key, value);
119+
}
120+
121+
@Override
122+
public void clearDefaultDimensions() {
123+
getOrCreateThreadLocalMetrics().clearDefaultDimensions();
124+
}
125+
126+
@Override
127+
public void flush() {
128+
// Always create instance to ensure validation and warnings are triggered. E.g. when raiseOnEmptyMetrics
129+
// is enabled.
130+
Metrics metrics = getOrCreateThreadLocalMetrics();
131+
metrics.flush();
132+
threadLocalMetrics.remove();
133+
}
134+
135+
@Override
136+
public void captureColdStartMetric(Context context, DimensionSet dimensions) {
137+
getOrCreateThreadLocalMetrics().captureColdStartMetric(context, dimensions);
138+
}
139+
140+
@Override
141+
public void captureColdStartMetric(DimensionSet dimensions) {
142+
getOrCreateThreadLocalMetrics().captureColdStartMetric(dimensions);
143+
}
144+
145+
@Override
146+
public void flushMetrics(Consumer<Metrics> metricsConsumer) {
147+
getOrCreateThreadLocalMetrics().flushMetrics(metricsConsumer);
148+
}
149+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ void tearDown() throws Exception {
6868
System.setOut(standardOut);
6969

7070
// Reset the singleton state between tests
71-
java.lang.reflect.Field field = MetricsFactory.class.getDeclaredField("metrics");
71+
java.lang.reflect.Field field = MetricsFactory.class.getDeclaredField("metricsProxy");
7272
field.setAccessible(true);
7373
field.set(null, null);
7474

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

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,10 +27,10 @@
2727
import com.fasterxml.jackson.databind.JsonNode;
2828
import com.fasterxml.jackson.databind.ObjectMapper;
2929

30+
import software.amazon.lambda.powertools.metrics.internal.ThreadLocalMetricsProxy;
3031
import software.amazon.lambda.powertools.metrics.model.DimensionSet;
3132
import software.amazon.lambda.powertools.metrics.model.MetricUnit;
3233
import software.amazon.lambda.powertools.metrics.provider.MetricsProvider;
33-
import software.amazon.lambda.powertools.metrics.testutils.TestMetrics;
3434
import software.amazon.lambda.powertools.metrics.testutils.TestMetricsProvider;
3535

3636
class MetricsBuilderTest {
@@ -49,7 +49,7 @@ void tearDown() throws Exception {
4949
System.setOut(standardOut);
5050

5151
// Reset the singleton state between tests
52-
java.lang.reflect.Field field = MetricsFactory.class.getDeclaredField("metrics");
52+
java.lang.reflect.Field field = MetricsFactory.class.getDeclaredField("metricsProxy");
5353
field.setAccessible(true);
5454
field.set(null, null);
5555

@@ -151,7 +151,7 @@ void shouldBuildWithMultipleDefaultDimensions() throws Exception {
151151
}
152152

153153
@Test
154-
void shouldBuildWithCustomMetricsProvider() {
154+
void shouldBuildWithCustomMetricsProvider() throws Exception {
155155
// Given
156156
MetricsProvider testProvider = new TestMetricsProvider();
157157

@@ -161,7 +161,13 @@ void shouldBuildWithCustomMetricsProvider() {
161161
.build();
162162

163163
// Then
164-
assertThat(metrics).isInstanceOf(TestMetrics.class);
164+
assertThat(metrics)
165+
.isInstanceOf(ThreadLocalMetricsProxy.class);
166+
167+
java.lang.reflect.Field providerField = metrics.getClass().getDeclaredField("provider");
168+
providerField.setAccessible(true);
169+
MetricsProvider actualProvider = (MetricsProvider) providerField.get(metrics);
170+
assertThat(actualProvider).isSameAs(testProvider);
165171
}
166172

167173
@Test

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

Lines changed: 82 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import java.io.ByteArrayOutputStream;
2121
import java.io.PrintStream;
2222
import java.lang.reflect.Method;
23+
import java.util.concurrent.CountDownLatch;
2324

2425
import org.junit.jupiter.api.AfterEach;
2526
import org.junit.jupiter.api.BeforeEach;
@@ -30,9 +31,9 @@
3031
import com.fasterxml.jackson.databind.ObjectMapper;
3132

3233
import software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor;
34+
import software.amazon.lambda.powertools.metrics.internal.ThreadLocalMetricsProxy;
3335
import software.amazon.lambda.powertools.metrics.model.MetricUnit;
3436
import software.amazon.lambda.powertools.metrics.provider.MetricsProvider;
35-
import software.amazon.lambda.powertools.metrics.testutils.TestMetrics;
3637
import software.amazon.lambda.powertools.metrics.testutils.TestMetricsProvider;
3738

3839
class MetricsFactoryTest {
@@ -64,7 +65,7 @@ void tearDown() throws Exception {
6465
System.setOut(standardOut);
6566

6667
// Reset the singleton state between tests
67-
java.lang.reflect.Field field = MetricsFactory.class.getDeclaredField("metrics");
68+
java.lang.reflect.Field field = MetricsFactory.class.getDeclaredField("metricsProxy");
6869
field.setAccessible(true);
6970
field.set(null, null);
7071

@@ -126,7 +127,7 @@ void shouldUseServiceNameFromEnvironmentVariable() throws Exception {
126127
}
127128

128129
@Test
129-
void shouldSetCustomMetricsProvider() {
130+
void shouldSetCustomMetricsProvider() throws Exception {
130131
// Given
131132
MetricsProvider testProvider = new TestMetricsProvider();
132133

@@ -135,7 +136,13 @@ void shouldSetCustomMetricsProvider() {
135136
Metrics metrics = MetricsFactory.getMetricsInstance();
136137

137138
// Then
138-
assertThat(metrics).isInstanceOf(TestMetrics.class);
139+
assertThat(metrics)
140+
.isInstanceOf(ThreadLocalMetricsProxy.class);
141+
142+
java.lang.reflect.Field providerField = metrics.getClass().getDeclaredField("provider");
143+
providerField.setAccessible(true);
144+
MetricsProvider actualProvider = (MetricsProvider) providerField.get(metrics);
145+
assertThat(actualProvider).isSameAs(testProvider);
139146
}
140147

141148
@Test
@@ -163,4 +170,75 @@ void shouldNotSetServiceDimensionWhenServiceUndefined() throws Exception {
163170
// Service dimension should not be present
164171
assertThat(rootNode.has("Service")).isFalse();
165172
}
173+
174+
@Test
175+
void concurrentInvocations_shouldIsolateDimensions() throws Exception {
176+
// GIVEN - Simulate real Lambda scenario: Metrics instance created outside handler
177+
Metrics metrics = MetricsFactory.getMetricsInstance();
178+
179+
CountDownLatch latch = new CountDownLatch(2);
180+
Exception[] exceptions = new Exception[2];
181+
182+
Thread thread1 = new Thread(() -> {
183+
try {
184+
latch.countDown();
185+
latch.await();
186+
187+
// Simulate handleRequest execution
188+
metrics.setNamespace("TestNamespace");
189+
metrics.addDimension("userId", "user123");
190+
metrics.addMetric("ProcessedOrder", 1, MetricUnit.COUNT);
191+
metrics.flush();
192+
} catch (Exception e) {
193+
exceptions[0] = e;
194+
}
195+
});
196+
197+
Thread thread2 = new Thread(() -> {
198+
try {
199+
latch.countDown();
200+
latch.await();
201+
202+
// Simulate handleRequest execution
203+
metrics.setNamespace("TestNamespace");
204+
metrics.addDimension("userId", "user456");
205+
metrics.addMetric("ProcessedOrder", 1, MetricUnit.COUNT);
206+
metrics.flush();
207+
} catch (Exception e) {
208+
exceptions[1] = e;
209+
}
210+
});
211+
212+
// WHEN
213+
thread1.start();
214+
thread2.start();
215+
thread1.join();
216+
thread2.join();
217+
218+
// THEN
219+
assertThat(exceptions[0]).isNull();
220+
assertThat(exceptions[1]).isNull();
221+
222+
String emfOutput = outputStreamCaptor.toString().trim();
223+
String[] jsonLines = emfOutput.split("\n");
224+
assertThat(jsonLines).hasSize(2);
225+
226+
JsonNode output1 = objectMapper.readTree(jsonLines[0]);
227+
JsonNode output2 = objectMapper.readTree(jsonLines[1]);
228+
229+
// Check if dimensions are leaking across threads
230+
boolean hasLeakage = false;
231+
232+
if (output1.has("userId") && output2.has("userId")) {
233+
String userId1 = output1.get("userId").asText();
234+
String userId2 = output2.get("userId").asText();
235+
// Both should have different userIds
236+
hasLeakage = userId1.equals(userId2);
237+
}
238+
239+
// Each thread should have isolated dimensions
240+
assertThat(hasLeakage)
241+
.as("Dimensions should NOT leak across threads - each thread should have its own userId")
242+
.isFalse();
243+
}
166244
}

0 commit comments

Comments
 (0)