feat: tracing support for timer calls (#2044)

aludwiko · franciscolopezsancho · web-flow · commit d2ce1e31657f · 2024-02-22T11:34:08.000+01:00
* feat: tracing support for timer calls

* fixing test

* Apply suggestions from code review

Co-authored-by: Francisco Lopez-Sancho &lt;franciscolopezsancho@gmail.com&gt;

* fixing compilation

---------

Co-authored-by: Francisco Lopez-Sancho &lt;franciscolopezsancho@gmail.com&gt;
diff --git a/samples/java-spring-choreography-saga-quickstart/src/it/java/user/registry/UserCreationIntegrationTest.java b/samples/java-spring-choreography-saga-quickstart/src/it/java/user/registry/UserCreationIntegrationTest.java
@@ -9,7 +9,7 @@
 import java.time.Duration;
 
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.testcontainers.shaded.org.awaitility.Awaitility.await;
+import static org.awaitility.Awaitility.await;
 
 /**
  * This is a skeleton for implementing integration tests for a Kalix application built with the Java SDK.
@@ -108,10 +108,7 @@ public void testUserCreationFailureDueToInvalidInput() throws Exception {
 
     await()
       .ignoreExceptions()
-      // timer will fire in 3 seconds and un-reserve the email
-      // see it/resources/application.conf for the configuration
-      // we only start to polling after 3 seconds to give the timer a chance to fire
-      .between(Duration.ofSeconds(3), Duration.ofSeconds(6))
+      .timeout(Duration.ofSeconds(10)) //3 seconds for the projection lag + 3 seconds for the timer to fire
       .untilAsserted(() -> {
         assertThat(callGetEmailInfo.execute())
           .succeedsWithin(timeout)
diff --git a/sdk/java-sdk-protobuf/src/main/java/kalix/javasdk/action/Action.java b/sdk/java-sdk-protobuf/src/main/java/kalix/javasdk/action/Action.java
@@ -103,7 +103,7 @@ public final TimerScheduler timers() {
     ActionContextImpl impl =
       (ActionContextImpl)
         actionContext("Timers can only be scheduled or cancelled when handling a message.");
-    return new TimerSchedulerImpl(impl.messageCodec(), impl.system());
+    return new TimerSchedulerImpl(impl.messageCodec(), impl.system(), impl.componentCallMetadata());
   }
 
   /**
diff --git a/sdk/java-sdk-protobuf/src/main/scala/kalix/javasdk/impl/telemetry/Telemetry.scala b/sdk/java-sdk-protobuf/src/main/scala/kalix/javasdk/impl/telemetry/Telemetry.scala
@@ -190,7 +190,7 @@ private final class TraceInstrumentation(
       if (logger.isTraceEnabled) logger.trace("`traceparent` found")
 
       val context = openTelemetry.getPropagators.getTextMapPropagator
-        .extract(OtelContext.current(), metadata, otelGetter.asInstanceOf[TextMapGetter[Object]])
+        .extract(OtelContext.current(), metadata, otelGetter)
 
       val span = openTelemetry
         .getTracer("java-sdk")
@@ -217,7 +217,7 @@ private final class TraceInstrumentation(
       if (logger.isTraceEnabled) logger.trace("`traceparent` found")
 
       val context = openTelemetry.getPropagators.getTextMapPropagator
-        .extract(OtelContext.current(), metadata, otelGetter.asInstanceOf[TextMapGetter[Object]])
+        .extract(OtelContext.current(), metadata, otelGetter)
 
       val span = getTracer()
         .spanBuilder(command.name)
diff --git a/sdk/java-sdk-protobuf/src/main/scala/kalix/javasdk/impl/timer/TimerSchedulerImpl.scala b/sdk/java-sdk-protobuf/src/main/scala/kalix/javasdk/impl/timer/TimerSchedulerImpl.scala
@@ -19,22 +19,32 @@ package kalix.javasdk.impl.timer
 import java.time.Duration
 import java.util.concurrent.CompletionStage
 
+import scala.jdk.CollectionConverters._
 import scala.jdk.FutureConverters.FutureOps
 
 import akka.Done
 import akka.actor.ActorSystem
+import akka.grpc.scaladsl.SingleResponseRequestBuilder
+import com.google.protobuf.any.{ Any => ScalaPbAny }
 import com.google.protobuf.duration.{ Duration => ProtoDuration }
 import com.google.protobuf.wrappers.StringValue
-import com.google.protobuf.any.{ Any => ScalaPbAny }
 import kalix.javasdk.DeferredCall
-import kalix.javasdk.impl.{ GrpcClients, GrpcDeferredCall, MessageCodec, RestDeferredCall }
+import kalix.javasdk.Metadata
+import kalix.javasdk.impl.GrpcClients
+import kalix.javasdk.impl.GrpcDeferredCall
+import kalix.javasdk.impl.MessageCodec
+import kalix.javasdk.impl.RestDeferredCall
 import kalix.javasdk.timer.TimerScheduler
 import kalix.timers.timers.Call
 import kalix.timers.timers.SingleTimer
 import kalix.timers.timers.TimerService
+import kalix.timers.timers.TimerServiceClient
 
 /** INTERNAL API */
-private[kalix] final class TimerSchedulerImpl(val messageCodec: MessageCodec, val system: ActorSystem)
+private[kalix] final class TimerSchedulerImpl(
+    val messageCodec: MessageCodec,
+    val system: ActorSystem,
+    val metadata: Metadata)
     extends TimerScheduler {
 
   override def startSingleTimer[I, O](
@@ -48,7 +58,8 @@ private[kalix] final class TimerSchedulerImpl(val messageCodec: MessageCodec, va
       delay: Duration,
       maxRetries: Int,
       deferredCall: DeferredCall[I, O]): CompletionStage[Done] = {
-    val timerServiceClient = GrpcClients(system).getProxyGrpcClient(classOf[TimerService])
+    val timerServiceClient =
+      GrpcClients(system).getProxyGrpcClient(classOf[TimerService]).asInstanceOf[TimerServiceClient]
 
     val call = deferredCall match {
       case grpcDeferredCall: GrpcDeferredCall[I, O] =>
@@ -64,13 +75,22 @@ private[kalix] final class TimerSchedulerImpl(val messageCodec: MessageCodec, va
     }
 
     val singleTimer = SingleTimer(name, Some(call), Some(ProtoDuration(delay)), maxRetries)
-
-    timerServiceClient.addSingle(singleTimer).asJava.thenApply(_ => Done)
+    addHeaders(timerServiceClient.addSingle(), metadata).invoke(singleTimer).asJava.thenApply(_ => Done)
   }
 
   def cancel(name: String): CompletionStage[Done] = {
-    val timerServiceClient = GrpcClients(system).getProxyGrpcClient(classOf[TimerService])
-    timerServiceClient.remove(StringValue(name)).asJava.thenApply(_ => Done)
+    val timerServiceClient =
+      GrpcClients(system).getProxyGrpcClient(classOf[TimerService]).asInstanceOf[TimerServiceClient]
+    addHeaders(timerServiceClient.remove(), metadata).invoke(StringValue(name)).asJava.thenApply(_ => Done)
+  }
+
+  private def addHeaders[I, O](
+      callBuilder: SingleResponseRequestBuilder[I, O],
+      metadata: Metadata): SingleResponseRequestBuilder[I, O] = {
+    metadata.asScala.foldLeft(callBuilder) { case (builder, entry) =>
+      if (entry.isText) builder.addHeader(entry.getKey, entry.getValue)
+      else builder
+    }
   }
 
 }
diff --git a/sdk/java-sdk-protobuf/src/main/scala/kalix/javasdk/impl/workflow/WorkflowImpl.scala b/sdk/java-sdk-protobuf/src/main/scala/kalix/javasdk/impl/workflow/WorkflowImpl.scala
@@ -282,7 +282,7 @@ final class WorkflowImpl(system: ActorSystem, val services: Map[String, Workflow
           val metadata = new MetadataImpl(command.metadata.map(_.entries.toVector).getOrElse(Nil))
 
           val context = new CommandContextImpl(workflowId, command.name, command.id, metadata, system)
-          val timerScheduler = new TimerSchedulerImpl(service.messageCodec, system)
+          val timerScheduler = new TimerSchedulerImpl(service.messageCodec, system, context.componentCallMetadata)
 
           val cmd =
             service.messageCodec.decodeMessage(
@@ -306,7 +306,7 @@ final class WorkflowImpl(system: ActorSystem, val services: Map[String, Workflow
         case Step(executeStep) =>
           val context =
             new CommandContextImpl(workflowId, executeStep.stepName, executeStep.commandId, Metadata.EMPTY, system)
-          val timerScheduler = new TimerSchedulerImpl(service.messageCodec, system)
+          val timerScheduler = new TimerSchedulerImpl(service.messageCodec, system, context.componentCallMetadata)
           val stepResponse =
             try {
               executeStep.userState.foreach { state =>
diff --git a/sdk/scala-sdk-protobuf/src/main/scala/kalix/scalasdk/action/Action.scala b/sdk/scala-sdk-protobuf/src/main/scala/kalix/scalasdk/action/Action.scala
@@ -22,7 +22,9 @@ import scala.concurrent.Future
 import kalix.scalasdk.{ DeferredCall, Metadata, SideEffect }
 import kalix.scalasdk.impl.action.ActionEffectImpl
 import io.grpc.Status
+import kalix.javasdk.impl
 import kalix.javasdk.impl.action.ActionContextImpl
+import kalix.scalasdk.impl.MetadataImpl
 import kalix.scalasdk.impl.action.ScalaActionContextAdapter
 import kalix.scalasdk.timer.TimerScheduler
 import kalix.scalasdk.impl.timer.TimerSchedulerImpl
@@ -259,7 +261,10 @@ abstract class Action {
             s"Incompatible ActionContext instance. Found ${other.getClass}, expecting ${classOf[ActionContextImpl].getName}")
       }
 
-    new TimerSchedulerImpl(javaActionContextImpl.messageCodec, javaActionContextImpl.system)
+    new TimerSchedulerImpl(
+      javaActionContextImpl.messageCodec,
+      javaActionContextImpl.system,
+      MetadataImpl(javaActionContextImpl.metadata.asInstanceOf[impl.MetadataImpl]))
   }
 
   protected final def effects[T]: Action.Effect.Builder =
diff --git a/sdk/scala-sdk-protobuf/src/main/scala/kalix/scalasdk/impl/timer/TimerSchedulerImpl.scala b/sdk/scala-sdk-protobuf/src/main/scala/kalix/scalasdk/impl/timer/TimerSchedulerImpl.scala
@@ -23,20 +23,24 @@ import scala.jdk.DurationConverters.ScalaDurationOps
 
 import akka.Done
 import akka.actor.ActorSystem
+import akka.grpc.scaladsl.SingleResponseRequestBuilder
 import com.google.protobuf.duration.{ Duration => ProtoDuration }
 import com.google.protobuf.wrappers.StringValue
 import kalix.javasdk.impl.GrpcDeferredCall
 import kalix.javasdk.impl.GrpcClients
 import kalix.javasdk.impl.MessageCodec
 import kalix.scalasdk.DeferredCall
+import kalix.scalasdk.Metadata
 import kalix.scalasdk.impl.ScalaDeferredCallAdapter
 import kalix.scalasdk.timer.TimerScheduler
 import kalix.timers.timers.Call
 import kalix.timers.timers.SingleTimer
 import kalix.timers.timers.TimerService
+import kalix.timers.timers.TimerServiceClient
 
 /** INTERNAL API */
-private[kalix] final class TimerSchedulerImpl(messageCodec: MessageCodec, system: ActorSystem) extends TimerScheduler {
+private[kalix] final class TimerSchedulerImpl(messageCodec: MessageCodec, system: ActorSystem, metadata: Metadata)
+    extends TimerScheduler {
 
   override def startSingleTimer[I, O](
       name: String,
@@ -49,7 +53,8 @@ private[kalix] final class TimerSchedulerImpl(messageCodec: MessageCodec, system
       delay: FiniteDuration,
       maxRetries: Int,
       deferredCall: DeferredCall[I, O]): Future[Done] = {
-    val timerServiceClient = GrpcClients(system).getProxyGrpcClient(classOf[TimerService])
+    val timerServiceClient =
+      GrpcClients(system).getProxyGrpcClient(classOf[TimerService]).asInstanceOf[TimerServiceClient]
 
     val deferredCallImpl =
       deferredCall match {
@@ -67,12 +72,25 @@ private[kalix] final class TimerSchedulerImpl(messageCodec: MessageCodec, system
         Some(messageCodec.encodeScala(deferredCall.message)))
 
     val singleTimer = SingleTimer(name, Some(call), Some(ProtoDuration(delay.toJava)))
-    timerServiceClient.addSingle(singleTimer).map(_ => Done)(ExecutionContext.parasitic)
-
+    addHeaders(timerServiceClient.addSingle(), metadata)
+      .invoke(singleTimer)
+      .map(_ => Done)(ExecutionContext.parasitic)
   }
 
   override def cancel(name: String): Future[Done] = {
-    val timerServiceClient = GrpcClients(system).getProxyGrpcClient(classOf[TimerService])
-    timerServiceClient.remove(StringValue(name)).map(_ => Done)(ExecutionContext.parasitic)
+    val timerServiceClient =
+      GrpcClients(system).getProxyGrpcClient(classOf[TimerService]).asInstanceOf[TimerServiceClient]
+    addHeaders(timerServiceClient.remove(), metadata)
+      .invoke(StringValue(name))
+      .map(_ => Done)(ExecutionContext.parasitic)
+  }
+
+  private def addHeaders[I, O](
+      callBuilder: SingleResponseRequestBuilder[I, O],
+      metadata: Metadata): SingleResponseRequestBuilder[I, O] = {
+    metadata.foldLeft(callBuilder) { case (builder, entry) =>
+      if (entry.isText) builder.addHeader(entry.key, entry.value)
+      else builder
+    }
   }
 }
diff --git a/sdk/scala-sdk-protobuf/src/main/scala/kalix/scalasdk/impl/workflow/WorkflowAdapters.scala b/sdk/scala-sdk-protobuf/src/main/scala/kalix/scalasdk/impl/workflow/WorkflowAdapters.scala
@@ -27,10 +27,11 @@ import scala.jdk.OptionConverters._
 import akka.stream.Materializer
 import com.google.protobuf.Descriptors
 import kalix.javasdk
+import kalix.javasdk.impl
 import kalix.javasdk.timer.TimerScheduler
-import kalix.javasdk.workflow.AbstractWorkflow.RecoverStrategy.MaxRetries
 import kalix.scalasdk.impl.InternalContext
 import kalix.scalasdk.impl.MetadataConverters
+import kalix.scalasdk.impl.MetadataImpl
 import kalix.scalasdk.impl.ScalaDeferredCallAdapter
 import kalix.scalasdk.impl.timer.TimerSchedulerImpl
 import kalix.scalasdk.workflow.AbstractWorkflow
@@ -57,7 +58,10 @@ private[scalasdk] final class JavaWorkflowAdapter[S >: Null](scalaSdkWorkflow: A
   override def _internalSetTimerScheduler(timerScheduler: Optional[TimerScheduler]): Unit = {
     scalaSdkWorkflow._internalSetTimerScheduler(timerScheduler.toScala.map {
       case javaTimerScheduler: kalix.javasdk.impl.timer.TimerSchedulerImpl =>
-        new TimerSchedulerImpl(javaTimerScheduler.messageCodec, javaTimerScheduler.system)
+        new TimerSchedulerImpl(
+          javaTimerScheduler.messageCodec,
+          javaTimerScheduler.system,
+          MetadataImpl(javaTimerScheduler.metadata.asInstanceOf[impl.MetadataImpl]))
     })
   }
 

Original file line number	Diff line number	Diff line change
`@@ -103,7 +103,7 @@ public final TimerScheduler timers() {`
`103`	`103`	`ActionContextImpl impl =`
`104`	`104`	`(ActionContextImpl)`
`105`	`105`	`actionContext("Timers can only be scheduled or cancelled when handling a message.");`
`106`		`- return new TimerSchedulerImpl(impl.messageCodec(), impl.system());`
	`106`	`+ return new TimerSchedulerImpl(impl.messageCodec(), impl.system(), impl.componentCallMetadata());`
`107`	`107`	`}`
`108`	`108`
`109`	`109`	`/**`