feat: Improved custom tracing span API

akka-javasdk-testkit/src/main/scala/akka/javasdk/testkit/impl/TestKitCommandContextTimed.scala

@@ -5,14 +5,13 @@
 package akka.javasdk.testkit.impl
 
 import akka.javasdk.Metadata
+import akka.javasdk.Tracing
 import akka.javasdk.impl.InternalContext
 import akka.javasdk.testkit.MockRegistry
 import akka.javasdk.timedaction.CommandContext
-import io.opentelemetry.api.OpenTelemetry
-import io.opentelemetry.api.trace.Tracer
 
 /**
- * INTERNAL API Used by the generated testkit
+ * INTERNAL API Used by the testkit
  */
 final class TestKitCommandContextTimed(metadata: Metadata, mockRegistry: MockRegistry = MockRegistry.EMPTY)
     extends AbstractTestKitContext(mockRegistry)
@@ -29,5 +28,5 @@ final class TestKitCommandContextTimed(metadata: Metadata, mockRegistry: MockReg
 
   override def metadata() = metadata
 
-  override def getTracer: Tracer = OpenTelemetry.noop().getTracer("noop")
+  override def tracing(): Tracing = TestKitTracing
 }

akka-javasdk-testkit/src/main/scala/akka/javasdk/testkit/impl/TestKitEventSourcedEntityCommandContext.scala

@@ -5,6 +5,7 @@
 package akka.javasdk.testkit.impl
 
 import akka.javasdk.Metadata
+import akka.javasdk.Tracing
 import akka.javasdk.eventsourcedentity.CommandContext
 import akka.javasdk.impl.InternalContext
 
@@ -22,6 +23,8 @@ final class TestKitEventSourcedEntityCommandContext(
     this(metadata = metadata, commandName = "stubCommandName")
   }
 
