Add ScheduledExecutorService support to CaptureContextOnSubmitInstrum…

…entation (#787)
kamon-io · Jun 15, 2020 · e390184 · e390184
1 parent d72b70a
commit e390184
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 1 deletion.
diff --git a/...s/src/main/java/kamon/instrumentation/executor/CaptureContextOnSubmitInstrumentation.java b/...s/src/main/java/kamon/instrumentation/executor/CaptureContextOnSubmitInstrumentation.java
@@ -59,6 +59,18 @@ public CaptureContextOnSubmitInstrumentation() {
                 .advise(method("submit").and(withArgument(Runnable.class)), RunnableWrapperAdvisor.class)
                 .advise(method("submit").and(withArgument(Callable.class)), CallableWrapperAdvisor.class)
                 .advise(anyMethods("invokeAny", "invokeAll").and(withArgument(Collection.class)), CallableCollectionWrapperAdvisor.class);
+
+        /**
+         * Instrument all implementations of:
+         *
+         * java.util.concurrent.ScheduledExecutorService::schedule(Runnable, long, TimeUnit)
+         * java.util.concurrent.ScheduledExecutorService::schedule(Callable, long, TimeUnit)
+         *
+         */
+        onSubTypesOf("java.util.concurrent.ScheduledExecutorService")
+                .advise(method("schedule").and(withArgument(0, Runnable.class)), RunnableWrapperAdvisor.class)
+                .advise(method("schedule").and(withArgument(0, Callable.class)), CallableWrapperAdvisor.class);
+
     }
 
     /**

diff --git a/...test/scala/kamon/instrumentation/executor/CaptureContextOnSubmitInstrumentationSpec.scala b/...test/scala/kamon/instrumentation/executor/CaptureContextOnSubmitInstrumentationSpec.scala
@@ -16,7 +16,7 @@
 
 package kamon.instrumentation.executor
 
-import java.util.concurrent.{Callable, CountDownLatch, Executors => JavaExecutors}
+import java.util.concurrent.{Callable, CountDownLatch, TimeUnit, Executors => JavaExecutors}
 
 import com.google.common.util.concurrent.MoreExecutors
 import kamon.Kamon
@@ -116,6 +116,19 @@ class CaptureContextOnSubmitInstrumentationSpec extends WordSpec with Matchers w
       ctx.value should be ("in-runnable-body")
     }
 
+    "capture the context when call schedule(Runnable,long,TimeUnit) in ScheduledThreadPool" in {
+      val executor = JavaExecutors.newScheduledThreadPool(1)
+
+      val ctx = Kamon.runWithContext(testContext("in-runnable-body")) {
+        val runnable = new SimpleRunnable
+        executor.schedule(runnable, 1, TimeUnit.MILLISECONDS)
+        runnable.latch.await()
+        runnable.ctx
+      }
+
+      ctx.value should be ("in-runnable-body")
+    }
+
     "capture the context when call submit(Callable) in ThreadPool" in {
       val executor = JavaExecutors.newSingleThreadExecutor()
       val ctx = Kamon.runWithContext(testContext("in-callable-body")) {
@@ -140,6 +153,19 @@ class CaptureContextOnSubmitInstrumentationSpec extends WordSpec with Matchers w
       ctx.value should be ("in-callable-body")
     }
 
+    "capture the context when call schedule(Callable,long,TimeUnit) in ScheduledThreadPool" in {
+      val executor = JavaExecutors.newScheduledThreadPool(1)
+
+      val ctx = Kamon.runWithContext(testContext("in-callable-body")) {
+        val callable = new SimpleCallable
+        executor.schedule(callable, 1, TimeUnit.MILLISECONDS)
+        callable.latch.await()
+        callable.ctx
+      }
+
+      ctx.value should be ("in-callable-body")
+    }
+
     "capture the context when call submit(Callable) in ForkJoinPool" in {
       val executor = JavaExecutors.newWorkStealingPool()