OxTelemetry - integration of ox, ox flows with OpenTelemetry

build.sbt

@@ -129,12 +129,26 @@ lazy val otelContext: Project = (project in file("otel-context"))
   .settings(
     name := "otel-context",
     libraryDependencies ++= Seq(
-      "io.opentelemetry" % "opentelemetry-api" % "1.66.0",
+      "io.opentelemetry" % "opentelemetry-api" % "1.67.0",
       scalaTest
     )
   )
   .dependsOn(core % "test->test;compile->compile")
 
+lazy val otel: Project = (project in file("otel"))
+  .settings(
+    name := "ox-telemetry",
+    Test / fork := true,
+    Test / javaOptions += "-Dio.opentelemetry.context.enableStrictContext=true",
+    libraryDependencies ++= Seq(
+      "io.opentelemetry" % "opentelemetry-api" % "1.67.0",
+      "io.opentelemetry" % "opentelemetry-sdk" % "1.67.0",
+      scalaTest,
+      "io.opentelemetry" % "opentelemetry-sdk-testing" % "1.67.0" % Test
+    )
+  )
+  .dependsOn(core % "test->test;compile->compile", otelContext % "test->test;compile->compile")
+
 lazy val documentation: Project = (project in file("generated-doc")) // important: it must not be doc/
   .enablePlugins(MdocPlugin)
   .settings(
@@ -157,5 +171,6 @@ lazy val documentation: Project = (project in file("generated-doc")) // importan
     flowReactiveStreams,
     flowJson,
     cron,
-    otelContext
+    otelContext,
+    otel
   )

core/src/main/scala/ox/flow/FlowOps.scala

@@ -293,8 +293,6 @@ class FlowOps[+T]:
     val base: Flow[U] = expand[U](u => Iterator.single(u) ++ extrapolator(u))
     initial.fold(base)(e => Flow.fromValues(e).concat(base))
 
