Update Mongo instrumentation to latest version

build.sbt

@@ -121,6 +121,7 @@ lazy val instrumentation = (project in file("instrumentation"))
     `kamon-jdbc`,
     `kamon-kafka`,
     `kamon-mongo`,
+    `kamon-mongo-legacy`,
     `kamon-cassandra`,
     `kamon-elasticsearch`,
     `kamon-spring`,
@@ -288,9 +289,7 @@ lazy val `kamon-kafka` = (project in file("instrumentation/kamon-kafka"))
   ).dependsOn(`kamon-core`, `kamon-executors`, `kamon-testkit` % "test")
 
 
-
-
-lazy val `kamon-mongo` = (project in file("instrumentation/kamon-mongo"))
+lazy val `kamon-mongo-legacy` = (project in file("instrumentation/kamon-mongo-legacy"))
   .disablePlugins(AssemblyPlugin)
   .enablePlugins(JavaAgent)
   .settings(instrumentationSettings)
@@ -307,6 +306,24 @@ lazy val `kamon-mongo` = (project in file("instrumentation/kamon-mongo"))
     )
   ).dependsOn(`kamon-core`, `kamon-instrumentation-common`, `kamon-testkit` % "test")
 
+
+lazy val `kamon-mongo` = (project in file("instrumentation/kamon-mongo"))
+  .disablePlugins(AssemblyPlugin)
+  .enablePlugins(JavaAgent)
+  .settings(instrumentationSettings)
+  .settings(
+    libraryDependencies ++= Seq(
+      kanelaAgent % "provided",
+      "org.mongodb"         %   "mongodb-driver-sync"             % "4.2.3" % "provided",
+      "org.mongodb.scala"   %%  "mongo-scala-driver"              % "4.2.3" % "provided",
+      "org.mongodb"         %   "mongodb-driver-reactivestreams"  % "4.2.3" % "provided",
+
+      scalatest % "test",
+      logbackClassic % "test",
+      "de.flapdoodle.embed" %   "de.flapdoodle.embed.mongo"       % "2.2.0" % "test"
+    )
+  ).dependsOn(`kamon-core`, `kamon-instrumentation-common`, `kamon-testkit` % "test")
+
 lazy val `kamon-cassandra` = (project in file("instrumentation/kamon-cassandra"))
   .disablePlugins(AssemblyPlugin)
   .enablePlugins(JavaAgent)
@@ -706,6 +723,7 @@ val `kamon-bundle` = (project in file("bundle/kamon-bundle"))
     `kamon-jdbc` % "shaded",
     `kamon-kafka` % "shaded",
     `kamon-mongo` % "shaded",
+    `kamon-mongo-legacy` % "shaded",
     `kamon-cassandra` % "shaded",
     `kamon-elasticsearch` % "shaded",
     `kamon-spring` % "shaded",

instrumentation/kamon-mongo-legacy/Readme.MD

@@ -0,0 +1,18 @@
+This project is here to preserve backwards compatibility with old
+versions of the mongo client.
+When updating versions of the supported clients, 
+we found that keeping the old instrumentation, and the new instrumentation in the same package
+would impose a great maintenance burden.
+
+The easiest solution was this: 
+If you're using old versions of the client, use this project.
+For newer ones, use `kamon-mongo`.
+
+Of course, none of this applies if you're using `kamon-bundle`, in which case you can 
+ignore all of this, since everything should work as it did.
+
+Versions supported by this project:
+```
+mongodb-driver-sync <= 3.11.0
+mongo-scala-driver <= 2.7.0
+```

