Merge pull request #471 from iRevive/sdk-trace/sampler-autoconfigure

sdk-trace: add `SamplerAutoConfigure`

sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/autoconfigure/SamplerAutoConfigure.scala

@@ -0,0 +1,156 @@
+/*
+ * Copyright 2023 Typelevel
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.typelevel.otel4s.sdk.trace.autoconfigure
+
+import cats.MonadThrow
+import cats.effect.Resource
+import cats.syntax.either._
+import org.typelevel.otel4s.sdk.autoconfigure.AutoConfigure
+import org.typelevel.otel4s.sdk.autoconfigure.Config
+import org.typelevel.otel4s.sdk.autoconfigure.ConfigurationError
+import org.typelevel.otel4s.sdk.trace.samplers.Sampler
+
+/** Autoconfigures a [[Sampler]].
+  *
+  * The configuration options:
+  * {{{
+  * | System property         | Environment variable    | Description                                                             |
+  * |-------------------------|-------------------------|-------------------------------------------------------------------------|
+  * | otel.traces.sampler     | OTEL_TRACES_SAMPLER     | The sampler to use for tracing. Defaults to `parentbased_always_on`     |
+  * | otel.traces.sampler.arg | OTEL_TRACES_SAMPLER_ARG | An argument to the configured tracer if supported, for example a ratio. |
+  * }}}
+  *
+  * @see
+  *   [[https://github.com/open-telemetry/opentelemetry-java/blob/main/sdk-extensions/autoconfigure/README.md#sampler]]
+  */
+private final class SamplerAutoConfigure[F[_]: MonadThrow]
+    extends AutoConfigure.WithHint[F, Sampler](
+      "Sampler",
+      SamplerAutoConfigure.ConfigKeys.All
+    ) {
+
+  import SamplerAutoConfigure.ConfigKeys
+  import SamplerAutoConfigure.Defaults
+
+  private val options = Set(
+    "always_on",
+    "always_off",
+    "traceidratio",
+    "parentbased_always_on",
+    "parentbased_always_off",
+    "parentbased_traceidratio"
+  )
+
+  def fromConfig(config: Config): Resource[F, Sampler] = {
+    val sampler = config.getOrElse(ConfigKeys.Sampler, Defaults.Sampler)
+    def traceIdRatioSampler =
+      config
+        .getOrElse(ConfigKeys.SamplerArg, Defaults.Ratio)
+        .flatMap { ratio =>
+          Either
+            .catchNonFatal(Sampler.traceIdRatioBased(ratio))
+            .leftMap { cause =>
+              ConfigurationError(
+                s"[${ConfigKeys.SamplerArg.name}] has invalid ratio [$ratio] - ${cause.getMessage}",
+                cause
+              )
+            }
+        }
+
+    def attempt = sampler.flatMap {
+      case "always_on" =>
+        Right(Sampler.AlwaysOn)
+
+      case "always_off" =>
+        Right(Sampler.AlwaysOff)
+
+      case "traceidratio" =>
+        traceIdRatioSampler
+
+      case "parentbased_always_on" =>
+        Right(Sampler.parentBased(Sampler.AlwaysOn))
+
+      case "parentbased_always_off" =>
+        Right(Sampler.parentBased(Sampler.AlwaysOff))
+
+      case "parentbased_traceidratio" =>
+        traceIdRatioSampler.map(s => Sampler.parentBased(s))
+
+      case other =>
+        Left(
+          ConfigurationError.unrecognized(
+            ConfigKeys.Sampler.name,
+            other,
+            options
+          )
+        )
+    }
+
+    Resource.eval(MonadThrow[F].fromEither(attempt))
+  }
+
+}
+
+private[sdk] object SamplerAutoConfigure {
+
+  private object ConfigKeys {
+    val Sampler: Config.Key[String] = Config.Key("otel.traces.sampler")
+    val SamplerArg: Config.Key[Double] = Config.Key("otel.traces.sampler.arg")
+
+    val All: Set[Config.Key[_]] = Set(Sampler, SamplerArg)
+  }
+
+  private object Defaults {
+    val Sampler = "parentbased_always_on"
+    val Ratio = 1.0
+  }
+
+  /** Autoconfigures a [[Sampler]].
+    *
+    * The configuration options:
+    * {{{
+    * | System property         | Environment variable    | Description                                                             |
+    * |-------------------------|-------------------------|-------------------------------------------------------------------------|
+    * | otel.traces.sampler     | OTEL_TRACES_SAMPLER     | The sampler to use for tracing. Defaults to `parentbased_always_on`     |
+    * | otel.traces.sampler.arg | OTEL_TRACES_SAMPLER_ARG | An argument to the configured tracer if supported, for example a ratio. |
+    * }}}
+    *
+    * Supported options for `otel.traces.sampler` are:
+    *   - `always_on` - [[Sampler.AlwaysOn]]
+    *
+    *   - `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_off` - [[Sampler.parentBased]] with
+    *     [[Sampler.AlwaysOff]]
+    *
+    *   - `parentbased_traceidratio`- [[Sampler.parentBased]] with
+    *     [[Sampler.traceIdRatioBased]], where `otel.traces.sampler.arg` sets
+    *     the ratio
+    *
+    * @see
+    *   [[https://github.com/open-telemetry/opentelemetry-java/blob/main/sdk-extensions/autoconfigure/README.md#sampler]]
+    */
+  def apply[F[_]: MonadThrow]: AutoConfigure[F, Sampler] =
+    new SamplerAutoConfigure[F]
+
+}

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

