Explorar el Código

Fix flaky finagle test (#11441)

Lauri Tulmin hace 9 meses
padre
commit
c9cbd8a97a

+ 24 - 0
instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/Function1Wrapper.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.finaglehttp.v23_11;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import scala.Function1;
+
+public final class Function1Wrapper {
+
+  public static <T1, R> Function1<T1, R> wrap(Function1<T1, R> function1) {
+    Context context = Context.current();
+    return value -> {
+      try (Scope ignored = context.makeCurrent()) {
+        return function1.apply(value);
+      }
+    };
+  }
+
+  private Function1Wrapper() {}
+}

+ 46 - 0
instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/PromiseMonitoredInstrumentation.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.finaglehttp.v23_11;
+
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import scala.Function1;
+
+public class PromiseMonitoredInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.twitter.util.Promise$Monitored");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isConstructor().and(takesArgument(1, named("scala.Function1"))),
+        this.getClass().getName() + "$WrapFunctionAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class WrapFunctionAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void wrap(
+        @Advice.Argument(value = 1, readOnly = false) Function1<?, ?> function1) {
+      if (function1 == null) {
+        return;
+      }
+
+      function1 = Function1Wrapper.wrap(function1);
+    }
+  }
+}

+ 3 - 2
instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/TwitterUtilCoreInstrumentationModule.java

@@ -5,7 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.finaglehttp.v23_11;
 
-import static java.util.Collections.singletonList;
+import static java.util.Arrays.asList;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
@@ -21,6 +21,7 @@ public class TwitterUtilCoreInstrumentationModule extends InstrumentationModule
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return singletonList(new LocalSchedulerActivationInstrumentation());
+    return asList(
+        new LocalSchedulerActivationInstrumentation(), new PromiseMonitoredInstrumentation());
   }
 }