-  //
-
   /** Applies the given mapping function `f` to each element emitted by this flow. The returned flow then emits the results.
     *
     * @param f

doc/index.md

@@ -93,6 +93,7 @@ In addition to this documentation, ScalaDocs can be browsed at [https://javadoc.
    integrations/mdc-logback
    integrations/cron4s
    integrations/otel-context
+   integrations/otel-tracing
    integrations/tapir
    integrations/sttp-client
 

doc/integrations/otel-context.md

@@ -13,7 +13,7 @@ into traces, or metrics without the appropriate context.
 
 To fix this problem, the context must be propagated whenever a new virtual thread is created. One way to achieve this
 is by using a custom thread factory, provided by this module - `PropagatingVirtualThreadFactory`. It can be set 
-for the whole app when using [`OxApp`](../utils/oxapp.md), or manually through `oxThreadFactory`:
+for the whole app when using [`OxApp`](../utils/oxapp.md), or manually through `setOxThreadFactory`:
 
 ```scala mdoc:compile-only
 import ox.*

doc/integrations/otel-tracing.md

@@ -0,0 +1,104 @@
+# OpenTelemetry tracing
+
+Dependency:
+
+```scala
+"com.softwaremill.ox" %% "ox-telemetry" % "@VERSION@"
+```
+
+The module depends on [otel-context](otel-context.md), so it doesn't need to be added separately. An OpenTelemetry SDK
+and exporter have to be configured by the application, as usual.
+
+Ox provides helpers for creating [OpenTelemetry](https://opentelemetry.io) spans around code that uses Ox's concurrency
+operators, including [Flows](../streaming/flows.md). An opened span becomes a child of the current span, and becomes the
+current span itself, until it's closed.
+
+Spans are created using an `OxTracing` instance, wrapping an OpenTelemetry `Tracer`:
+
+```scala mdoc:compile-only
+import io.opentelemetry.api.trace.Tracer
+import ox.telemetry.tracing.OxTracing
+
+def tracing(tracer: Tracer): OxTracing = OxTracing.native(tracer)
+```
+
+## Propagating the context
+
+The current span is kept in the OpenTelemetry `Context`, which by default isn't propagated to new virtual threads, e.g.
+forks created in [`supervised`](../structured-concurrency/fork-join.md) scopes. `OxTracing.native` relies on the
+`PropagatingVirtualThreadFactory` from [otel-context](otel-context.md) to do that, and prints a warning to stderr if
+it's not set up. The factory can be set for the whole app when using [`OxApp`](../utils/oxapp.md), or manually through
+`setOxThreadFactory`, before the thread factory is first used (that is, before any scope is started, or `OxTracing.native`
+is used):
+
+```scala mdoc:compile-only
+import ox.*
+import ox.otel.context.PropagatingVirtualThreadFactory
+
+object MyApp extends OxApp:
+  override def settings: OxApp.Settings = OxApp.Settings.Default.copy(
+    threadFactory = Some(PropagatingVirtualThreadFactory())
+  )
+
+  def run(args: Vector[String])(using Ox): ExitCode = ExitCode.Success
+```
+
+If you want to take care of context propagation yourself, use `OxTracing.custom`.
+
+## Spans around code and flows
+
+`byNameSpan` evaluates a code block within a span. The span is ended once the block completes, and if the block throws,
+the exception is recorded on the span, and its status is set to `ERROR` (this includes an interruption). `Flow.span`
+does the same for a [`Flow`](../streaming/flows.md): the span is opened when the flow is run, and ended when the flow
+completes, fails, or is interrupted. Downstream stages are run within the span as well, and running the flow multiple
+times opens a new span each time.
+
+In the example below, the `process-numbers` span is created in a fork, and thanks to the context being propagated, it's
+a child of the `request` span:
+
+```scala mdoc:compile-only
+import ox.*
+import ox.flow.Flow
+import ox.telemetry.tracing.*
+
+def example(tracing: OxTracing): Unit =
+  val result = supervised {
+    fork(Flow.fromValues(1, 2, 3).span(tracing)("process-numbers").runToList()).join()
+  }.byNameSpan(tracing)("request")
+  println(result)
+```
+
+Both accept optional `spanKind`, `attributes` and `links`, as well as a `beforeClose` callback, called right before the
+span is closed, e.g. to set attributes or add events. The callback receives the span and the exception thrown by the
+spanned code, if any, and has to return the span. If it throws, the exception is added as suppressed to the one thrown
+by the spanned code; the span is closed in either case:
+
+```scala mdoc:compile-only
+import ox.flow.Flow
+import ox.telemetry.tracing.*
+
+def withAttributes(tracing: OxTracing): List[Int] =
+  Flow
+    .fromValues(1, 2, 3)
+    .span(tracing)("process-numbers", beforeClose = (span, error) => span.setAttribute("failed", error.isDefined))
+    .runToList()
+```
+
+## Manual spans
+
+`spanUnsafe` opens a span, which has to be closed by the caller. Closing the returned `AutoCloseableSpan` ends the span
+and restores the previously current span. It's unsafe, as forgetting to close it leaves the span open and the current
+span unrestored; use it with [`useCloseableInScope`](../utils/resources.md), so that closing is guaranteed.
+`currentSpanUnsafe` returns the current span, e.g. to add attributes:
+
+```scala mdoc:compile-only
+import ox.*
+import ox.telemetry.tracing.*
+
+def manual(tracing: OxTracing): Unit =
+  resourceScope {
+    val opened = useCloseableInScope(spanUnsafe(tracing)("manual-span"))
+    opened.span.setAttribute("key", "value")
+    currentSpanUnsafe(tracing).addEvent("working"): Unit
+  }
+```

otel/src/main/scala/ox/telemetry/ContextHolder.scala

@@ -0,0 +1,48 @@
+package ox.telemetry
+
+import io.opentelemetry.context.{Context, Scope}
+import ox.otel.context.PropagatingVirtualThreadFactory
+
+/** Abstraction over how [[ox.telemetry.tracing.OxTracing]] interacts with OpenTelemetry [[Context]].
+  */
+sealed trait ContextHolder:
+  /** Returns the current [[Context]]. */
+  def get(): Context
+
+  /** Makes the given `context` current. Closing the returned [[Scope]] restores the previously current context. */
+  def set(context: Context): Scope
+
+private[telemetry] object ContextHolder:
+
+  /** Delegates to OpenTelemetry Context. Propagating the context to new (virtual) threads is left to the user.
+    */
+  private[telemetry] object Delegating extends ContextHolder:
+
+    override def get(): Context = Context.current()
+
+    override def set(context: Context): Scope = context.makeCurrent()
+
+  /** As [[Delegating]], but warns when [[ox.oxThreadFactory]] is not a [[PropagatingVirtualThreadFactory]], as in that case the context
+    * won't be propagated to forks.
+    */
+  private[telemetry] object CheckingThreadFactory extends ContextHolder:
+    @volatile private var checked = false
+    private val wrapped: ContextHolder = Delegating
+
+    override def get(): Context =
+      checkThreadFactory()
+      wrapped.get()
+
+    override def set(context: Context): Scope =
+      checkThreadFactory()
+      wrapped.set(context)
+
+    private def checkThreadFactory() =
+      if !checked then
+        if !ox.oxThreadFactory.isInstanceOf[PropagatingVirtualThreadFactory] then
+          System.err.println(
+            s"Default OpenTelemetry integration should use ${classOf[PropagatingVirtualThreadFactory].getName}, if you want to use OpenTelemetry with your own oxThreadFactory, you need to take care of propagation of OpenTelemetry context to newly created VirtualThreads"
+          )
+        checked = true
+  end CheckingThreadFactory
+end ContextHolder