@@ -30,7 +30,7 @@ import scodec.bits.ByteVector
   * @see
   *   [[https://opentelemetry.io/docs/specs/otel/trace/sdk/#parentbased]]
   */
-private final class ParentBasedSampler private (
+private final case class ParentBasedSampler private (
     root: Sampler,
     remoteParentSampled: Sampler,
     remoteParentNotSampled: Sampler,
@@ -153,7 +153,7 @@ object ParentBasedSampler {
       copy(localParentNotSampled = Some(sampler))
 
     def build: Sampler =
-      new ParentBasedSampler(
+      ParentBasedSampler(
         root,
         remoteParentSampled.getOrElse(Sampler.AlwaysOn),
         remoteParentNotSampled.getOrElse(Sampler.AlwaysOff),

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

@@ -39,7 +39,7 @@ import scodec.bits.ByteVector
   * @see
   *   [[https://opentelemetry.io/docs/specs/otel/trace/sdk/#traceidratiobased]]
   */
-private final class TraceIdRatioBasedSampler private (
+private final case class TraceIdRatioBasedSampler private (
     ratio: Double,
     idUpperBound: Long
 ) extends Sampler {
@@ -91,7 +91,7 @@ private object TraceIdRatioBasedSampler {
       else if (ratio == 1.0) Long.MaxValue
       else (ratio * Long.MaxValue).toLong
 
-    new TraceIdRatioBasedSampler(ratio, idUpperBound)
+    TraceIdRatioBasedSampler(ratio, idUpperBound)
   }
 
 }

sdk/trace/src/test/scala/org/typelevel/otel4s/sdk/trace/autoconfigure/SamplerAutoConfigureSuite.scala

@@ -0,0 +1,186 @@
+/*
+ * Copyright 2023 Typelevel
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.typelevel.otel4s
+package sdk
+package trace
+package autoconfigure
+
+import cats.effect.IO
+import cats.syntax.either._
+import munit.CatsEffectSuite
+import org.typelevel.otel4s.sdk.autoconfigure.Config
+import org.typelevel.otel4s.sdk.trace.samplers.Sampler
+
+class SamplerAutoConfigureSuite extends CatsEffectSuite {
+
+  private val default = Sampler.parentBased(Sampler.AlwaysOn)
+
+  test("load from an empty config - load default") {
+    val config = Config(Map.empty, Map.empty, Map.empty)
+    SamplerAutoConfigure[IO].configure(config).use { sampler =>
+      IO(assertEquals(sampler, default))
+    }
+  }
+
+  test("load from the config (empty string) - load default") {
+    val props = Map("otel.traces.sampler" -> "")
+    val config = Config(props, Map.empty, Map.empty)
+    SamplerAutoConfigure[IO].configure(config).use { sampler =>
+      IO(assertEquals(sampler, default))
+    }
+  }
+
+  test("load from the config - always_on") {
+    val props = Map("otel.traces.sampler" -> "always_on")
+    val config = Config(props, Map.empty, Map.empty)
+    SamplerAutoConfigure[IO].configure(config).use { sampler =>
+      IO(assertEquals(sampler, Sampler.AlwaysOn))
+    }
+  }
+
+  test("load from the config - always_off") {
+    val props = Map("otel.traces.sampler" -> "always_on")
+    val config = Config(props, Map.empty, Map.empty)
+    SamplerAutoConfigure[IO].configure(config).use { sampler =>
+      IO(assertEquals(sampler, Sampler.AlwaysOn))
+    }
+  }
+
+  test("load from the config - traceidratio - use default ratio") {
+    val props = Map("otel.traces.sampler" -> "traceidratio")
+    val config = Config(props, Map.empty, Map.empty)
+    SamplerAutoConfigure[IO].configure(config).use { sampler =>
+      IO(assertEquals(sampler, Sampler.traceIdRatioBased(1.0)))
+    }
+  }
+
+  test("load from the config - traceidratio - use given ratio") {
+    val props = Map(
+      "otel.traces.sampler" -> "traceidratio",
+      "otel.traces.sampler.arg" -> "0.1"
+    )
+    val config = Config(props, Map.empty, Map.empty)
+    SamplerAutoConfigure[IO].configure(config).use { sampler =>
+      IO(assertEquals(sampler, Sampler.traceIdRatioBased(0.1)))
+    }
+  }
+
+  test("load from the config - traceidratio - fail when ratio is incorrect") {
+    val props = Map(
+      "otel.traces.sampler" -> "traceidratio",
+      "otel.traces.sampler.arg" -> "-0.1"
+    )
+    val config = Config(props, Map.empty, Map.empty)
+
+    SamplerAutoConfigure[IO]
+      .configure(config)
+      .use_
+      .attempt
+      .map(_.leftMap(_.getMessage))
+      .assertEquals(
+        Left(
+          """Cannot autoconfigure [Sampler].
+            |Cause: [otel.traces.sampler.arg] has invalid ratio [-0.1] - requirement failed: ratio must be >= 0 and <= 1.0.
+            |Config:
+            |1) `otel.traces.sampler` - traceidratio
+            |2) `otel.traces.sampler.arg` - -0.1""".stripMargin
+        )
+      )
+  }
+
+  test("load from the config - parentbased_always_on") {
+    val props = Map("otel.traces.sampler" -> "parentbased_always_on")
+    val config = Config(props, Map.empty, Map.empty)
+    SamplerAutoConfigure[IO].configure(config).use { sampler =>
+      IO(assertEquals(sampler, Sampler.parentBased(Sampler.AlwaysOn)))
+    }
+  }
+
+  test("load from the config - parentbased_always_off") {
+    val props = Map("otel.traces.sampler" -> "parentbased_always_off")
+    val config = Config(props, Map.empty, Map.empty)
+    SamplerAutoConfigure[IO].configure(config).use { sampler =>
+      IO(assertEquals(sampler, Sampler.parentBased(Sampler.AlwaysOff)))
+    }
+  }
+
+  test("load from the config - parentbased_traceidratio - use default ratio") {
+    val props = Map("otel.traces.sampler" -> "parentbased_traceidratio")
+    val config = Config(props, Map.empty, Map.empty)
+    val expected = Sampler.parentBased(Sampler.traceIdRatioBased(1.0))
+    SamplerAutoConfigure[IO].configure(config).use { sampler =>
+      IO(assertEquals(sampler, expected))
+    }
+  }
+
+  test("load from the config - parentbased_traceidratio - use given ratio") {
+    val props = Map(
+      "otel.traces.sampler" -> "parentbased_traceidratio",
+      "otel.traces.sampler.arg" -> "0.1"
+    )
+    val config = Config(props, Map.empty, Map.empty)
+    val expected = Sampler.parentBased(Sampler.traceIdRatioBased(0.1))
+    SamplerAutoConfigure[IO].configure(config).use { sampler =>
+      IO(assertEquals(sampler, expected))
+    }
+  }
+
+  test(
+    "load from the config - parentbased_traceidratio - fail when ratio is incorrect"
+  ) {
+    val props = Map(
+      "otel.traces.sampler" -> "parentbased_traceidratio",
+      "otel.traces.sampler.arg" -> "-0.1"
+    )
+    val config = Config(props, Map.empty, Map.empty)
+
+    SamplerAutoConfigure[IO]
+      .configure(config)
+      .use_
+      .attempt
+      .map(_.leftMap(_.getMessage))
+      .assertEquals(
+        Left(
+          """Cannot autoconfigure [Sampler].
+            |Cause: [otel.traces.sampler.arg] has invalid ratio [-0.1] - requirement failed: ratio must be >= 0 and <= 1.0.
+            |Config:
+            |1) `otel.traces.sampler` - parentbased_traceidratio
+            |2) `otel.traces.sampler.arg` - -0.1""".stripMargin
+        )
+      )
+  }
+
+  test("load from the config - unknown sampler - fail") {
+    val props = Map("otel.traces.sampler" -> "some-sampler")
+    val config = Config(props, Map.empty, Map.empty)
+
+    SamplerAutoConfigure[IO]
+      .configure(config)
+      .use_
+      .attempt
+      .map(_.leftMap(_.getMessage))
+      .assertEquals(
+        Left(
+          """Cannot autoconfigure [Sampler].
+            |Cause: Unrecognized value for [otel.traces.sampler]: some-sampler. Supported options [parentbased_traceidratio, traceidratio, parentbased_always_off, always_on, always_off, parentbased_always_on].
+            |Config:
+            |1) `otel.traces.sampler` - some-sampler
+            |2) `otel.traces.sampler.arg` - N/A""".stripMargin
+        )
+      )
+  }
+}