1 files changed, 41 insertions, 0 deletions
diff --git a/kamon-core/src/main/scala/kamon/instrumentation/ActorInstrumentation.scala b/kamon-core/src/main/scala/kamon/instrumentation/ActorInstrumentation.scala
new file mode 100644
index 00000000..4e47c2a4
--- /dev/null
+++ b/kamon-core/src/main/scala/kamon/instrumentation/ActorInstrumentation.scala
@@ -0,0 +1,41 @@
+package kamon.instrumentation
+
+import akka.actor.{Props, ActorSystem, ActorRef}
+import akka.dispatch.{MessageDispatcher, Envelope}
+import kamon.{Tracer, TraceContext}
+import kamon.instrumentation.SimpleContextPassingInstrumentation.SimpleTraceMessage
+
+trait ActorInstrumentationConfiguration {
+  def sendMessageTransformation(from: ActorRef, to: ActorRef, message: Any): Any
+  def receiveInvokeInstrumentation(system: ActorSystem, self: ActorRef, props: Props, dispatcher: MessageDispatcher, parent: ActorRef): ActorReceiveInvokeInstrumentation
+}
+
+
+trait ActorReceiveInvokeInstrumentation {
+  def preReceive(envelope: Envelope): (Envelope, Option[TraceContext])
+}
+
+object ActorReceiveInvokeInstrumentation {
+  val noopPreReceive = new ActorReceiveInvokeInstrumentation{
+    def preReceive(envelope: Envelope): (Envelope, Option[TraceContext]) = (envelope, None)
+  }
+}
+
+class SimpleContextPassingInstrumentation extends ActorInstrumentationConfiguration {
+  def sendMessageTransformation(from: ActorRef, to: ActorRef, message: Any): Any = SimpleTraceMessage(message, Tracer.context)
+
+  def receiveInvokeInstrumentation(system: ActorSystem, self: ActorRef, props: Props, dispatcher: MessageDispatcher, parent: ActorRef): ActorReceiveInvokeInstrumentation = {
+    new ActorReceiveInvokeInstrumentation {
+      def preReceive(envelope: Envelope): (Envelope, Option[TraceContext]) = envelope match {
+        case env @ Envelope(SimpleTraceMessage(msg, ctx), _) => (env.copy(message = msg), ctx)
+        case anyOther                                        => (anyOther, None)
+      }
+    }
+  }
+}
+
+object SimpleContextPassingInstrumentation {
+  case class SimpleTraceMessage(message: Any, context: Option[TraceContext])
+}
+
+