From b019fb4beb90397c8498a3f126d14392eefe0ca3 Mon Sep 17 00:00:00 2001
From: Diego <diegolparra@gmail.com>
Date: Fri, 27 Nov 2015 14:46:32 -0300
Subject: ! akka: avoid error thrown in dispatcher instrumentation when using
 custom dispatchers and close #290

---
 .../DispatcherInstrumentation.scala                | 28 +++++++++++++++-------
 1 file changed, 20 insertions(+), 8 deletions(-)

(limited to 'kamon-akka/src/main/scala/kamon/akka/instrumentation')

diff --git a/kamon-akka/src/main/scala/kamon/akka/instrumentation/DispatcherInstrumentation.scala b/kamon-akka/src/main/scala/kamon/akka/instrumentation/DispatcherInstrumentation.scala
index 0d504343..691cae13 100644
--- a/kamon-akka/src/main/scala/kamon/akka/instrumentation/DispatcherInstrumentation.scala
+++ b/kamon-akka/src/main/scala/kamon/akka/instrumentation/DispatcherInstrumentation.scala
@@ -1,6 +1,6 @@
 /*
  * =========================================================================================
- * Copyright © 2013-2014 the kamon project <http://kamon.io/>
+ * Copyright © 2013-2015 the kamon project <http://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
@@ -16,6 +16,7 @@
 
 package akka.kamon.instrumentation
 
+import java.lang.reflect.Method
 import java.util.concurrent.{ ExecutorService, ThreadPoolExecutor }
 
 import akka.actor.{ ActorContext, Props, ActorSystem, ActorSystemImpl }
@@ -47,15 +48,26 @@ class DispatcherInstrumentation {
 
     // Yes, reflection sucks, but this piece of code is only executed once on ActorSystem's startup.
     val defaultDispatcher = system.dispatcher
-    val executorServiceDelegateField = defaultDispatcher.getClass.getDeclaredField("executorServiceDelegate")
-    executorServiceDelegateField.setAccessible(true)
+    val defaultDispatcherExecutor = extractExecutor(defaultDispatcher.asInstanceOf[MessageDispatcher])
+    registerDispatcher(Dispatchers.DefaultDispatcherId, defaultDispatcherExecutor, system)
+  }
 
-    val lazyExecutorServiceDelegate = executorServiceDelegateField.get(defaultDispatcher)
-    val executorField = lazyExecutorServiceDelegate.getClass.getMethod("executor")
-    executorField.setAccessible(true)
+  private def extractExecutor(dispatcher: MessageDispatcher): ExecutorService = {
+    val executorServiceMethod: Method = {
+      // executorService is protected
+      val method = classOf[Dispatcher].getDeclaredMethod("executorService")
+      method.setAccessible(true)
+      method
+    }
 
-    val defaultDispatcherExecutor = executorField.invoke(lazyExecutorServiceDelegate).asInstanceOf[ExecutorService]
-    registerDispatcher(Dispatchers.DefaultDispatcherId, defaultDispatcherExecutor, system)
+    dispatcher match {
+      case x: Dispatcher ⇒
+        val executor = executorServiceMethod.invoke(x) match {
+          case delegate: ExecutorServiceDelegate ⇒ delegate.executor
+          case other                             ⇒ other
+        }
+        executor.asInstanceOf[ExecutorService]
+    }
   }
 
   private def registerDispatcher(dispatcherName: String, executorService: ExecutorService, system: ActorSystem): Unit =
-- 
cgit v1.2.3