Skip to content

Commit 095dd7c

Browse files
committed
create thread local instrumentation context only when needed
1 parent bc71a6b commit 095dd7c

File tree

4 files changed

+40
-63
lines changed

4 files changed

+40
-63
lines changed

instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/semconv/db/SqlStatementSanitizerUtil.java

-4
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,6 @@ class SqlStatementSanitizerUtil {
1818
private static final SqlStatementSanitizer sanitizer = SqlStatementSanitizer.create(true);
1919

2020
static SqlStatementInfo sanitize(String queryText) {
21-
if (!InstrumenterContext.isActive()) {
22-
return sanitizer.sanitize(queryText);
23-
}
24-
2521
Map<String, SqlStatementInfo> map =
2622
InstrumenterContext.computeIfAbsent("sanitized-sql-map", unused -> new HashMap<>());
2723
return map.computeIfAbsent(queryText, sanitizer::sanitize);

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/Instrumenter.java

+5-1
Original file line numberDiff line numberDiff line change
@@ -165,7 +165,11 @@ Context startAndEnd(
165165
}
166166

167167
private Context doStart(Context parentContext, REQUEST request, @Nullable Instant startTime) {
168-
return InstrumenterContext.withContext(() -> doStartImpl(parentContext, request, startTime));
168+
try {
169+
return doStartImpl(parentContext, request, startTime);
170+
} finally {
171+
InstrumenterContext.reset();
172+
}
169173
}
170174

171175
private Context doStartImpl(Context parentContext, REQUEST request, @Nullable Instant startTime) {

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/internal/InstrumenterContext.java

+5-27
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111
import java.util.HashMap;
1212
import java.util.Map;
1313
import java.util.function.Function;
14-
import java.util.function.Supplier;
1514

1615
/**
1716
* Helper class for sharing computed values between different {@link AttributesExtractor}s and
@@ -21,45 +20,24 @@
2120
* at any time.
2221
*/
2322
public final class InstrumenterContext {
24-
private static final ThreadLocal<InstrumenterContext> instrumenterContext = new ThreadLocal<>();
23+
private static final ThreadLocal<InstrumenterContext> instrumenterContext =
24+
ThreadLocal.withInitial(InstrumenterContext::new);
2525

2626
private final Map<String, Object> map = new HashMap<>();
27-
private int useCount;
2827

2928
private InstrumenterContext() {}
3029

3130
@SuppressWarnings("unchecked")
3231
public static <T> T computeIfAbsent(String key, Function<String, T> function) {
33-
InstrumenterContext context = instrumenterContext.get();
34-
if (context == null) {
35-
return function.apply(key);
36-
}
37-
return (T) context.map.computeIfAbsent(key, function);
32+
return (T) get().computeIfAbsent(key, function);
3833
}
3934

4035
// visible for testing
4136
static Map<String, Object> get() {
4237
return instrumenterContext.get().map;
4338
}
4439

45-
public static boolean isActive() {
46-
return instrumenterContext.get() != null;
47-
}
48-
49-
public static <T> T withContext(Supplier<T> action) {
50-
InstrumenterContext context = instrumenterContext.get();
51-
if (context == null) {
52-
context = new InstrumenterContext();
53-
instrumenterContext.set(context);
54-
}
55-
context.useCount++;
56-
try {
57-
return action.get();
58-
} finally {
59-
context.useCount--;
60-
if (context.useCount == 0) {
61-
instrumenterContext.remove();
62-
}
63-
}
40+
public static void reset() {
41+
instrumenterContext.remove();
6442
}
6543
}

instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/internal/InstrumenterContextTest.java

+30-31
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,16 @@
1717
import io.opentelemetry.instrumentation.api.incubator.semconv.db.SqlStatementInfo;
1818
import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
1919
import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
20+
import io.opentelemetry.instrumentation.testing.internal.AutoCleanupExtension;
2021
import io.opentelemetry.semconv.incubating.DbIncubatingAttributes;
2122
import java.util.Collection;
2223
import java.util.Collections;
2324
import java.util.Map;
2425
import org.junit.jupiter.api.Test;
26+
import org.junit.jupiter.api.extension.RegisterExtension;
2527

2628
class InstrumenterContextTest {
29+
@RegisterExtension static final AutoCleanupExtension cleanup = AutoCleanupExtension.create();
2730

2831
@SuppressWarnings({"unchecked", "deprecation"}) // using deprecated DB_SQL_TABLE
2932
@Test
@@ -41,39 +44,35 @@ public Collection<String> getRawQueryTexts(Object request) {
4144
AttributesExtractor<Object, Void> attributesExtractor =
4245
SqlClientAttributesExtractor.create(getter);
4346

44-
assertThat(InstrumenterContext.isActive()).isFalse();
45-
InstrumenterContext.withContext(
46-
() -> {
47-
assertThat(InstrumenterContext.isActive()).isTrue();
48-
assertThat(InstrumenterContext.get()).isEmpty();
49-
assertThat(spanNameExtractor.extract(null)).isEqualTo("SELECT test");
50-
// verify that sanitized statement was cached, see SqlStatementSanitizerUtil
51-
assertThat(InstrumenterContext.get()).containsKey("sanitized-sql-map");
52-
Map<String, SqlStatementInfo> sanitizedMap =
53-
(Map<String, SqlStatementInfo>) InstrumenterContext.get().get("sanitized-sql-map");
54-
assertThat(sanitizedMap).containsKey(testQuery);
47+
InstrumenterContext.reset();
48+
cleanup.deferCleanup(InstrumenterContext::reset);
5549

56-
// replace cached sanitization result to verify it is used
57-
sanitizedMap.put(
58-
testQuery,
59-
SqlStatementInfo.create("SELECT name2 FROM test2 WHERE id = ?", "SELECT", "test2"));
60-
{
61-
AttributesBuilder builder = Attributes.builder();
62-
attributesExtractor.onStart(builder, Context.root(), null);
63-
assertThat(builder.build().get(maybeStable(DbIncubatingAttributes.DB_SQL_TABLE)))
64-
.isEqualTo("test2");
65-
}
50+
assertThat(InstrumenterContext.get()).isEmpty();
51+
assertThat(spanNameExtractor.extract(null)).isEqualTo("SELECT test");
52+
// verify that sanitized statement was cached, see SqlStatementSanitizerUtil
53+
assertThat(InstrumenterContext.get()).containsKey("sanitized-sql-map");
54+
Map<String, SqlStatementInfo> sanitizedMap =
55+
(Map<String, SqlStatementInfo>) InstrumenterContext.get().get("sanitized-sql-map");
56+
assertThat(sanitizedMap).containsKey(testQuery);
6657

67-
// clear cached value to see whether it gets recomputed correctly
68-
sanitizedMap.clear();
69-
{
70-
AttributesBuilder builder = Attributes.builder();
71-
attributesExtractor.onStart(builder, Context.root(), null);
72-
assertThat(builder.build().get(maybeStable(DbIncubatingAttributes.DB_SQL_TABLE)))
73-
.isEqualTo("test");
74-
}
58+
// replace cached sanitization result to verify it is used
59+
sanitizedMap.put(
60+
testQuery,
61+
SqlStatementInfo.create("SELECT name2 FROM test2 WHERE id = ?", "SELECT", "test2"));
62+
{
63+
AttributesBuilder builder = Attributes.builder();
64+
attributesExtractor.onStart(builder, Context.root(), null);
65+
assertThat(builder.build().get(maybeStable(DbIncubatingAttributes.DB_SQL_TABLE)))
66+
.isEqualTo("test2");
67+
}
7568

76-
return null;
77-
});
69+
// clear cached value to see whether it gets recomputed correctly
70+
sanitizedMap.clear();
71+
{
72+
AttributesBuilder builder = Attributes.builder();
73+
attributesExtractor.onStart(builder, Context.root(), null);
74+
assertThat(builder.build().get(maybeStable(DbIncubatingAttributes.DB_SQL_TABLE)))
75+
.isEqualTo("test");
76+
}
7877
}
7978
}

0 commit comments

Comments
 (0)