Merge branch 'issue#271/fix-balancing-pool-metrics'

1 files changed, 32 insertions, 0 deletions

diff --git a/kamon-akka/src/main/scala/kamon/akka/instrumentation/EnvelopeInstrumentation.scala b/kamon-akka/src/main/scala/kamon/akka/instrumentation/EnvelopeInstrumentation.scala
new file mode 100644
index 00000000..0bb50dc2
--- /dev/null
+++ b/kamon-akka/src/main/scala/kamon/akka/instrumentation/EnvelopeInstrumentation.scala
@@ -0,0 +1,32 @@
+package akka.kamon.instrumentation
+
+import kamon.trace.{ EmptyTraceContext, TraceContext }
+import kamon.util.RelativeNanoTimestamp
+import org.aspectj.lang.annotation.{ DeclareMixin, Aspect }
+
+case class EnvelopeContext(nanoTime: RelativeNanoTimestamp, context: TraceContext)
+
+object EnvelopeContext {
+  val Empty = EnvelopeContext(RelativeNanoTimestamp.zero, EmptyTraceContext)
+}
+
+trait InstrumentedEnvelope {
+  def envelopeContext(): EnvelopeContext
+  def setEnvelopeContext(envelopeContext: EnvelopeContext): Unit
+}
+
+object InstrumentedEnvelope {
+  def apply(): InstrumentedEnvelope = new InstrumentedEnvelope {
+    var envelopeContext: EnvelopeContext = _
+
+    def setEnvelopeContext(envelopeContext: EnvelopeContext): Unit =
+      this.envelopeContext = envelopeContext
+  }
+}
+
+@Aspect
+class EnvelopeContextIntoEnvelopeMixin {
+
+  @DeclareMixin("akka.dispatch.Envelope")
+  def mixinInstrumentationToEnvelope: InstrumentedEnvelope = InstrumentedEnvelope()
+}
\ No newline at end of file