Make clickhouse instrumentation indy compatible (#12181)

open-telemetry · Sep 13, 2024 · 0b04b87 · 0b04b87
1 parent e8c5c06
commit 0b04b87
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 32 deletions.
diff --git a/...o/opentelemetry/javaagent/instrumentation/clickhouse/ClickHouseClientInstrumentation.java b/...o/opentelemetry/javaagent/instrumentation/clickhouse/ClickHouseClientInstrumentation.java
@@ -7,7 +7,6 @@
 
 import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
-import static io.opentelemetry.javaagent.instrumentation.clickhouse.ClickHouseSingletons.instrumenter;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
@@ -17,7 +16,6 @@
 import com.clickhouse.client.ClickHouseRequest;
 import com.clickhouse.client.config.ClickHouseDefaults;
 import io.opentelemetry.context.Context;
-import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.bootstrap.CallDepth;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
@@ -43,19 +41,12 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class ClickHouseExecuteAndWaitAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void onEnter(
-        @Advice.Argument(0) ClickHouseRequest<?> clickHouseRequest,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope,
-        @Advice.Local("otelCallDepth") CallDepth callDepth) {
+    public static ClickHouseScope onEnter(
+        @Advice.Argument(0) ClickHouseRequest<?> clickHouseRequest) {
 
-      callDepth = CallDepth.forClass(ClickHouseClient.class);
-      if (callDepth.getAndIncrement() > 0) {
-        return;
-      }
-
-      if (clickHouseRequest == null) {
-        return;
+      CallDepth callDepth = CallDepth.forClass(ClickHouseClient.class);
+      if (callDepth.getAndIncrement() > 0 || clickHouseRequest == null) {
+        return null;
       }
 
       Context parentContext = currentContext();
@@ -70,32 +61,19 @@ public static void onEnter(
                   .orElse(ClickHouseDefaults.DATABASE.getDefaultValue().toString()),
               clickHouseRequest.getPreparedQuery().getOriginalQuery());
 
-      if (!instrumenter().shouldStart(parentContext, request)) {
-        return;
-      }
-
-      context = instrumenter().start(parentContext, request);
-      scope = context.makeCurrent();
+      return ClickHouseScope.start(parentContext, request);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void onExit(
-        @Advice.Thrown Throwable throwable,
-        @Advice.Local("otelRequest") ClickHouseDbRequest clickHouseRequest,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope,
-        @Advice.Local("otelCallDepth") CallDepth callDepth) {
-
-      if (callDepth.decrementAndGet() > 0) {
-        return;
-      }
+        @Advice.Thrown Throwable throwable, @Advice.Enter ClickHouseScope scope) {
 
-      if (scope == null) {
+      CallDepth callDepth = CallDepth.forClass(ClickHouseClient.class);
+      if (callDepth.decrementAndGet() > 0 || scope == null) {
         return;
       }
 
-      scope.close();
-      instrumenter().end(context, clickHouseRequest, null, throwable);
+      scope.end(throwable);
     }
   }
 }
diff --git a/.../src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/ClickHouseScope.java b/.../src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/ClickHouseScope.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.clickhouse;
+
+import static io.opentelemetry.javaagent.instrumentation.clickhouse.ClickHouseSingletons.instrumenter;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+
+/** Container used to carry state between enter and exit advices */
+public final class ClickHouseScope {
+  private final ClickHouseDbRequest clickHouseDbRequest;
+  private final Context context;
+  private final Scope scope;
+
+  private ClickHouseScope(ClickHouseDbRequest clickHouseDbRequest, Context context, Scope scope) {
+    this.clickHouseDbRequest = clickHouseDbRequest;
+    this.context = context;
+    this.scope = scope;
+  }
+
+  public static ClickHouseScope start(
+      Context parentContext, ClickHouseDbRequest clickHouseDbRequest) {
+    if (!instrumenter().shouldStart(parentContext, clickHouseDbRequest)) {
+      return null;
+    }
+
+    Context context = instrumenter().start(parentContext, clickHouseDbRequest);
+    return new ClickHouseScope(clickHouseDbRequest, context, context.makeCurrent());
+  }
+
+  public void end(Throwable throwable) {
+    scope.close();
+    instrumenter().end(context, clickHouseDbRequest, null, throwable);
+  }
+}