feat: Improved custom tracing span API (#10)

* feat: Improved custom tracing span API

* Drop the tracing context, provide access to parent span if needed, update docs

* Improved javadoc

* review feedback

---------

Co-authored-by: Johan Andrén <[email protected]>

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,20 @@
+/*
+ * 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()
+
+  override def parentSpan(): Optional[Span] = Optional.empty()
+}

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

@@ -213,13 +213,6 @@ public interface Metadata extends Iterable<Metadata.MetadataEntry> {
    */
   CloudEvent asCloudEvent(String id, URI source, String type);
 
-
-  /**
-   * Get the trace context associated with this request metadata.
-   * @return The trace context.
-   */
-  TraceContext traceContext();
-
   /**
    * Merge the given Metadata entries with this Metadata. If the same key is present in both, both values will be kept.
    *

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

@@ -0,0 +1,38 @@
+/*
+ * 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. Injectable into endpoint constructors or available through component
+ * command contexts.
+ */
+@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);
+
+  /**
+   * If tracing is enabled, this returns the current parent span, to use for propagating trace
+   * parent through third party integrations. This span should only be used for observing, ending it
+   * or marking it as failed etc. is managed by the SDK and the runtime.
+   *
+   * @see {{@link #startSpan(String)}} for creating a custom span tied to some logic in a service.
+   */
+  Optional<Span> parentSpan();
+}

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

@@ -8,8 +8,7 @@
 import akka.javasdk.Context;
 import akka.javasdk.JwtClaims;
 import akka.javasdk.Principals;
-
-import java.util.Optional;
+import akka.javasdk.Tracing;
 
 /**
  * Not for user extension, can be injected as constructor parameter into HTTP endpoint components
@@ -26,4 +25,7 @@ public interface RequestContext extends Context {
 
   /** @return The JWT claims, if any, associated with this request. */
   JwtClaims getJwtClaims();
+
+  /** Access to tracing for custom app specific tracing. */
+  Tracing tracing();
 }

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/MetadataImpl.scala

@@ -18,7 +18,6 @@ import scala.jdk.OptionConverters._
 import akka.annotation.InternalApi
 import akka.javasdk.CloudEvent
 import akka.javasdk.Metadata
-import akka.javasdk.TraceContext
 import akka.javasdk.impl.telemetry.Telemetry
 import akka.javasdk.impl.telemetry.Telemetry.metadataGetter
 import com.google.protobuf.ByteString
@@ -207,22 +206,16 @@ private[javasdk] class MetadataImpl private (val entries: Seq[MetadataEntry]) ex
 
   override def asMetadata(): Metadata = this
 
-  override lazy val traceContext: TraceContext = new TraceContext {
-    override def asOpenTelemetryContext(): OtelContext = W3CTraceContextPropagator
+  lazy val traceId: Option[String] = {
+    val otelContext = W3CTraceContextPropagator
       .getInstance()
       .extract(OtelContext.current(), asMetadata(), metadataGetter)
 
-    override def traceId(): Optional[String] = {
-      Span.fromContext(asOpenTelemetryContext()).getSpanContext.getTraceId match {
-        case "00000000000000000000000000000000" =>
-          Optional.empty() // when no traceId returns io.opentelemetry.api.trace.TraceId.INVALID
-        case traceId => Some(traceId).toJava
-      }
+    Span.fromContext(otelContext).getSpanContext.getTraceId match {
+      case "00000000000000000000000000000000" =>
+        None // when no traceId returns io.opentelemetry.api.trace.TraceId.INVALID
+      case traceId => Some(traceId)
     }
-
-    override def traceParent(): Optional[String] = getScala(Telemetry.TRACE_PARENT_KEY).toJava
-
-    override def traceState(): Optional[String] = getScala(Telemetry.TRACE_STATE_KEY).toJava
   }
 
   override def merge(other: Metadata): Metadata = {