diff --git a/instrumentation/kamon-executors/src/main/java/kamon/instrumentation/executor/CaptureContextOnSubmitInstrumentation.java b/instrumentation/kamon-executors/src/main/java/kamon/instrumentation/executor/CaptureContextOnSubmitInstrumentation.java index ed1ce785b..04e0b490a 100644 --- a/instrumentation/kamon-executors/src/main/java/kamon/instrumentation/executor/CaptureContextOnSubmitInstrumentation.java +++ b/instrumentation/kamon-executors/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/instrumentation/kamon-executors/src/test/scala/kamon/instrumentation/executor/CaptureContextOnSubmitInstrumentationSpec.scala b/instrumentation/kamon-executors/src/test/scala/kamon/instrumentation/executor/CaptureContextOnSubmitInstrumentationSpec.scala index 61e12f4d5..5024a3856 100644 --- a/instrumentation/kamon-executors/src/test/scala/kamon/instrumentation/executor/CaptureContextOnSubmitInstrumentationSpec.scala +++ b/instrumentation/kamon-executors/src/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()