otel/src/main/scala/ox/telemetry/tracing/AutoCloseableSpan.scala

@@ -0,0 +1,19 @@
+package ox.telemetry.tracing
+
+import io.opentelemetry.api.trace.Span
+import io.opentelemetry.context.Scope
+
+/** An opened [[Span]], which is also the current one, together with the [[Scope]] that made it current.
+  *
+  * Calling [[close]] ends the span and closes the scope, restoring the previously current span. Prefer scoped usage, e.g.
+  * [[ox.useCloseableInScope]], so that the span is always closed.
+  */
+sealed trait AutoCloseableSpan extends AutoCloseable:
+  /** The opened span, e.g. to set attributes or add events. */
+  def span: Span
+
+private object AutoCloseableSpan:
+  private[ox] final class Impl(val span: Span, openedScope: Scope) extends AutoCloseableSpan:
+    def close(): Unit =
+      span.end()
+      openedScope.close()

otel/src/main/scala/ox/telemetry/tracing/OxTracing.scala

@@ -0,0 +1,105 @@
+package ox.telemetry.tracing
+
+import io.opentelemetry.api.common.Attributes
+import io.opentelemetry.api.trace.{Span, SpanContext, SpanKind, StatusCode, Tracer}
+import ox.telemetry.ContextHolder
+import ox.telemetry.tracing.flow.OxFlowTracing
+
+sealed trait OxTracing extends OxFlowTracing:
+
+  protected def contextHolder: ContextHolder
+
+  protected def tracer: Tracer
+
+  /** Returns the current span. This is not unsafe per se, but the caller is responsible for using the span properly, e.g. not ending a span
+    * that is managed elsewhere.
+    */
+  def currentSpanUnsafe(): Span = Span.fromContext(contextHolder.get())
+
+  /** Opens a new [[Span]] as a child of the current span. The new span becomes the current one.
+    *
+    * Unsafe, because the caller is responsible for closing the returned [[AutoCloseableSpan]], which ends the span and restores the
+    * previously current one. Scoped usage, e.g. [[ox.useCloseableInScope]], guarantees that.
+    *
+    * @return
+    *   An [[AutoCloseableSpan]], which closes the opened span.
+    */
+  def spanUnsafe(
+      spanName: String,
+      spanKind: SpanKind = SpanKind.INTERNAL,
+      attributes: Attributes = Attributes.empty(),
+      links: Seq[SpanContext] = Seq.empty
+  ): AutoCloseableSpan =
+    val previousContext = contextHolder.get()
+    val builder = tracer.spanBuilder(spanName).setSpanKind(spanKind).setAllAttributes(attributes).setParent(previousContext)
+    val withLinks = links.foldLeft(builder)((b, l) => b.addLink(l))
+    val started = withLinks.startSpan()
+    val scope = contextHolder.set(previousContext.`with`(started))
+
+    AutoCloseableSpan.Impl(started, scope)
+  end spanUnsafe
+
+  /** Evaluates `t` within a new [[Span]]. The span is opened before `t` is evaluated, as a child of the current span, and becomes the
+    * current one. It is ended once `t` completes, successfully or not, and the previously current span is restored.
+    *
+    * If `t` throws, the exception is recorded on the span and the span's status is set to [[io.opentelemetry.api.trace.StatusCode.ERROR]],
+    * before `beforeClose` is called. The exception is then re-thrown.
+    *
+    * To span a [[ox.flow.Flow]], use [[OxFlowTracing.span]].
+    *
+    * @param beforeClose
+    *   Run right before the span is closed, e.g. to set attributes or add events. The second element of the argument is the exception
+    *   thrown by `t`, if any. If `beforeClose` itself throws, the exception is added as suppressed to the one thrown by `t`, or thrown if
+    *   `t` succeeded. The span is closed in either case.
+    */
+  def byNameSpan[T](t: => T)(
+      spanName: String,
+      spanKind: SpanKind = SpanKind.INTERNAL,
+      attributes: Attributes = Attributes.empty(),
+      links: Seq[SpanContext] = Seq.empty,
+      beforeClose: ((Span, Option[Throwable])) => Span = _._1
+  ): T =
+    var error: Throwable | Null = null
+    var finalizer: AutoCloseableSpan | Null = null
+    try
+      val finalize = spanUnsafe(spanName, spanKind, attributes, links)
+      finalizer = finalize
+      t
+    catch
+      case e: Throwable =>
+        error = e
+        if finalizer ne null then
+          finalizer.nn.span.recordException(e)
+          finalizer.nn.span.setStatus(StatusCode.ERROR, e.getMessage): Unit
+        throw e
+    finally
+      try if finalizer ne null then beforeClose(finalizer.nn.span, Option(error)): Unit
+      catch
+        case e: Throwable =>
+          if error ne null then error.nn.addSuppressed(e)
+          else throw e
+      finally if finalizer ne null then finalizer.nn.close()
+      end try
+    end try
+
+  end byNameSpan
+end OxTracing
+
+object OxTracing:
+
+  /** Tracing that relies on the [[ox.otel.context.PropagatingVirtualThreadFactory]] to propagate the context to forks. A warning is printed
+    * to stderr on first use, if the thread factory is not set up.
+    *
+    * @see
+    *   [[ox.OxApp.Settings.threadFactory]], [[ox.setOxThreadFactory]]
+    */
+  def native(tracer: Tracer): OxTracing = Impl(tracer, ContextHolder.CheckingThreadFactory)
+
+  /** Tracing which leaves the propagation of the [[io.opentelemetry.context.Context]] to new threads to the user. For an out-of-the-box
+    * integration, use [[native]] together with the [[ox.otel.context.PropagatingVirtualThreadFactory]].
+    */
+  def custom(tracer: Tracer): OxTracing = Impl(tracer, ContextHolder.Delegating)
+
+  private final class Impl(val tracer: Tracer, val contextHolder: ContextHolder) extends OxTracing
+
+end OxTracing

