PlayWS strict context check (#4365)

* PlayWS strict context check

* add comment

instrumentation/play-ws/play-ws-2.0/javaagent/build.gradle.kts

@@ -45,7 +45,3 @@ dependencies {
 
   latestDepTestLibrary("com.typesafe.play:play-ahc-ws-standalone_$scalaVersion:2.0.+")
 }
-
-tasks.withType<Test>().configureEach {
-  jvmArgs("-Dio.opentelemetry.javaagent.shaded.io.opentelemetry.context.enableStrictContext=false")
-}

instrumentation/play-ws/play-ws-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/playws/v2_0/PlayWsInstrumentationModule.java

@@ -13,6 +13,7 @@ import com.google.auto.service.AutoService;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.instrumentation.playws.AbstractBootstrapInstrumentation;
 import io.opentelemetry.javaagent.instrumentation.playws.AsyncHttpClientInstrumentation;
 import io.opentelemetry.javaagent.instrumentation.playws.HandlerPublisherInstrumentation;
 import java.util.List;
@@ -32,7 +33,8 @@ public class PlayWsInstrumentationModule extends InstrumentationModule {
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(
         new AsyncHttpClientInstrumentation(this.getClass().getName() + "$ClientAdvice"),
-        new HandlerPublisherInstrumentation());
+        new HandlerPublisherInstrumentation(),
+        new AbstractBootstrapInstrumentation());
   }
 
   @SuppressWarnings("unused")

instrumentation/play-ws/play-ws-2.1/javaagent/build.gradle.kts

@@ -42,7 +42,3 @@ dependencies {
   testInstrumentation(project(":instrumentation:akka-http-10.0:javaagent"))
   testInstrumentation(project(":instrumentation:akka-actor-2.5:javaagent"))
 }
-
-tasks.withType<Test>().configureEach {
-  jvmArgs("-Dio.opentelemetry.javaagent.shaded.io.opentelemetry.context.enableStrictContext=false")
-}

instrumentation/play-ws/play-ws-2.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/playws/v2_1/PlayWsInstrumentationModule.java

@@ -13,6 +13,7 @@ import com.google.auto.service.AutoService;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.instrumentation.playws.AbstractBootstrapInstrumentation;
 import io.opentelemetry.javaagent.instrumentation.playws.AsyncHttpClientInstrumentation;
 import io.opentelemetry.javaagent.instrumentation.playws.HandlerPublisherInstrumentation;
 import java.util.List;
@@ -32,7 +33,8 @@ public class PlayWsInstrumentationModule extends InstrumentationModule {
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(
         new AsyncHttpClientInstrumentation(this.getClass().getName() + "$ClientAdvice"),
-        new HandlerPublisherInstrumentation());
+        new HandlerPublisherInstrumentation(),
+        new AbstractBootstrapInstrumentation());
   }
 
   @SuppressWarnings("unused")

instrumentation/play-ws/play-ws-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/playws/AbstractBootstrapInstrumentation.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.playws;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class AbstractBootstrapInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("play.shaded.ahc.io.netty.bootstrap.AbstractBootstrap");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("initAndRegister"),
+        AbstractBootstrapInstrumentation.class.getName() + "$DisablePropagationAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class DisablePropagationAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static Scope onEnter() {
+      // Prevent context from leaking by running this method under root context.
+      // Root context is not propagated by executor instrumentation.
+      if (Java8BytecodeBridge.currentContext() != Java8BytecodeBridge.rootContext()) {
+        return Java8BytecodeBridge.rootContext().makeCurrent();
+      }
+      return null;
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.Enter Scope scope) {
+      if (scope != null) {
+        scope.close();
+      }
+    }
+  }
+}