Merge upstream

core/trace/src/main/scala/org/typelevel/otel4s/trace/SpanContext.scala

@@ -163,19 +163,24 @@ object SpanContext {
     * @param traceFlags
     *   the trace flags of the span context
     *
+    * @param traceState
+    *   the trace state of the span context
+    *
     * @param remote
     *   whether the span is propagated from the remote parent or not
     */
   def create(
       traceId: ByteVector,
       spanId: ByteVector,
       traceFlags: TraceFlags,
+      traceState: TraceState,
       remote: Boolean
   ): SpanContext =
     createInternal(
       traceId = traceId,
       spanId = spanId,
       traceFlags = traceFlags,
+      traceState = traceState,
       remote = remote,
       skipIdValidation = false
     )
@@ -202,6 +207,9 @@ object SpanContext {
     * @param traceFlags
     *   the trace flags of the span context
     *
+    * @param traceState
+    *   the trace state of the span context
+    *
     * @param remote
     *   whether the span is propagated from the remote parent or not
     *
@@ -213,6 +221,7 @@ object SpanContext {
       traceId: ByteVector,
       spanId: ByteVector,
       traceFlags: TraceFlags,
+      traceState: TraceState,
       remote: Boolean,
       skipIdValidation: Boolean
   ): SpanContext = {
@@ -225,6 +234,7 @@ object SpanContext {
         spanId = spanId,
         spanIdHex = spanId.toHex,
         traceFlags = traceFlags,
+        traceState = traceState,
         isRemote = remote,
         isValid = true
       )
@@ -235,6 +245,7 @@ object SpanContext {
         spanId = SpanId.Invalid,
         spanIdHex = SpanId.Invalid.toHex,
         traceFlags = traceFlags,
+        traceState = traceState,
         isRemote = remote,
         isValid = false
       )
@@ -251,6 +262,7 @@ object SpanContext {
       traceId: ByteVector,
       spanId: ByteVector,
       traceFlags: TraceFlags,
+      traceState: TraceState,
       remote: Boolean,
       isValid: Boolean
   ): Delegate[A] =
@@ -261,6 +273,7 @@ object SpanContext {
       spanId = spanId,
       spanIdHex = spanId.toHex,
       traceFlags = traceFlags,
+      traceState = traceState,
       isRemote = remote,
       isValid = isValid,
     )
@@ -272,7 +285,13 @@ object SpanContext {
 
   implicit val spanContextShow: Show[SpanContext] =
     Show.show { ctx =>
-      show"SpanContext{traceId=${ctx.traceIdHex}, spanId=${ctx.spanIdHex}, traceFlags=${ctx.traceFlags}, remote=${ctx.isRemote}, valid=${ctx.isValid}}"
+      show"SpanContext{" +
+        show"traceId=${ctx.traceIdHex}, " +
+        show"spanId=${ctx.spanIdHex}, " +
+        show"traceFlags=${ctx.traceFlags}, " +
+        show"traceState=${ctx.traceState}, " +
+        show"remote=${ctx.isRemote}, " +
+        show"valid=${ctx.isValid}}"
     }
 
   private final case class SpanContextImpl(

java/trace/src/main/scala/org/typelevel/otel4s/java/trace/SpanBackendImpl.scala

@@ -104,7 +104,7 @@ private[java] class SpanBackendImpl[F[_]: Sync](
 
 private[java] object SpanBackendImpl {
   def fromJSpan[F[_]: Sync](jSpan: JSpan): SpanBackendImpl[F] =
-    new SpanBackendImpl(jSpan, WrappedSpanContext(jSpan.getSpanContext))
+    new SpanBackendImpl(jSpan, WrappedSpanContext.wrap(jSpan.getSpanContext))
 
   private def toJStatus(status: Status): JStatusCode =
     status match {

sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/SdkSpanBuilder.scala

@@ -37,6 +37,7 @@ import org.typelevel.otel4s.trace.SpanFinalizer
 import org.typelevel.otel4s.trace.SpanKind
 import org.typelevel.otel4s.trace.SpanOps
 import org.typelevel.otel4s.trace.TraceFlags
+import org.typelevel.otel4s.trace.TraceState
 
 import scala.concurrent.duration.FiniteDuration
 
@@ -163,10 +164,15 @@ private[trace] final case class SdkSpanBuilder[F[_]: Temporal](
           if (samplingDecision.isSampled) TraceFlags.Sampled
           else TraceFlags.Default
 
+        val traceState = parentSpanContext.fold(TraceState.empty) { ctx =>
+          samplingResult.traceStateUpdater.update(ctx.traceState)
+        }
+
         val spanContext = SpanContext.createInternal(
           traceId = traceId,
           spanId = spanId,
           traceFlags = traceFlags,
+          traceState = traceState,
           remote = false,
           skipIdValidation = tracerSharedState.idGenerator.canSkipIdValidation
         )

sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/SdkTracer.scala

@@ -25,6 +25,7 @@ import org.typelevel.otel4s.context.propagation.TextMapGetter
 import org.typelevel.otel4s.context.propagation.TextMapUpdater
 import org.typelevel.otel4s.sdk.common.InstrumentationScopeInfo
 import org.typelevel.otel4s.sdk.context.Context
+import org.typelevel.otel4s.trace.Span
 import org.typelevel.otel4s.trace.SpanBuilder
 import org.typelevel.otel4s.trace.SpanContext
 import org.typelevel.otel4s.trace.Tracer
@@ -41,6 +42,9 @@ final class SdkTracer[F[_]: Temporal] private[trace] (
   def currentSpanContext: F[Option[SpanContext]] =
     scope.current.map(current => current.filter(_.isValid))
 
+  def currentSpanOrNoop: F[Span[F]] =
+    ???
+
   def spanBuilder(name: String): SpanBuilder[F] =
     new SdkSpanBuilder[F](name, scopeInfo, sharedState, scope)
 

sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/propagation/W3CTraceContextPropagator.scala

@@ -26,6 +26,7 @@ import org.typelevel.otel4s.trace.SpanContext
 import org.typelevel.otel4s.trace.SpanContext.SpanId
 import org.typelevel.otel4s.trace.SpanContext.TraceId
 import org.typelevel.otel4s.trace.TraceFlags
+import org.typelevel.otel4s.trace.TraceState
 
 import scala.util.matching.Regex
 
@@ -79,7 +80,8 @@ object W3CTraceContextPropagator extends TextMapPropagator[Context] {
           SpanId.fromHex(spanIdHex),
           TraceFlags.fromHex(traceFlagsHex)
         ).mapN { (traceId, spanId, traceFlags) =>
-          SpanContext.create(traceId, spanId, traceFlags, remote = true)
+          val state = TraceState.empty
+          SpanContext.create(traceId, spanId, traceFlags, state, remote = true)
         }
 
       case _ =>

sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/samplers/SamplingResult.scala

@@ -20,8 +20,12 @@ import cats.Hash
 import cats.Show
 import cats.syntax.show._
 import org.typelevel.otel4s.sdk.Attributes
+import org.typelevel.otel4s.trace.TraceState
 
 /** Sampling result returned by [[Sampler.shouldSample]].
+  *
+  * @see
+  *   [[https://opentelemetry.io/docs/specs/otel/trace/sdk/#shouldsample]]
   */
 sealed trait SamplingResult {
 
@@ -30,10 +34,21 @@ sealed trait SamplingResult {
     */
   def decision: SamplingDecision
 
-  /** The attributes that will be attached to the span
+  /** The attributes that will be attached to the span.
     */
   def attributes: Attributes
 
+  /** A modifier of the parent's TraceState.
+    *
+    * It may return the same trace state that was provided originally, or an
+    * updated one.
+    *
+    * '''Note''': If an empty trace state is returned, the trace state will be
+    * cleared, so samplers should normally return the passed-in trace state if
+    * they do not intend to change it.
+    */
+  def traceStateUpdater: SamplingResult.TraceStateUpdater
+
   override final def hashCode(): Int =
     Hash[SamplingResult].hash(this)
 
@@ -49,40 +64,106 @@ sealed trait SamplingResult {
 
 object SamplingResult {
 
+  /** A modifier of the parent's TraceState.
+    *
+    * It may return the same trace state that was provided originally, or an
+    * updated one.
+    *
+    * '''Note''': If an empty trace state is returned, the trace state will be
+    * cleared, so samplers should normally return the passed-in trace state if
+    * they do not intend to change it.
+    */
+  sealed trait TraceStateUpdater {
+    def update(state: TraceState): TraceState
+  }
+
+  object TraceStateUpdater {
+
+    /** Always returns the given trace state without modifications. */
+    case object Identity extends TraceStateUpdater {
+      def update(state: TraceState): TraceState = state
+    }
+
+    /** Returns the given trace state modified by the `modify` function. */
+    final case class Modifier(modify: TraceState => TraceState)
+        extends TraceStateUpdater {
+      def update(state: TraceState): TraceState =
+        modify(state)
+    }
+
+    /** Always returns the `const` state. No matter the input. */
+    final case class Const(const: TraceState) extends TraceStateUpdater {
+      def update(state: TraceState): TraceState = const
+    }
+
+    implicit val traceStateUpdaterHash: Hash[TraceStateUpdater] =
+      Hash.fromUniversalHashCode
+
+    implicit val traceStateUpdaterShow: Show[TraceStateUpdater] =
+      Show.show {
+        case Identity     => "Identity"
+        case Modifier(_)  => "Modifier(f)"
+        case Const(const) => show"Const($const)"
+      }
+  }
+
   val RecordAndSample: SamplingResult =
-    SamplingResultImpl(SamplingDecision.RecordAndSample, Attributes.Empty)
+    SamplingResultImpl(
+      SamplingDecision.RecordAndSample,
+      Attributes.Empty,
+      TraceStateUpdater.Identity
+    )
 
   val RecordOnly: SamplingResult =
-    SamplingResultImpl(SamplingDecision.RecordOnly, Attributes.Empty)
+    SamplingResultImpl(
+      SamplingDecision.RecordOnly,
+      Attributes.Empty,
+      TraceStateUpdater.Identity
+    )
 
   val Drop: SamplingResult =
-    SamplingResultImpl(SamplingDecision.Drop, Attributes.Empty)
+    SamplingResultImpl(
+      SamplingDecision.Drop,
+      Attributes.Empty,
+      TraceStateUpdater.Identity
+    )
 
-  def create(decision: SamplingDecision): SamplingResult =
+  def apply(decision: SamplingDecision): SamplingResult =
     decision match {
       case SamplingDecision.RecordAndSample => RecordAndSample
       case SamplingDecision.RecordOnly      => RecordOnly
       case SamplingDecision.Drop            => Drop
     }
 
-  def create(
+  def apply(
       decision: SamplingDecision,
       attributes: Attributes
   ): SamplingResult =
-    if (attributes.isEmpty) create(decision)
-    else SamplingResultImpl(decision, attributes)
+    if (attributes.isEmpty) apply(decision)
+    else SamplingResultImpl(decision, attributes, TraceStateUpdater.Identity)
+
+  def apply(
+      decision: SamplingDecision,
+      attributes: Attributes,
+      traceStateUpdater: TraceStateUpdater
+  ): SamplingResult =
+    if (traceStateUpdater == TraceStateUpdater.Identity)
+      apply(decision, attributes)
+    else
+      SamplingResultImpl(decision, attributes, traceStateUpdater)
 
   implicit val samplingResultHash: Hash[SamplingResult] =
-    Hash.by(r => (r.decision, r.attributes))
+    Hash.by(r => (r.decision, r.attributes, r.traceStateUpdater))
 
   implicit val samplingResultShow: Show[SamplingResult] =
     Show.show { r =>
-      show"SamplingResult{decision=${r.decision}, attributes=${r.attributes}}"
+      show"SamplingResult{decision=${r.decision}, attributes=${r.attributes}, traceStateUpdater=${r.traceStateUpdater}}"
     }
 
   private final case class SamplingResultImpl(
       decision: SamplingDecision,
-      attributes: Attributes
+      attributes: Attributes,
+      traceStateUpdater: TraceStateUpdater
   ) extends SamplingResult
 
 }

sdk/trace/src/test/scala/org/typelevel/otel4s/sdk/trace/propagation/W3CTraceContextPropagatorSuite.scala

@@ -21,6 +21,7 @@ import org.typelevel.otel4s.sdk.context.Context
 import org.typelevel.otel4s.sdk.trace.SdkTraceScope
 import org.typelevel.otel4s.trace.SpanContext
 import org.typelevel.otel4s.trace.TraceFlags
+import org.typelevel.otel4s.trace.TraceState
 import scodec.bits.ByteVector
 
 class W3CTraceContextPropagatorSuite extends FunSuite {
@@ -32,6 +33,7 @@ class W3CTraceContextPropagatorSuite extends FunSuite {
   private val spanId = ByteVector.fromValidHex(spanIdHex)
 
   private val flags = List(TraceFlags.Sampled, TraceFlags.Default)
+  private val state = TraceState.empty
 
   private val propagator = W3CTraceContextPropagator
 
@@ -47,7 +49,7 @@ class W3CTraceContextPropagatorSuite extends FunSuite {
   test("inject context info") {
     flags.foreach { flag =>
       val spanContext =
-        SpanContext.create(traceId, spanId, flag, remote = false)
+        SpanContext.create(traceId, spanId, flag, state, remote = false)
       val ctx = SdkTraceScope.storeInContext(Context.root, spanContext)
       val result = propagator.inject(ctx, Map.empty[String, String])
 
@@ -67,12 +69,13 @@ class W3CTraceContextPropagatorSuite extends FunSuite {
   test("extract span context") {
     flags.foreach { flag =>
       val spanContext =
-        SpanContext.create(traceId, spanId, flag, remote = false)
+        SpanContext.create(traceId, spanId, flag, state, remote = false)
       val carrier = Map("traceparent" -> toTraceParent(spanContext))
 
       val ctx = propagator.extract(Context.root, carrier)
 
-      val expected = SpanContext.create(traceId, spanId, flag, remote = true)
+      val expected =
+        SpanContext.create(traceId, spanId, flag, state, remote = true)
 
       assertEquals(SdkTraceScope.fromContext(ctx), Some(expected))
     }