Merge pull request #782 from iRevive/sdk-trace/sampler-f

sdk-trace: make `Sampler` effectful
typelevel · Sep 24, 2024 · ff25785 · ff25785
2 parents a5bfda8 + 6b4b382
commit ff25785
Show file tree

Hide file tree

Showing 19 changed files with 243 additions and 229 deletions.
diff --git a/sdk/all/src/main/scala/org/typelevel/otel4s/sdk/OpenTelemetrySdk.scala b/sdk/all/src/main/scala/org/typelevel/otel4s/sdk/OpenTelemetrySdk.scala
@@ -279,9 +279,7 @@ object OpenTelemetrySdk {
         * @param configurer
         *   the configurer to add
         */
-      def addSamplerConfigurer(
-          configurer: AutoConfigure.Named[F, Sampler]
-      ): Builder[F]
+      def addSamplerConfigurer(configurer: AutoConfigure.Named[F, Sampler[F]]): Builder[F]
 
       /** Adds the text map propagator configurer. Can be used to register propagators that aren't included in the SDK.
         *
@@ -330,7 +328,7 @@ object OpenTelemetrySdk {
           AutoConfigure.Named[F, MetricExporter[F]]
         ],
         spanExporterConfigurers: Set[AutoConfigure.Named[F, SpanExporter[F]]],
-        samplerConfigurers: Set[AutoConfigure.Named[F, Sampler]],
+        samplerConfigurers: Set[AutoConfigure.Named[F, Sampler[F]]],
         textMapPropagatorConfigurers: Set[
           AutoConfigure.Named[F, TextMapPropagator[Context]]
         ]
@@ -387,9 +385,7 @@ object OpenTelemetrySdk {
       ): Builder[F] =
         copy(spanExporterConfigurers = spanExporterConfigurers + configurer)
 
-      def addSamplerConfigurer(
-          configurer: AutoConfigure.Named[F, Sampler]
-      ): Builder[F] =
+      def addSamplerConfigurer(configurer: AutoConfigure.Named[F, Sampler[F]]): Builder[F] =
         copy(samplerConfigurers = samplerConfigurers + configurer)
 
       def addTextMapPropagatorConfigurer(

diff --git a/sdk/all/src/test/scala/org/typelevel/otel4s/sdk/OpenTelemetrySdkSuite.scala b/sdk/all/src/test/scala/org/typelevel/otel4s/sdk/OpenTelemetrySdkSuite.scala
@@ -56,7 +56,7 @@ class OpenTelemetrySdkSuite extends CatsEffectSuite {
   private val DefaultSdk =
     sdkToString(
       TelemetryResource.default,
-      sampler = Sampler.parentBased(Sampler.AlwaysOn)
+      sampler = Sampler.parentBased(Sampler.alwaysOn)
     )
 
   private val NoopSdk =
@@ -131,7 +131,7 @@ class OpenTelemetrySdkSuite extends CatsEffectSuite {
       )
     )
 
-    val sampler = Sampler.AlwaysOff
+    val sampler = Sampler.alwaysOff[IO]
 
     OpenTelemetrySdk
       .autoConfigured[IO](
@@ -222,16 +222,16 @@ class OpenTelemetrySdkSuite extends CatsEffectSuite {
       )
     )
 
-    val sampler: Sampler = new Sampler {
+    val sampler: Sampler[IO] = new Sampler[IO] {
       def shouldSample(
           parentContext: Option[SpanContext],
           traceId: ByteVector,
           name: String,
           spanKind: SpanKind,
           attributes: Attributes,
           parentLinks: Vector[LinkData]
-      ): SamplingResult =
-        SamplingResult.Drop
+      ): IO[SamplingResult] =
+        IO.pure(SamplingResult.Drop)
 
       def description: String = "CustomSampler"
     }
@@ -389,7 +389,7 @@ class OpenTelemetrySdkSuite extends CatsEffectSuite {
   private def sdkToString(
       resource: TelemetryResource = TelemetryResource.default,
       spanLimits: SpanLimits = SpanLimits.default,
-      sampler: Sampler = Sampler.parentBased(Sampler.AlwaysOn),
+      sampler: Sampler[IO] = Sampler.parentBased(Sampler.alwaysOn),
       propagators: ContextPropagators[Context] = ContextPropagators.of(
         W3CTraceContextPropagator.default,
         W3CBaggagePropagator.default

diff --git a/sdk/all/src/test/scala/org/typelevel/otel4s/sdk/trace/SdkTracerSuite.scala b/sdk/all/src/test/scala/org/typelevel/otel4s/sdk/trace/SdkTracerSuite.scala
@@ -120,7 +120,7 @@ class SdkTracerSuite extends BaseTracerSuite[Context, Context.Key] {
 
   sdkTest(
     "keep propagating non-recording spans, but don't record them",
-    _.withSampler(Sampler.AlwaysOff)
+    _.withSampler(Sampler.alwaysOff)
   ) { sdk =>
     for {
       tracer <- sdk.provider.get("tracer")

diff --git a/sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/SdkSpanBuilder.scala b/sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/SdkSpanBuilder.scala
@@ -133,10 +133,7 @@ private final case class SdkSpanBuilder[F[_]: Temporal: Console] private (
         .filter(_.isValid)
         .fold(idGenerator.generateTraceId)(ctx => Temporal[F].pure(ctx.traceId))
 
-    def sample(
-        parent: Option[SpanContext],
-        traceId: ByteVector
-    ): SamplingResult =
+    def sample(parent: Option[SpanContext], traceId: ByteVector): F[SamplingResult] =
       tracerSharedState.sampler.shouldSample(
         parentContext = parent,
         traceId = traceId,
@@ -151,8 +148,8 @@ private final case class SdkSpanBuilder[F[_]: Temporal: Console] private (
       spanId <- idGenerator.generateSpanId
       traceId <- genTraceId(parentSpanContext)
 
+      samplingResult <- sample(parentSpanContext, traceId)
       backend <- {
-        val samplingResult = sample(parentSpanContext, traceId)
         val samplingDecision = samplingResult.decision
 
         val traceFlags =

diff --git a/sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/SdkTracerProvider.scala b/sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/SdkTracerProvider.scala
@@ -37,7 +37,7 @@ private class SdkTracerProvider[F[_]: Temporal: Parallel: Console](
     idGenerator: IdGenerator[F],
     resource: TelemetryResource,
     spanLimits: SpanLimits,
-    sampler: Sampler,
+    sampler: Sampler[F],
     propagators: ContextPropagators[Context],
     spanProcessors: List[SpanProcessor[F]],
     traceScope: TraceScope[F, Context],
@@ -126,17 +126,15 @@ object SdkTracerProvider {
       * @param sampler
       *   the [[org.typelevel.otel4s.sdk.trace.samplers.Sampler Sampler]] to use
       */
-    def withSampler(sampler: Sampler): Builder[F]
+    def withSampler(sampler: Sampler[F]): Builder[F]
 
     /** Adds [[org.typelevel.otel4s.context.propagation.TextMapPropagator TextMapPropagator]]s to use for the context
       * propagation.
       *
       * @param propagators
       *   the propagators to add
       */
-    def addTextMapPropagators(
-        propagators: TextMapPropagator[Context]*
-    ): Builder[F]
+    def addTextMapPropagators(propagators: TextMapPropagator[Context]*): Builder[F]
 
     /** Adds a [[org.typelevel.otel4s.sdk.trace.processor.SpanProcessor SpanProcessor]] to the span processing pipeline
       * that will be built.
@@ -167,7 +165,7 @@ object SdkTracerProvider {
       idGenerator = IdGenerator.random,
       resource = TelemetryResource.default,
       spanLimits = SpanLimits.default,
-      sampler = Sampler.parentBased(Sampler.AlwaysOn),
+      sampler = Sampler.parentBased(Sampler.alwaysOn),
       propagators = Nil,
       spanProcessors = Nil
     )
@@ -178,7 +176,7 @@ object SdkTracerProvider {
       idGenerator: IdGenerator[F],
       resource: TelemetryResource,
       spanLimits: SpanLimits,
-      sampler: Sampler,
+      sampler: Sampler[F],
       propagators: List[TextMapPropagator[Context]],
       spanProcessors: List[SpanProcessor[F]]
   ) extends Builder[F] {
@@ -195,7 +193,7 @@ object SdkTracerProvider {
     def withSpanLimits(limits: SpanLimits): Builder[F] =
       copy(spanLimits = limits)
 
-    def withSampler(sampler: Sampler): Builder[F] =
+    def withSampler(sampler: Sampler[F]): Builder[F] =
       copy(sampler = sampler)
 
     def addTextMapPropagators(

diff --git a/sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/SdkTraces.scala b/sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/SdkTraces.scala
@@ -199,9 +199,7 @@ object SdkTraces {
         * @param configurer
         *   the configurer to add
         */
-      def addSamplerConfigurer(
-          configurer: AutoConfigure.Named[F, Sampler]
-      ): Builder[F]
+      def addSamplerConfigurer(configurer: AutoConfigure.Named[F, Sampler[F]]): Builder[F]
 
       /** Adds the text map propagator configurer. Can be used to register propagators that aren't included in the SDK.
         *
@@ -244,7 +242,7 @@ object SdkTraces {
         tracerProviderCustomizer: Customizer[SdkTracerProvider.Builder[F]],
         resourceDetectors: Set[TelemetryResourceDetector[F]],
         exporterConfigurers: Set[AutoConfigure.Named[F, SpanExporter[F]]],
-        samplerConfigurers: Set[AutoConfigure.Named[F, Sampler]],
+        samplerConfigurers: Set[AutoConfigure.Named[F, Sampler[F]]],
         textMapPropagatorConfigurers: Set[
           AutoConfigure.Named[F, TextMapPropagator[Context]]
         ]
@@ -283,9 +281,7 @@ object SdkTraces {
       ): Builder[F] =
         copy(exporterConfigurers = this.exporterConfigurers + configurer)
 
-      def addSamplerConfigurer(
-          configurer: AutoConfigure.Named[F, Sampler]
-      ): Builder[F] =
+      def addSamplerConfigurer(configurer: AutoConfigure.Named[F, Sampler[F]]): Builder[F] =
         copy(samplerConfigurers = this.samplerConfigurers + configurer)
 
       def addTextMapPropagatorConfigurer(

diff --git a/sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/TracerSharedState.scala b/sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/TracerSharedState.scala
@@ -24,6 +24,6 @@ private final case class TracerSharedState[F[_]](
     idGenerator: IdGenerator[F],
     resource: TelemetryResource,
     spanLimits: SpanLimits,
-    sampler: Sampler,
+    sampler: Sampler[F],
     spanProcessor: SpanProcessor[F]
 )
diff --git a/...ce/src/main/scala/org/typelevel/otel4s/sdk/trace/autoconfigure/SamplerAutoConfigure.scala b/...ce/src/main/scala/org/typelevel/otel4s/sdk/trace/autoconfigure/SamplerAutoConfigure.scala
@@ -38,35 +38,32 @@ import org.typelevel.otel4s.sdk.trace.samplers.Sampler
   *   [[https://opentelemetry.io/docs/languages/java/configuration/#sampler]]
   */
 private final class SamplerAutoConfigure[F[_]: MonadThrow](
-    extra: Set[AutoConfigure.Named[F, Sampler]]
-) extends AutoConfigure.WithHint[F, Sampler](
-      "Sampler",
-      SamplerAutoConfigure.ConfigKeys.All
-    ) {
+    extra: Set[AutoConfigure.Named[F, Sampler[F]]]
+) extends AutoConfigure.WithHint[F, Sampler[F]]("Sampler", SamplerAutoConfigure.ConfigKeys.All) {
 
   import SamplerAutoConfigure.ConfigKeys
   import SamplerAutoConfigure.Defaults
 
   private val configurers = {
-    val default: Set[AutoConfigure.Named[F, Sampler]] = Set(
-      AutoConfigure.Named.const("always_on", Sampler.AlwaysOn),
-      AutoConfigure.Named.const("always_off", Sampler.AlwaysOff),
+    val default: Set[AutoConfigure.Named[F, Sampler[F]]] = Set(
+      AutoConfigure.Named.const("always_on", Sampler.alwaysOn),
+      AutoConfigure.Named.const("always_off", Sampler.alwaysOff),
       traceIdRatioSampler("traceidratio")(identity),
       AutoConfigure.Named.const(
         "parentbased_always_on",
-        Sampler.parentBased(Sampler.AlwaysOn)
+        Sampler.parentBased(Sampler.alwaysOn)
       ),
       AutoConfigure.Named.const(
         "parentbased_always_off",
-        Sampler.parentBased(Sampler.AlwaysOff)
+        Sampler.parentBased(Sampler.alwaysOff)
       ),
       traceIdRatioSampler("parentbased_traceidratio")(Sampler.parentBased)
     )
 
     default ++ extra
   }
 
-  def fromConfig(config: Config): Resource[F, Sampler] =
+  def fromConfig(config: Config): Resource[F, Sampler[F]] =
     config.getOrElse(ConfigKeys.Sampler, Defaults.Sampler) match {
       case Right(name) =>
         configurers.find(_.name == name) match {
@@ -89,11 +86,11 @@ private final class SamplerAutoConfigure[F[_]: MonadThrow](
 
   private def traceIdRatioSampler(
       samplerName: String
-  )(make: Sampler => Sampler): AutoConfigure.Named[F, Sampler] =
-    new AutoConfigure.Named[F, Sampler] {
+  )(make: Sampler[F] => Sampler[F]): AutoConfigure.Named[F, Sampler[F]] =
+    new AutoConfigure.Named[F, Sampler[F]] {
       def name: String = samplerName
 
-      def configure(config: Config): Resource[F, Sampler] = {
+      def configure(config: Config): Resource[F, Sampler[F]] = {
         val attempt = config
           .getOrElse(ConfigKeys.SamplerArg, Defaults.Ratio)
           .flatMap { ratio =>
@@ -111,7 +108,6 @@ private final class SamplerAutoConfigure[F[_]: MonadThrow](
           case Right(sampler) => Resource.pure(make(sampler))
           case Left(error)    => Resource.raiseError(error: Throwable)
         }
-
       }
     }
 
@@ -142,25 +138,23 @@ private[sdk] object SamplerAutoConfigure {
     * }}}
     *
     * The following options for `otel.traces.sampler` are supported out of the box:
-    *   - `always_on` - [[Sampler.AlwaysOn]]
+    *   - `always_on` - [[Sampler.alwaysOn]]
     *
-    *   - `always_off` - [[Sampler.AlwaysOff]]
+    *   - `always_off` - [[Sampler.alwaysOff]]
     *
     *   - `traceidratio` - [[Sampler.traceIdRatioBased]], where `otel.traces.sampler.arg` sets the ratio
     *
-    *   - `parentbased_always_on` - [[Sampler.parentBased]] with [[Sampler.AlwaysOn]]
+    *   - `parentbased_always_on` - [[Sampler.parentBased]] with [[Sampler.alwaysOn]]
     *
-    *   - `parentbased_always_off` - [[Sampler.parentBased]] with [[Sampler.AlwaysOff]]
+    *   - `parentbased_always_off` - [[Sampler.parentBased]] with [[Sampler.alwaysOff]]
     *
     *   - `parentbased_traceidratio`- [[Sampler.parentBased]] with [[Sampler.traceIdRatioBased]], where
     *     `otel.traces.sampler.arg` sets the ratio
     *
     * @see
     *   [[https://opentelemetry.io/docs/languages/java/configuration/#sampler]]
     */
-  def apply[F[_]: MonadThrow](
-      extra: Set[AutoConfigure.Named[F, Sampler]]
-  ): AutoConfigure[F, Sampler] =
+  def apply[F[_]: MonadThrow](extra: Set[AutoConfigure.Named[F, Sampler[F]]]): AutoConfigure[F, Sampler[F]] =
     new SamplerAutoConfigure[F](extra)
 
 }
diff --git a/...main/scala/org/typelevel/otel4s/sdk/trace/autoconfigure/TracerProviderAutoConfigure.scala b/...main/scala/org/typelevel/otel4s/sdk/trace/autoconfigure/TracerProviderAutoConfigure.scala
@@ -42,7 +42,7 @@ private final class TracerProviderAutoConfigure[
     resource: TelemetryResource,
     contextPropagators: ContextPropagators[Context],
     customizer: Customizer[SdkTracerProvider.Builder[F]],
-    samplerConfigurers: Set[AutoConfigure.Named[F, Sampler]],
+    samplerConfigurers: Set[AutoConfigure.Named[F, Sampler[F]]],
     exporterConfigurers: Set[AutoConfigure.Named[F, SpanExporter[F]]]
 ) extends AutoConfigure.WithHint[F, TracerProvider[F]](
       "TracerProvider",
@@ -136,7 +136,7 @@ private[sdk] object TracerProviderAutoConfigure {
       resource: TelemetryResource,
       contextPropagators: ContextPropagators[Context],
       tracerProviderBuilderCustomizer: Customizer[SdkTracerProvider.Builder[F]],
-      samplerConfigurers: Set[AutoConfigure.Named[F, Sampler]],
+      samplerConfigurers: Set[AutoConfigure.Named[F, Sampler[F]]],
       exporterConfigurers: Set[AutoConfigure.Named[F, SpanExporter[F]]]
   ): AutoConfigure[F, TracerProvider[F]] =
     new TracerProviderAutoConfigure[F](