From 851cfb21ab4289b3326a32d07448922602c7bbe9 Mon Sep 17 00:00:00 2001
From: Lauri Tulmin <ltulmin@splunk.com>
Date: Mon, 20 May 2024 15:50:36 +0300
Subject: [PATCH] Fix flaky finagle http client tests (#11400)

---
 .../javaagent/build.gradle.kts                |  2 +
 ...calSchedulerActivationInstrumentation.java | 47 +++++++++++++++++++
 .../TwitterUtilCoreInstrumentationModule.java | 26 ++++++++++
 3 files changed, 75 insertions(+)
 create mode 100644 instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/LocalSchedulerActivationInstrumentation.java
 create mode 100644 instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/TwitterUtilCoreInstrumentationModule.java

diff --git a/instrumentation/finagle-http-23.11/javaagent/build.gradle.kts b/instrumentation/finagle-http-23.11/javaagent/build.gradle.kts
index 24b29e4f89ad..292ca2ac477a 100644
--- a/instrumentation/finagle-http-23.11/javaagent/build.gradle.kts
+++ b/instrumentation/finagle-http-23.11/javaagent/build.gradle.kts
@@ -30,6 +30,8 @@ val scalified = fun(pack: String): String {
 }
 
 dependencies {
+  bootstrap(project(":instrumentation:executors:bootstrap"))
+
   library("${scalified("com.twitter:finagle-http")}:$finagleVersion")
 
   // should wire netty contexts
diff --git a/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/LocalSchedulerActivationInstrumentation.java b/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/LocalSchedulerActivationInstrumentation.java
new file mode 100644
index 000000000000..f996916f816b
--- /dev/null
+++ b/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/LocalSchedulerActivationInstrumentation.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.finaglehttp.v23_11;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.bootstrap.executors.ContextPropagatingRunnable;
+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;
+
+public class LocalSchedulerActivationInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.twitter.concurrent.LocalScheduler$Activation");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("submit").and(takesArguments(Runnable.class)),
+        this.getClass().getName() + "$WrapRunnableAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class WrapRunnableAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void wrap(@Advice.Argument(value = 0, readOnly = false) Runnable task) {
+      if (task == null) {
+        return;
+      }
+
+      Context context = Java8BytecodeBridge.currentContext();
+      task = ContextPropagatingRunnable.propagateContext(task, context);
+    }
+  }
+}
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
new file mode 100644
index 000000000000..bc9e4ba5a5ed
--- /dev/null
+++ b/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/TwitterUtilCoreInstrumentationModule.java
@@ -0,0 +1,26 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.finaglehttp.v23_11;
+
+import static java.util.Collections.singletonList;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import java.util.List;
+
+@AutoService(InstrumentationModule.class)
+public class TwitterUtilCoreInstrumentationModule extends InstrumentationModule {
+
+  public TwitterUtilCoreInstrumentationModule() {
+    super("twitter-util-core");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return singletonList(new LocalSchedulerActivationInstrumentation());
+  }
+}