create thread local instrumentation context only when needed

Author: laurit

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

@@ -18,10 +18,6 @@ class SqlStatementSanitizerUtil {
   private static final SqlStatementSanitizer sanitizer = SqlStatementSanitizer.create(true);
 
   static SqlStatementInfo sanitize(String queryText) {
-    if (!InstrumenterContext.isActive()) {
-      return sanitizer.sanitize(queryText);
-    }
-
     Map<String, SqlStatementInfo> map =
         InstrumenterContext.computeIfAbsent("sanitized-sql-map", unused -> new HashMap<>());
     return map.computeIfAbsent(queryText, sanitizer::sanitize);

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

@@ -165,7 +165,11 @@ Context startAndEnd(
   }
 
   private Context doStart(Context parentContext, REQUEST request, @Nullable Instant startTime) {
-    return InstrumenterContext.withContext(() -> doStartImpl(parentContext, request, startTime));
+    try {
+      return doStartImpl(parentContext, request, startTime);
+    } finally {
+      InstrumenterContext.reset();
+    }
   }
 
   private Context doStartImpl(Context parentContext, REQUEST request, @Nullable Instant startTime) {

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

@@ -11,7 +11,6 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.function.Function;
-import java.util.function.Supplier;
 
 /**
  * Helper class for sharing computed values between different {@link AttributesExtractor}s and
@@ -21,45 +20,24 @@
  * at any time.
  */
 public final class InstrumenterContext {
-  private static final ThreadLocal<InstrumenterContext> instrumenterContext = new ThreadLocal<>();
+  private static final ThreadLocal<InstrumenterContext> instrumenterContext =
+      ThreadLocal.withInitial(InstrumenterContext::new);
 
   private final Map<String, Object> map = new HashMap<>();
-  private int useCount;
 
   private InstrumenterContext() {}
 
   @SuppressWarnings("unchecked")
   public static <T> T computeIfAbsent(String key, Function<String, T> function) {
-    InstrumenterContext context = instrumenterContext.get();
-    if (context == null) {
-      return function.apply(key);
-    }
-    return (T) context.map.computeIfAbsent(key, function);
+    return (T) get().computeIfAbsent(key, function);
   }
 
   // visible for testing
   static Map<String, Object> get() {
     return instrumenterContext.get().map;
   }
 
-  public static boolean isActive() {
-    return instrumenterContext.get() != null;
-  }
-
-  public static <T> T withContext(Supplier<T> action) {
-    InstrumenterContext context = instrumenterContext.get();
-    if (context == null) {
-      context = new InstrumenterContext();
-      instrumenterContext.set(context);
-    }
-    context.useCount++;
-    try {
-      return action.get();
-    } finally {
-      context.useCount--;
-      if (context.useCount == 0) {
-        instrumenterContext.remove();
-      }
-    }
+  public static void reset() {
+    instrumenterContext.remove();
   }
 }

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

@@ -17,13 +17,16 @@
 import io.opentelemetry.instrumentation.api.incubator.semconv.db.SqlStatementInfo;
 import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
+import io.opentelemetry.instrumentation.testing.internal.AutoCleanupExtension;
 import io.opentelemetry.semconv.incubating.DbIncubatingAttributes;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Map;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 class InstrumenterContextTest {
+  @RegisterExtension static final AutoCleanupExtension cleanup = AutoCleanupExtension.create();
 
   @SuppressWarnings({"unchecked", "deprecation"}) // using deprecated DB_SQL_TABLE
   @Test
@@ -41,39 +44,35 @@ public Collection<String> getRawQueryTexts(Object request) {
     AttributesExtractor<Object, Void> attributesExtractor =
         SqlClientAttributesExtractor.create(getter);
 
-    assertThat(InstrumenterContext.isActive()).isFalse();
-    InstrumenterContext.withContext(
-        () -> {
-          assertThat(InstrumenterContext.isActive()).isTrue();
-          assertThat(InstrumenterContext.get()).isEmpty();
-          assertThat(spanNameExtractor.extract(null)).isEqualTo("SELECT test");
-          // verify that sanitized statement was cached, see SqlStatementSanitizerUtil
-          assertThat(InstrumenterContext.get()).containsKey("sanitized-sql-map");
-          Map<String, SqlStatementInfo> sanitizedMap =
-              (Map<String, SqlStatementInfo>) InstrumenterContext.get().get("sanitized-sql-map");
-          assertThat(sanitizedMap).containsKey(testQuery);
+    InstrumenterContext.reset();
+    cleanup.deferCleanup(InstrumenterContext::reset);
 
-          // replace cached sanitization result to verify it is used
-          sanitizedMap.put(
-              testQuery,
-              SqlStatementInfo.create("SELECT name2 FROM test2 WHERE id = ?", "SELECT", "test2"));
-          {
-            AttributesBuilder builder = Attributes.builder();
-            attributesExtractor.onStart(builder, Context.root(), null);
-            assertThat(builder.build().get(maybeStable(DbIncubatingAttributes.DB_SQL_TABLE)))
-                .isEqualTo("test2");
-          }
+    assertThat(InstrumenterContext.get()).isEmpty();
+    assertThat(spanNameExtractor.extract(null)).isEqualTo("SELECT test");
+    // verify that sanitized statement was cached, see SqlStatementSanitizerUtil
+    assertThat(InstrumenterContext.get()).containsKey("sanitized-sql-map");
+    Map<String, SqlStatementInfo> sanitizedMap =
+        (Map<String, SqlStatementInfo>) InstrumenterContext.get().get("sanitized-sql-map");
+    assertThat(sanitizedMap).containsKey(testQuery);
 
-          // clear cached value to see whether it gets recomputed correctly
-          sanitizedMap.clear();
-          {
-            AttributesBuilder builder = Attributes.builder();
-            attributesExtractor.onStart(builder, Context.root(), null);
-            assertThat(builder.build().get(maybeStable(DbIncubatingAttributes.DB_SQL_TABLE)))
-                .isEqualTo("test");
-          }
+    // replace cached sanitization result to verify it is used
+    sanitizedMap.put(
+        testQuery,
+        SqlStatementInfo.create("SELECT name2 FROM test2 WHERE id = ?", "SELECT", "test2"));
+    {
+      AttributesBuilder builder = Attributes.builder();
+      attributesExtractor.onStart(builder, Context.root(), null);
+      assertThat(builder.build().get(maybeStable(DbIncubatingAttributes.DB_SQL_TABLE)))
+          .isEqualTo("test2");
+    }
 
-          return null;
-        });
+    // clear cached value to see whether it gets recomputed correctly
+    sanitizedMap.clear();
+    {
+      AttributesBuilder builder = Attributes.builder();
+      attributesExtractor.onStart(builder, Context.root(), null);
+      assertThat(builder.build().get(maybeStable(DbIncubatingAttributes.DB_SQL_TABLE)))
+          .isEqualTo("test");
+    }
   }
 }