run callbacks in the context of the parent span

open-telemetry · May 6, 2024 · 389a438 · 389a438
1 parent fbba347
commit 389a438
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 8 deletions.
diff --git a/...io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbImplInstrumentation.java b/...io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbImplInstrumentation.java
@@ -24,6 +24,7 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
 import net.bytebuddy.matcher.ElementMatcher;
 import okhttp3.HttpUrl;
 import org.influxdb.dto.BatchPoints;
@@ -74,6 +75,7 @@ public static class InfluxDbQueryAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void onEnter(
         @Advice.Argument(0) Query query,
+        @Advice.AllArguments(readOnly = false, typing = Assigner.Typing.DYNAMIC) Object[] arguments,
         @Advice.FieldValue(value = "retrofit") Retrofit retrofit,
         @Advice.Local("otelCallDepth") CallDepth callDepth,
         @Advice.Local("otelRequest") InfluxDbRequest influxDbRequest,
@@ -98,6 +100,17 @@ public static void onEnter(
         return;
       }
 
+      // wrap callbacks so they'd run in the context of the parent span
+      Object[] newArguments = new Object[arguments.length];
+      boolean hasChangedArgument = false;
+      for (int i = 0; i < arguments.length; i++) {
+        newArguments[i] = InfluxDbObjetWrapper.wrap(arguments[i], parentContext);
+        hasChangedArgument |= newArguments[i] != arguments[i];
+      }
+      if (hasChangedArgument) {
+        arguments = newArguments;
+      }
+
       context = instrumenter().start(parentContext, influxDbRequest);
       scope = context.makeCurrent();
     }

diff --git a/...est/java/io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbClientTest.java b/...est/java/io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbClientTest.java
@@ -211,12 +211,8 @@ void testQueryWithFiveArguments() throws InterruptedException {
           influxDb.query(
               query,
               10,
-              (cancellable, queryResult) -> {
-                countDownLatch.countDown();
-              },
-              () -> {
-                testing.runWithSpan("child", () -> {});
-              },
+              (cancellable, queryResult) -> countDownLatch.countDown(),
+              () -> testing.runWithSpan("child", () -> {}),
               throwable -> {});
         });
     assertThat(countDownLatch.await(10, TimeUnit.SECONDS)).isTrue();
@@ -235,7 +231,7 @@ void testQueryWithFiveArguments() throws InterruptedException {
                                 "SELECT",
                                 databaseName)),
                 span ->
-                    span.hasName("child").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(1))));
+                    span.hasName("child").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(0))));
   }
 
   @Test
@@ -269,7 +265,7 @@ void testQueryFailedWithFiveArguments() throws InterruptedException {
                             attributeAssertions(
                                 "SELECT MEAN(water_level) FROM;", "SELECT", databaseName)),
                 span ->
-                    span.hasName("child").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(1))));
+                    span.hasName("child").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(0))));
   }
 
   @Test