instrument the Akka scheduler.scheduleOnce function, fixes #1093 (#1135)

* instrument the Akka scheduler.scheduleOnce function, fixes #1093

* add license headers

instrumentation/kamon-akka/src/common/java/kamon/instrumentation/akka/instrumentations/SchedulerRunnableAdvice.java

@@ -0,0 +1,50 @@
+/* =========================================================================================
+ * Copyright © 2013-2022 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
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+ * either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ * =========================================================================================
+ */
+
+package kamon.instrumentation.akka.instrumentations;
+
+import kamon.Kamon;
+import kamon.context.Context;
+import kamon.context.Storage;
+import kanela.agent.libs.net.bytebuddy.asm.Advice;
+
+public class SchedulerRunnableAdvice {
+
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static void enter(@Advice.Argument(value = 1, readOnly = false) Runnable runnable) {
+    runnable = new ContextAwareRunnable(Kamon.currentContext(), runnable);
+  }
+
+  public static class ContextAwareRunnable implements Runnable {
+    private final Context context;
+    private final Runnable underlyingRunnable;
+
+    public ContextAwareRunnable(Context context, Runnable underlyingRunnable) {
+      this.context = context;
+      this.underlyingRunnable = underlyingRunnable;
+    }
+
+    @Override
+    public void run() {
+      final Storage.Scope scope = Kamon.storeContext(context);
+
+      try {
+        underlyingRunnable.run();
+      } finally {
+        scope.close();
+      }
+    }
+  }
+}

instrumentation/kamon-akka/src/common/resources/reference.conf

@@ -143,7 +143,8 @@ kanela.modules {
       "kamon.instrumentation.akka.instrumentations.ActorRefInstrumentation",
       "kamon.instrumentation.akka.instrumentations.akka_25.DispatcherInstrumentation",
       "kamon.instrumentation.akka.instrumentations.akka_26.DispatcherInstrumentation",
-      "kamon.instrumentation.akka.instrumentations.akka_26.ActorMonitorInstrumentation"
+      "kamon.instrumentation.akka.instrumentations.akka_26.ActorMonitorInstrumentation",
+      "kamon.instrumentation.akka.instrumentations.SchedulerInstrumentation"
     ]
 
     within = [

instrumentation/kamon-akka/src/common/scala/kamon/instrumentation/akka/instrumentations/SchedulerInstrumentation.scala

@@ -0,0 +1,28 @@
+/* =========================================================================================
+ * Copyright © 2013-2022 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
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+ * either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ * =========================================================================================
+ */
+
+package kamon.instrumentation.akka.instrumentations
+
+import kanela.agent.api.instrumentation.InstrumentationBuilder
+
+class SchedulerInstrumentation extends InstrumentationBuilder {
+
+  /**
+    * Captures the current context when calling `scheduler.scheduleOnce` and restores it when the submitted runnable
+    * runs. This ensures that certain Akka patterns like retry and after work as expected.
+    */
+  onSubTypesOf("akka.actor.Scheduler")
+    .advise(method("scheduleOnce").and(withArgument(1, classOf[Runnable])), classOf[SchedulerRunnableAdvice])
+}

instrumentation/kamon-akka/src/test-common/scala/kamon/instrumentation/akka/SchedulerInstrumentationSpec.scala

@@ -0,0 +1,50 @@
+/* =========================================================================================
+ * Copyright © 2013-2022 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
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+ * either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ * =========================================================================================
+ */
+
+package kamon.instrumentation.akka
+
+import akka.actor.ActorSystem
+import akka.testkit.{ImplicitSender, TestKit}
+import kamon.Kamon
+import kamon.tag.Lookups.plain
+import kamon.testkit.InitAndStopKamonAfterAll
+import org.scalatest.concurrent.Eventually
+import org.scalatest.matchers.should.Matchers
+import org.scalatest.wordspec.AnyWordSpecLike
+
+import scala.concurrent.Promise
+import scala.concurrent.duration._
+
+class SchedulerInstrumentationSpec extends TestKit(ActorSystem("SchedulerInstrumentationSpec")) with AnyWordSpecLike
+    with Matchers with InitAndStopKamonAfterAll with ImplicitSender with Eventually {
+
+
+  "the Akka Scheduler instrumentation" should {
+    "propagate the current context in calls to scheduler.scheduleOnce" in {
+      val contextTagPromise = Promise[String]()
+      val tagValueFuture = contextTagPromise.future
+
+      Kamon.runWithContextTag("key", "one") {
+        system.scheduler.scheduleOnce(100 millis) {
+          contextTagPromise.success(Kamon.currentContext().getTag(plain("key")))
+        } (system.dispatcher)
+      }
+
+      eventually(timeout(5 seconds)) {
+        tagValueFuture.value.get.get shouldBe "one"
+      }
+    }
+  }
+}