otel/src/main/scala/ox/telemetry/tracing/flow/OxFlowTracing.scala

@@ -0,0 +1,29 @@
+package ox.telemetry.tracing.flow
+
+import io.opentelemetry.api.common.Attributes
+import io.opentelemetry.api.trace.{Span, SpanContext, SpanKind}
+import ox.flow.Flow
+import ox.telemetry.tracing.OxTracing
+
+private[tracing] trait OxFlowTracing:
+  self: OxTracing =>
+
+  /** Returns a flow, which when run, evaluates `flow` within a new span. The span is opened when the flow is run, as a child of the span
+    * current at that point. It is ended when the flow completes, fails or is interrupted. Exceptions are recorded on the span as in
+    * [[OxTracing.byNameSpan]].
+    *
+    * The parameters are the same as for [[OxTracing.byNameSpan]].
+    */
+  def span[T](flow: Flow[T])(
+      spanName: String,
+      spanKind: SpanKind = SpanKind.INTERNAL,
+      attributes: Attributes = Attributes.empty(),
+      links: Seq[SpanContext] = Seq.empty,
+      beforeClose: ((Span, Option[Throwable])) => Span = _._1
+  ): Flow[T] = Flow.usingEmit: emit =>
+    self.byNameSpan(
+      flow.runToEmit(
+        emit
+      )
+    )(spanName, spanKind, attributes, links, beforeClose)
+end OxFlowTracing

otel/src/main/scala/ox/telemetry/tracing/flow/flowTracing.scala

@@ -0,0 +1,19 @@
+package ox.telemetry.tracing.flow
+
+import io.opentelemetry.api.common.Attributes
+import io.opentelemetry.api.trace.{Span, SpanContext, SpanKind}
+import ox.flow.Flow
+import ox.telemetry.tracing.flow.OxFlowTracing
+
+extension [T](flow: Flow[T])
+  /** Opens a span around the run of this flow, using the given [[ox.telemetry.tracing.OxTracing]], see
+    * [[ox.telemetry.tracing.OxTracing.span]].
+    */
+  def span(oxTracing: OxFlowTracing)(
+      spanName: String,
+      spanKind: SpanKind = SpanKind.INTERNAL,
+      attributes: Attributes = Attributes.empty(),
+      links: Seq[SpanContext] = Seq.empty,
+      beforeClose: ((Span, Option[Throwable])) => Span = _._1
+  ): Flow[T] = oxTracing.span(flow)(spanName, spanKind, attributes, links, beforeClose)
+end extension