+  override def tracing(): Tracing = TestKitTracing
+
 }
 
 object TestKitEventSourcedEntityCommandContext {

akka-javasdk-testkit/src/main/scala/akka/javasdk/testkit/impl/TestKitEventSourcedEntityContext.scala

@@ -8,7 +8,7 @@ import akka.javasdk.eventsourcedentity.EventSourcedEntityContext
 import akka.javasdk.testkit.MockRegistry
 
 /**
- * INTERNAL API Used by the generated testkit
+ * INTERNAL API Used by the testkit
  */
 final class TestKitEventSourcedEntityContext(
     override val entityId: String,

akka-javasdk-testkit/src/main/scala/akka/javasdk/testkit/impl/TestKitEventSourcedEntityEventContext.scala

@@ -7,7 +7,7 @@ package akka.javasdk.testkit.impl
 import akka.javasdk.eventsourcedentity.EventContext
 
 /**
- * INTERNAL API Used by the generated testkit
+ * INTERNAL API Used by the testkit
  */
 final class TestKitEventSourcedEntityEventContext extends EventContext {
   override def entityId = "testkit-entity-id"

akka-javasdk-testkit/src/main/scala/akka/javasdk/testkit/impl/TestKitKeyValueEntityCommandContext.scala

@@ -5,12 +5,13 @@
 package akka.javasdk.testkit.impl
 
 import akka.javasdk.Metadata
+import akka.javasdk.Tracing
 import akka.javasdk.keyvalueentity.CommandContext
 import akka.javasdk.keyvalueentity.KeyValueEntityContext
 import akka.javasdk.testkit.MockRegistry
 
 /**
- * INTERNAL API Used by the generated testkit
+ * INTERNAL API Used by the testkit
  */
 final class TestKitKeyValueEntityCommandContext(
     override val entityId: String,
@@ -26,4 +27,6 @@ final class TestKitKeyValueEntityCommandContext(
     this(entityId = entityId, metadata = metadata, commandName = "stubCommandName")
   }
 
+  override def tracing(): Tracing = TestKitTracing
+
 }

akka-javasdk-testkit/src/main/scala/akka/javasdk/testkit/impl/TestKitKeyValueEntityContext.scala

@@ -8,7 +8,7 @@ import akka.javasdk.keyvalueentity.KeyValueEntityContext
 import akka.javasdk.testkit.MockRegistry
 
 /**
- * INTERNAL API Used by the generated testkit
+ * INTERNAL API Used by the testkit
  */
 final class TestKitKeyValueEntityContext(override val entityId: String, mockRegistry: MockRegistry = MockRegistry.EMPTY)
     extends AbstractTestKitContext(mockRegistry)

akka-javasdk-testkit/src/main/scala/akka/javasdk/testkit/impl/TestKitTracing.scala

@@ -0,0 +1,19 @@
+/*
+ * Copyright (C) 2021-2024 Lightbend Inc. <https://www.lightbend.com>
+ */
+
+package akka.javasdk.testkit.impl
+
+import akka.javasdk.Tracing
+import io.opentelemetry.api.trace.Span
+
+import java.util.Optional
+
+/**
+ * INTERNAL API
+ */
+object TestKitTracing extends Tracing {
+
+  override def startSpan(name: String): Optional[Span] = Optional.empty()
+
+}

akka-javasdk/src/main/java/akka/javasdk/Tracing.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright (C) 2021-2024 Lightbend Inc. <https://www.lightbend.com>
+ */
+
+package akka.javasdk;
+
+import akka.annotation.DoNotInherit;
+import io.opentelemetry.api.trace.Span;
+
+import java.util.Optional;
+
+/**
+ * Factory for manually creating open telemetry spans in addition to those automatically provided by
+ * the runtime and SDK.
+ *
+ * <p>Not for user extension. Is injected into constructors of supported components by the SDK
+ */
+@DoNotInherit
+public interface Tracing {
+  /**
+   * If tracing is enabled, create and start a new custom span with the given name, setting a parent
+   * for the span is done automatically so that the span is a child of the incoming request or
+   * component call.
+   *
+   * @return Optional of the span if tracing is enabled, empty option if tracing is not enabled.
+   */
+  Optional<Span> startSpan(String name);
+}

akka-javasdk/src/main/java/akka/javasdk/consumer/MessageContext.java

@@ -6,6 +6,7 @@
 
 import akka.javasdk.CloudEvent;
 import akka.javasdk.MetadataContext;
+import akka.javasdk.Tracing;
 import io.opentelemetry.api.trace.Tracer;
 
 import java.util.Optional;
@@ -19,11 +20,6 @@ public interface MessageContext extends MetadataContext {
    */
   Optional<String> eventSubject();
 
-  /**
-   * Get an OpenTelemetry tracer for the current message. This will allow for building and automatic
-   * exporting of spans.
-   *
-   * @return A tracer for the current message, if tracing is configured. Otherwise, a noops tracer.
-   */
-  Tracer getTracer();
+  /** Access to tracing for custom app specific tracing. */
+  Tracing tracing();
 }

akka-javasdk/src/main/java/akka/javasdk/eventsourcedentity/CommandContext.java

@@ -5,6 +5,7 @@
 package akka.javasdk.eventsourcedentity;
 
 import akka.javasdk.MetadataContext;
+import akka.javasdk.Tracing;
 
 /** An event sourced command context. */
 public interface CommandContext extends MetadataContext {
@@ -35,4 +36,7 @@ public interface CommandContext extends MetadataContext {
    * @return The entity id.
    */
   String entityId();
+
+  /** Access to tracing for custom app specific tracing. */
+  Tracing tracing();
 }

akka-javasdk/src/main/java/akka/javasdk/http/RequestContext.java

@@ -9,8 +9,6 @@
 import akka.javasdk.JwtClaims;
 import akka.javasdk.Principals;
 
-import java.util.Optional;
-
 /**
  * Not for user extension, can be injected as constructor parameter into HTTP endpoint components
  */

akka-javasdk/src/main/java/akka/javasdk/keyvalueentity/CommandContext.java

@@ -5,6 +5,7 @@
 package akka.javasdk.keyvalueentity;
 
 import akka.javasdk.MetadataContext;
+import akka.javasdk.Tracing;
 
 /** A value based entity command context. */
 public interface CommandContext extends MetadataContext {
@@ -29,4 +30,7 @@ public interface CommandContext extends MetadataContext {
    * @return The entity id.
    */
   String entityId();
+
+  /** Access to tracing for custom app specific tracing. */
+  Tracing tracing();
 }

akka-javasdk/src/main/java/akka/javasdk/timedaction/CommandContext.java

@@ -5,16 +5,12 @@
 package akka.javasdk.timedaction;
 
 import akka.javasdk.MetadataContext;
+import akka.javasdk.Tracing;
 import io.opentelemetry.api.trace.Tracer;
 
 /** Context for action calls. */
 public interface CommandContext extends MetadataContext {
 
-  /**
-   * Get an OpenTelemetry tracer for the current action. This will allow for building and automatic
-   * exporting of spans.
-   *
-   * @return A tracer for the current action, if tracing is configured. Otherwise, a noops tracer.
-   */
-  Tracer getTracer();
+  /** Access to tracing for custom app specific tracing. */
+  Tracing tracing();
 }

akka-javasdk/src/main/java/akka/javasdk/workflow/CommandContext.java

@@ -5,6 +5,7 @@
 package akka.javasdk.workflow;
 
 import akka.javasdk.MetadataContext;
+import akka.javasdk.Tracing;
 
 /** A value based workflow command context. */
 public interface CommandContext extends MetadataContext {
@@ -29,4 +30,7 @@ public interface CommandContext extends MetadataContext {
    * @return The workflow id.
    */
   String workflowId();
+
+  /** Access to tracing for custom app specific tracing. */
+  Tracing tracing();
 }

akka-javasdk/src/main/scala/akka/javasdk/impl/SdkRunner.scala

@@ -58,7 +58,10 @@ import akka.javasdk.view.View
 import akka.javasdk.workflow.Workflow
 import akka.javasdk.workflow.WorkflowContext
 import akka.javasdk.JwtClaims
+import akka.javasdk.Tracing
 import akka.javasdk.impl.http.JwtClaimsImpl
+import akka.javasdk.impl.telemetry.SpanTracingImpl
+import akka.javasdk.impl.telemetry.TraceInstrumentation
 import akka.runtime.sdk.spi.ComponentClients
 import akka.runtime.sdk.spi.HttpEndpointConstructionContext
 import akka.runtime.sdk.spi.HttpEndpointDescriptor
@@ -75,7 +78,7 @@ import com.typesafe.config.Config
 import com.typesafe.config.ConfigFactory
 import io.opentelemetry.api.trace.Span
 import io.opentelemetry.api.trace.Tracer
-import io.opentelemetry.context.Context
+import io.opentelemetry.context.{ Context => OtelContext }
 import kalix.protocol.action.Actions
 import kalix.protocol.discovery.Discovery
 import kalix.protocol.event_sourced_entity.EventSourcedEntities
@@ -267,6 +270,8 @@ private final class Sdk(
   private val applicationConfig = ApplicationConfig(system).getConfig
   private val sdkSettings = Settings(applicationConfig.getConfig("akka.javasdk"))
 
+  private val sdkTracerFactory = () => tracerFactory(TraceInstrumentation.InstrumentationScopeName)
+
   private val httpClientProvider = new HttpClientProviderImpl(
     system,
     None,
@@ -351,7 +356,6 @@ private final class Sdk(
     case h if h == classOf[HttpClientProvider] => httpClientProvider(span)
     case t if t == classOf[TimerScheduler]     => timerScheduler(span)
     case m if m == classOf[Materializer]       => sdkMaterializer
-    case s if s == classOf[Span]               => span.getOrElse(Span.current())
   }
 
   // FIXME mixing runtime config with sdk with user project config is tricky
@@ -380,7 +384,7 @@ private final class Sdk(
           actionAndConsumerServices,
           runtimeComponentClients.timerClient,
           sdkExecutionContext,
-          tracerFactory))
+          sdkTracerFactory))
     }
 
     services.groupBy(_._2.getClass).foreach {
@@ -393,23 +397,23 @@ private final class Sdk(
             eventSourcedServices,
             sdkSettings,
             sdkDispatcherName,
-            tracerFactory)
+            sdkTracerFactory)
         eventSourcedEntitiesEndpoint = Some(eventSourcedImpl)
 
       case (serviceClass, entityServices: Map[String, KeyValueEntityService[_, _]] @unchecked)
           if serviceClass == classOf[KeyValueEntityService[_, _]] =>
-        valueEntitiesEndpoint =
-          Some(new KeyValueEntitiesImpl(classicSystem, entityServices, sdkSettings, sdkDispatcherName, tracerFactory))
+        valueEntitiesEndpoint = Some(
+          new KeyValueEntitiesImpl(classicSystem, entityServices, sdkSettings, sdkDispatcherName, sdkTracerFactory))
 
       case (serviceClass, workflowServices: Map[String, WorkflowService[_, _]] @unchecked)
           if serviceClass == classOf[WorkflowService[_, _]] =>
         workflowEntitiesEndpoint = Some(
           new WorkflowImpl(
-            classicSystem,
             workflowServices,
             runtimeComponentClients.timerClient,
             sdkExecutionContext,
-            sdkDispatcherName))
+            sdkDispatcherName,
+            sdkTracerFactory))
 
       case (serviceClass, _: Map[String, TimedActionService[_]] @unchecked)
           if serviceClass == classOf[TimedActionService[_]] =>
@@ -571,6 +575,8 @@ private final class Sdk(
                       "There are no JWT claims defined but trying accessing the JWT claims. The class or the method needs to be annotated with @JWT.")
                 }
             }
+
+          case c if c == classOf[Tracing] => new SpanTracingImpl(context.openTelemetrySpan, sdkTracerFactory)
         }
       }
   }
@@ -652,7 +658,7 @@ private final class Sdk(
   private def httpClientProvider(openTelemetrySpan: Option[Span]): HttpClientProvider =
     openTelemetrySpan match {
       case None       => httpClientProvider
-      case Some(span) => httpClientProvider.withTraceContext(Context.current().`with`(span))
+      case Some(span) => httpClientProvider.withTraceContext(OtelContext.current().`with`(span))
     }
 
 }