From 25491d7d6182ee30fc5c50203e2f1e02c18f546b Mon Sep 17 00:00:00 2001
From: Lauri Tulmin <ltulmin@splunk.com>
Date: Wed, 13 Oct 2021 23:03:59 +0300
Subject: [PATCH] Mongo4 strict context check (#4364)

* Mongo4 strict context check

* Update instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/AsyncWorkManagerInstrumentation.java

Co-authored-by: Trask Stalnaker <trask.stalnaker@gmail.com>

Co-authored-by: Trask Stalnaker <trask.stalnaker@gmail.com>
---
 .../mongo-4.0/javaagent/build.gradle.kts      |  3 --
 .../v4_0/AsyncWorkManagerInstrumentation.java | 53 +++++++++++++++++++
 .../MongoClientInstrumentationModule.java     |  3 +-
 3 files changed, 55 insertions(+), 4 deletions(-)
 create mode 100644 instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/AsyncWorkManagerInstrumentation.java

diff --git a/instrumentation/mongo/mongo-4.0/javaagent/build.gradle.kts b/instrumentation/mongo/mongo-4.0/javaagent/build.gradle.kts
index a919702cfb..6c3030c42a 100644
--- a/instrumentation/mongo/mongo-4.0/javaagent/build.gradle.kts
+++ b/instrumentation/mongo/mongo-4.0/javaagent/build.gradle.kts
@@ -26,8 +26,5 @@ dependencies {
 tasks {
   test {
     usesService(gradle.sharedServices.registrations["testcontainersBuildService"].getService())
-    if (findProperty("testLatestDeps") as Boolean) {
-      jvmArgs("-Dio.opentelemetry.javaagent.shaded.io.opentelemetry.context.enableStrictContext=false")
-    }
   }
 }
diff --git a/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/AsyncWorkManagerInstrumentation.java b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/AsyncWorkManagerInstrumentation.java
new file mode 100644
index 0000000000..0c39631dea
--- /dev/null
+++ b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/AsyncWorkManagerInstrumentation.java
@@ -0,0 +1,53 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.mongo.v4_0;
+
+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 AsyncWorkManagerInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.mongodb.internal.connection.DefaultConnectionPool$AsyncWorkManager");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    // this method sets up a new thread pool and submits a task to it, we need to avoid context
+    // propagating there
+    transformer.applyAdviceToMethod(
+        named("initUnlessClosed"),
+        AsyncWorkManagerInstrumentation.class.getName() + "$DisablePropagationAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class DisablePropagationAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static Scope onEnter() {
+      if (Java8BytecodeBridge.currentContext() != Java8BytecodeBridge.rootContext()) {
+        // Prevent context from leaking by running this method under root context.
+        // Root context is not propagated by executor instrumentation.
+        return Java8BytecodeBridge.rootContext().makeCurrent();
+      }
+      return null;
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.Enter Scope scope) {
+      if (scope != null) {
+        scope.close();
+      }
+    }
+  }
+}
diff --git a/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/MongoClientInstrumentationModule.java b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/MongoClientInstrumentationModule.java
index 4d32ec8475..928a07e1e6 100644
--- a/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/MongoClientInstrumentationModule.java
+++ b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/MongoClientInstrumentationModule.java
@@ -25,6 +25,7 @@ public class MongoClientInstrumentationModule extends InstrumentationModule {
         new MongoClientSettingsBuilderInstrumentation(),
         new InternalStreamConnectionInstrumentation(),
         new BaseClusterInstrumentation(),
-        new DefaultConnectionPoolInstrumentation());
+        new DefaultConnectionPoolInstrumentation(),
+        new AsyncWorkManagerInstrumentation());
   }
 }