From 9028dae799d2668f3e217b0bf7f65c2b623056d7 Mon Sep 17 00:00:00 2001
From: Lauri Tulmin <ltulmin@splunk.com>
Date: Thu, 23 May 2024 16:56:37 +0300
Subject: [PATCH] Fix flaky finagle test

---
 .../finaglehttp/v23_11/Function1Wrapper.java  | 24 ++++++++++
 .../PromiseMonitoredInstrumentation.java      | 46 +++++++++++++++++++
 .../TwitterUtilCoreInstrumentationModule.java |  5 +-
 3 files changed, 73 insertions(+), 2 deletions(-)
 create mode 100644 instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/Function1Wrapper.java
 create mode 100644 instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/PromiseMonitoredInstrumentation.java

diff --git a/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/Function1Wrapper.java b/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/Function1Wrapper.java
new file mode 100644
index 000000000000..586be1efc0fe
--- /dev/null
+++ b/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() {}
+}
diff --git a/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/PromiseMonitoredInstrumentation.java b/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/PromiseMonitoredInstrumentation.java
new file mode 100644
index 000000000000..2cd637f4d385
--- /dev/null
+++ b/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);
+    }
+  }
+}
diff --git a/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/TwitterUtilCoreInstrumentationModule.java b/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/TwitterUtilCoreInstrumentationModule.java
index bc9e4ba5a5ed..022dea36455d 100644
--- a/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/TwitterUtilCoreInstrumentationModule.java
+++ b/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 TwitterUtilCoreInstrumentationModule() {
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return singletonList(new LocalSchedulerActivationInstrumentation());
+    return asList(
+        new LocalSchedulerActivationInstrumentation(), new PromiseMonitoredInstrumentation());
   }
 }