instrumentation/kamon-mongo-legacy/src/main/java/kamon/instrumentation/legacy/mongo/AsyncBatchCursorGetMoreAdvice.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2013-2020 The Kamon Project <https://kamon.io>
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kamon.instrumentation.legacy.mongo;
+
+import com.mongodb.MongoNamespace;
+import com.mongodb.async.SingleResultCallback;
+import kamon.instrumentation.context.HasContext;
+import kamon.trace.Span;
+import kanela.agent.libs.net.bytebuddy.asm.Advice;
+
+import java.util.List;
+
+public class AsyncBatchCursorGetMoreAdvice {
+
+  @Advice.OnMethodEnter
+  public static <T> void enter(
+      @Advice.This Object batchCursor,
+      @Advice.FieldValue("namespace") MongoNamespace namespace,
+      @Advice.Argument(value = 2, readOnly = false) SingleResultCallback<T> callback) {
+
+    final Span parentSpan = ((HasContext) batchCursor).context().get(Span.Key());
+    final Span getMoreSpan = MongoClientInstrumentation.getMoreSpanBuilder(parentSpan, namespace).start();
+    callback = spanCompletingCallback(callback, getMoreSpan);
+  }
+
+  public static <T> SingleResultCallback<T> spanCompletingCallback(SingleResultCallback<T> originalCallback, Span span) {
+    return new SingleResultCallback<T>() {
+
+      @Override
+      public void onResult(T result, Throwable t) {
+        try {
+          if(t == null) {
+            span.finish();
+          } else {
+            span.fail(t).finish();
+          }
+
+        } finally {
+          originalCallback.onResult(result, t);
+        }
+      }
+    };
+  }
+}

instrumentation/kamon-mongo-legacy/src/main/java/kamon/instrumentation/legacy/mongo/BatchCursorGetMoreAdvice.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2013-2020 The Kamon Project <https://kamon.io>
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kamon.instrumentation.legacy.mongo;
+
+import com.mongodb.MongoNamespace;
+import kamon.Kamon;
+import kamon.context.Storage;
+import kamon.instrumentation.context.HasContext;
+import kamon.trace.Span;
+import kanela.agent.libs.net.bytebuddy.asm.Advice;
+
+public class BatchCursorGetMoreAdvice {
+
+  @Advice.OnMethodEnter
+  public static Storage.Scope enter(
+      @Advice.This Object batchCursor,
+      @Advice.FieldValue("namespace") MongoNamespace namespace) {
+
+    final Span parentSpan = ((HasContext) batchCursor).context().get(Span.Key());
+    final Span getMoreSpan = MongoClientInstrumentation.getMoreSpanBuilder(parentSpan, namespace).start();
+
+    return Kamon.storeContext(Kamon.currentContext().withEntry(Span.Key(), getMoreSpan));
+  }
+
+  @Advice.OnMethodExit(onThrowable = Throwable.class)
+  public static void exit(@Advice.Enter Storage.Scope scope, @Advice.Thrown Throwable t) {
+    final Span span = scope.context().get(Span.Key());
+
+    if(t == null) {
+      span.finish();
+    } else {
+      span.fail(t).finish();
+    }
+
+    scope.close();
+  }
+}

instrumentation/kamon-mongo-legacy/src/main/java/kamon/instrumentation/legacy/mongo/ExecuteAsyncOperationAdvice.java

