ProducerMessage instrumentation for Alpakka + Kafka applications

Author: ivantopo

build.sbt

@@ -143,7 +143,8 @@ val instrumentationProjects = Seq[ProjectReference](
   `kamon-caffeine`,
   `kamon-lagom`,
   `kamon-finagle`,
-  `kamon-aws-sdk`
+  `kamon-aws-sdk`,
+  `kamon-alpakka-kafka`
 )
 
 lazy val instrumentation = (project in file("instrumentation"))
@@ -526,7 +527,7 @@ lazy val `kamon-tapir` = (project in file("instrumentation/kamon-tapir"))
   .enablePlugins(JavaAgent)
   .settings(
     instrumentationSettings,
-    crossScalaVersions := Seq("2.12.11", "2.13.1"),
+    crossScalaVersions := Seq(`scala_2.12_version`, `scala_2.13_version`),
     libraryDependencies ++= Seq(
       kanelaAgent % "provided",
       "com.softwaremill.sttp.tapir" %% "tapir-core" % "0.17.9" % "provided",
@@ -615,6 +616,22 @@ lazy val `kamon-aws-sdk` = (project in file("instrumentation/kamon-aws-sdk"))
     )
   ).dependsOn(`kamon-core`, `kamon-executors`, `kamon-testkit` % "test")
 
+lazy val `kamon-alpakka-kafka` = (project in file("instrumentation/kamon-alpakka-kafka"))
+  .disablePlugins(AssemblyPlugin)
+  .enablePlugins(JavaAgent)
+  .settings(instrumentationSettings)
+  .settings(
+    crossScalaVersions := Seq(`scala_2.12_version`, `scala_2.13_version`),
+    libraryDependencies ++= Seq(
+      kanelaAgent % "provided",
+      "com.typesafe.akka" %% "akka-stream-kafka" % "2.1.1" % "provided",
+      "com.typesafe.akka" %% "akka-stream" % "2.6.19" % "provided",
+
+      scalatest % "test",
+      logbackClassic % "test"
+    )
+  ).dependsOn(`kamon-core`, `kamon-akka`, `kamon-testkit` % "test")
+
 /**
  * Reporters
  */
@@ -912,6 +929,7 @@ lazy val `kamon-bundle-dependencies-2-12-and-up` = (project in file("bundle/kamo
     `kamon-akka-grpc`,
     `kamon-finagle`,
     `kamon-tapir`,
+    `kamon-alpakka-kafka`
   )
 
 lazy val `kamon-bundle` = (project in file("bundle/kamon-bundle"))

instrumentation/kamon-alpakka-kafka/src/main/resources/reference.conf

@@ -0,0 +1,26 @@
+# ===================================== #
+# Kamon Alpakka Reference Configuration #
+# ===================================== #
+
+kamon.instrumentation.alpakka {
+
+}
+
+kanela {
+  modules {
+    alpakka {
+
+      name = "Alpakka"
+      description = "PREVIEW. Provides context propagation for Alpakka applications"
+      instrumentations = [
+        "kamon.instrumentation.alpakka.kafka.ProducerMessageInstrumentation"
+      ]
+
+      within = [
+        "akka.kafka.ProducerMessage\\$Message",
+        "akka.kafka.ProducerMessage\\$MultiMessage",
+        "akka.kafka.internal.DefaultProducerStageLogic"
+      ]
+    }
+  }
+}

instrumentation/kamon-alpakka-kafka/src/main/scala/kamon/instrumentation/alpakka/kafka/ProducerMessageInstrumentation.scala

@@ -0,0 +1,55 @@
+/*
+ *  ==========================================================================================
+ *  Copyright © 2013-2022 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
+package instrumentation
+package alpakka
+package kafka
+
+import kamon.Kamon
+import kamon.context.Storage
+import kamon.context.Storage.Scope
+import kamon.instrumentation.context.HasContext
+import kanela.agent.api.instrumentation.InstrumentationBuilder
+import kanela.agent.libs.net.bytebuddy.asm.Advice
+
+class ProducerMessageInstrumentation extends InstrumentationBuilder {
+
+  /**
+    * Captures the current context the a Message or MultiMessage is created and restores it while
+    * the ProducerLogic is running, so the proper context gets propagated to the Kafka Producer.
+    */
+  onTypes("akka.kafka.ProducerMessage$Message", "akka.kafka.ProducerMessage$MultiMessage")
+    .mixin(classOf[HasContext.MixinWithInitializer])
+
+  onTypes("akka.kafka.internal.DefaultProducerStageLogic", "akka.kafka.internal.CommittingProducerSinkStageLogic")
+    .advise(method("produce"), ProduceWithEnvelopeContext)
+}
+
+object ProduceWithEnvelopeContext {
+
+  @Advice.OnMethodEnter
+  def enter(@Advice.Argument(0) envelope: Any): Storage.Scope = {
+    envelope match {
+      case hasContext: HasContext => Kamon.storeContext(hasContext.context)
+      case _                      => Scope.Empty
+    }
+  }
+
+  @Advice.OnMethodExit(onThrowable = classOf[Throwable])
+  def exit(@Advice.Enter scope: Storage.Scope): Unit =
+    scope.close()
+}

project/Build.scala

@@ -272,6 +272,7 @@ object AssemblyTweaks extends AutoPlugin {
     assembly / assemblyMergeStrategy := {
       case s if s.startsWith("LICENSE") => MergeStrategy.discard
       case s if s.startsWith("about") => MergeStrategy.discard
+      case "version.conf" => MergeStrategy.concat
       case x => (assembly / assemblyMergeStrategy).value(x)
     }
   ) ++ inConfig(Shaded)(Defaults.configSettings)