From c3c9cee55f3f4247abe38a6818d87a27a30e705d Mon Sep 17 00:00:00 2001
From: Maksym Ochenashko <maxochenashko@gmail.com>
Date: Tue, 14 Nov 2023 12:50:52 +0200
Subject: [PATCH] sdk-trace: add `IdGenerator`

---
 .../otel4s/sdk/trace/IdGenerator.scala        | 81 +++++++++++++++++++
 .../otel4s/sdk/trace/IdGeneratorSuite.scala   | 52 ++++++++++++
 2 files changed, 133 insertions(+)
 create mode 100644 sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/IdGenerator.scala
 create mode 100644 sdk/trace/src/test/scala/org/typelevel/otel4s/sdk/trace/IdGeneratorSuite.scala

diff --git a/sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/IdGenerator.scala b/sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/IdGenerator.scala
new file mode 100644
index 000000000..89fe4db54
--- /dev/null
+++ b/sdk/trace/src/main/scala/org/typelevel/otel4s/sdk/trace/IdGenerator.scala
@@ -0,0 +1,81 @@
+/*
+ * 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
+
+import cats.Monad
+import cats.effect.std.Random
+import cats.syntax.flatMap._
+import cats.syntax.functor._
+import cats.syntax.monad._
+import org.typelevel.otel4s.trace.SpanContext
+import scodec.bits.ByteVector
+
+/** Generates new span and trace ids.
+  *
+  * @see
+  *   [[https://opentelemetry.io/docs/specs/otel/trace/sdk/#id-generators]]
+  *
+  * @tparam F
+  *   the higher-kinded type of a polymorphic effect
+  */
+trait IdGenerator[F[_]] {
+
+  /** Generates a valid span id.
+    */
+  def generateSpanId: F[ByteVector]
+
+  /** Generates a valid trace id.
+    */
+  def generateTraceId: F[ByteVector]
+
+  /** Whether it's safe to skip the ID validation: we are sure the generated ids
+    * are valid.
+    */
+  private[trace] def canSkipIdValidation: Boolean =
+    false
+}
+
+object IdGenerator {
+
+  private final val InvalidId = 0
+
+  /** Creates [[IdGenerator]] that uses [[cats.effect.std.Random]] under the
+    * hood.
+    */
+  def random[F[_]: Monad: Random]: IdGenerator[F] =
+    new RandomIdGenerator[F]
+
+  private final class RandomIdGenerator[F[_]: Monad: Random]
+      extends IdGenerator[F] {
+
+    def generateSpanId: F[ByteVector] =
+      Random[F].nextLong
+        .iterateUntil(_ != InvalidId)
+        .map(SpanContext.SpanId.fromLong)
+
+    def generateTraceId: F[ByteVector] =
+      for {
+        hi <- Random[F].nextLong
+        lo <- Random[F].nextLong.iterateUntil(_ != InvalidId)
+      } yield SpanContext.TraceId.fromLongs(hi, lo)
+
+    // we trust ourselves
+    override private[trace] def canSkipIdValidation: Boolean =
+      true
+  }
+
+}
diff --git a/sdk/trace/src/test/scala/org/typelevel/otel4s/sdk/trace/IdGeneratorSuite.scala b/sdk/trace/src/test/scala/org/typelevel/otel4s/sdk/trace/IdGeneratorSuite.scala
new file mode 100644
index 000000000..a4cb92b19
--- /dev/null
+++ b/sdk/trace/src/test/scala/org/typelevel/otel4s/sdk/trace/IdGeneratorSuite.scala
@@ -0,0 +1,52 @@
+/*
+ * 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
+
+import cats.effect.IO
+import cats.effect.std.Random
+import munit.CatsEffectSuite
+import org.typelevel.otel4s.trace.SpanContext
+
+class IdGeneratorSuite extends CatsEffectSuite {
+  private val Attempts = 1_000_000
+
+  generatorTest("generate a valid trace id") { generator =>
+    generator.generateTraceId
+      .map(v => assert(SpanContext.TraceId.isValid(v)))
+      .replicateA_(Attempts)
+  }
+
+  generatorTest("generate a valid span id") { generator =>
+    generator.generateSpanId
+      .map(v => assert(SpanContext.SpanId.isValid(v)))
+      .replicateA_(Attempts)
+  }
+
+  generatorTest("should be trusted") { generator =>
+    IO(assert(generator.canSkipIdValidation))
+  }
+
+  private def generatorTest(
+      name: String
+  )(body: IdGenerator[IO] => IO[Unit]): Unit =
+    test(name) {
+      Random
+        .scalaUtilRandom[IO]
+        .flatMap { implicit random: Random[IO] => body(IdGenerator.random[IO]) }
+    }
+
+}