@@ -0,0 +1,83 @@
+/*
+ * Copyright 2013-2020 The Kamon Project <https://kamon.io>
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kamon.instrumentation.legacy.mongo;
+
+import com.mongodb.MongoNamespace;
+import com.mongodb.async.AsyncBatchCursor;
+import com.mongodb.async.SingleResultCallback;
+import com.mongodb.bulk.BulkWriteResult;
+import com.mongodb.operation.BatchCursor;
+import kamon.Kamon;
+import kamon.context.Context;
+import kamon.context.Storage;
+import kamon.instrumentation.context.HasContext;
+import kamon.trace.Span;
+import kanela.agent.libs.net.bytebuddy.asm.Advice;
+import kamon.instrumentation.legacy.mongo.MongoClientInstrumentation.HasOperationName;
+
+public class ExecuteAsyncOperationAdvice {
+
+  @Advice.OnMethodEnter()
+  public static <T> void enter(
+      @Advice.This Object operation,
+      @Advice.FieldValue("namespace") MongoNamespace namespace,
+      @Advice.Origin("#t") String operationClassName,
+      @Advice.Argument(value = 1, readOnly = false) SingleResultCallback<T> callback) {
+
+    final String operationClass = operation instanceof HasOperationName ?
+        ((HasOperationName) operation).name() :
+        operationClassName;
+
+    final Span clientSpan = MongoClientInstrumentation.clientSpanBuilder(namespace, operationClass).start();
+    callback = spanCompletingCallback(callback, Kamon.currentContext().withEntry(Span.Key(), clientSpan));
+  }
+
+  public static <T> SingleResultCallback<T> spanCompletingCallback(SingleResultCallback<T> originalCallback, Context context) {
+    return new SingleResultCallback<T>() {
+
+      @Override
+      public void onResult(T result, Throwable t) {
+        try {
+          final Span span = context.get(Span.Key());
+
+          if (t == null) {
+            if (result instanceof BulkWriteResult) {
+              final BulkWriteResult bulkResult = (BulkWriteResult) result;
+
+              span
+                .tag("mongo.bulk.ack", bulkResult.wasAcknowledged())
+                .tag("mongo.bulk.inserted", bulkResult.getInsertedCount())
+                .tag("mongo.bulk.modified", bulkResult.getModifiedCount())
+                .tag("mongo.bulk.matched", bulkResult.getMatchedCount())
+                .tag("mongo.bulk.deleted", bulkResult.getDeletedCount());
+            }
+
+            if (result instanceof AsyncBatchCursor && result instanceof HasContext) {
+              ((HasContext) result).setContext(Context.of(Span.Key(), span));
+            }
+
+            span.finish();
+          } else {
+            span.fail(t).finish();
+          }
+        } finally {
+          originalCallback.onResult(result, t);
+        }
+      }
+    };
+  }
+}

instrumentation/kamon-mongo-legacy/src/main/java/kamon/instrumentation/legacy/mongo/ExecuteOperationAdvice.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2013-2020 The Kamon Project <https://kamon.io>
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kamon.instrumentation.legacy.mongo;
+
+import com.mongodb.MongoNamespace;
+import com.mongodb.async.SingleResultCallback;
+import com.mongodb.bulk.BulkWriteResult;
+import com.mongodb.operation.BatchCursor;
+import kamon.Kamon;
+import kamon.context.Context;
+import kamon.context.Storage;
+import kamon.instrumentation.context.HasContext;
+import kamon.trace.Span;
+import kanela.agent.libs.net.bytebuddy.asm.Advice;
+
+public class ExecuteOperationAdvice {
+
+  @Advice.OnMethodEnter()
+  public static <T> Storage.Scope enter(
+      @Advice.This Object operation,
+      @Advice.FieldValue("namespace") MongoNamespace namespace,
+      @Advice.Origin("#t") String operationClassName) {
+
+    final String operationClass = operation instanceof MongoClientInstrumentation.HasOperationName ?
+        ((MongoClientInstrumentation.HasOperationName) operation).name() :
+        operationClassName;
+
+    final Span clientSpan = MongoClientInstrumentation.clientSpanBuilder(namespace, operationClass).start();
+
+    return Kamon.storeContext(Kamon.currentContext().withEntry(Span.Key(), clientSpan));
+  }
+
+  @Advice.OnMethodExit
+  public static void exit(@Advice.Enter Storage.Scope scope, @Advice.Return Object result) {
+    final Span span = scope.context().get(Span.Key());
+
+    if(result instanceof BatchCursor && result instanceof HasContext) {
+      ((HasContext) result).setContext(Context.of(Span.Key(), span));
+    }
+
+    if(result instanceof BulkWriteResult) {
+      final BulkWriteResult bulkResult = (BulkWriteResult) result;
+
+      span
+        .tag("mongo.bulk.ack", bulkResult.wasAcknowledged())
+        .tag("mongo.bulk.inserted", bulkResult.getInsertedCount())
+        .tag("mongo.bulk.modified", bulkResult.getModifiedCount())
+        .tag("mongo.bulk.matched", bulkResult.getMatchedCount())
+        .tag("mongo.bulk.deleted", bulkResult.getDeletedCount());
+    }
+
+    span.finish();
+    scope.close();
+  }
+}

instrumentation/kamon-mongo-legacy/src/main/resources/reference.conf

@@ -0,0 +1,17 @@
+kanela.modules {
+  mongo-driver {
+    name = "Mongo Driver Instrumentation"
+    description = "Provides automatic tracing of client operations on the official Mongo driver"
+
+    instrumentations = [
+      "kamon.instrumentation.legacy.mongo.MongoClientInstrumentation"
+    ]
+
+    within = [
+      "^com.mongodb.client.*"
+      "^com.mongodb.async.client.*"
+      "^com.mongodb.internal.*"
+      "^com.mongodb.operation.*"
+    ]
+  }
+}