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",
+      scalaTest,
+      "io.opentelemetry" % "opentelemetry-sdk-testing" % "1.67.0" % Test,
+      "io.opentelemetry" % "opentelemetry-exporter-otlp" % "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,131 @@
+# 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(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` therefore **requires** the
+`PropagatingVirtualThreadFactory` from [otel-context](otel-context.md) to be used; without it, spans created in forks
+won't be children of the span active when the fork was created. 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):
+
+```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
+```
+
+## Spans around flows
+
+`spanStages` runs each element of a [`Flow`](../streaming/flows.md) through the stages added by the given function,
+within a span opened for that element. The span is opened when the element enters the stages, and ended when the last
+stage finishes processing everything derived from it. Elements are processed one after another, so each one gets its own
+span, and the flow returned by the function is run for each element separately: stateful stages, e.g. `grouped`, don't
+carry their state across elements.
+
+If the stages fail, the exception is recorded on the span of the element, and its status is set to `ERROR`. An
+interruption and control-flow exits (`boundary.break`) end the span normally, as they aren't failures. The function is
+called within the span, so spans opened by it, and by the stages, are the span's children, and if it throws, the
+exception is recorded as well.
+
+The stages and the span run in a fork, and the results are passed to the stages placed after `spanStages` through a
+buffer, of the size given by the `BufferCapacity` in scope. Because of that boundary, the stages placed after it are not
+counted, spans created there are not the span's children, and their failures are not recorded on it. This holds as long
+as the stages emit at most `BufferCapacity + 1` elements for one input element, as they are processed without waiting
+for downstream. Beyond that, the span also includes the time of waiting for the following stages to free up space in the
+buffer. The buffers of the stages placed after it add up, so the more of them, the later it happens. To include more
+stages in the span, put them inside the function.
+
+In the example below, the `process-numbers` spans are created in a fork, and thanks to the context being propagated,
+they are children 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 = resourceScope {
+    useCloseableInScope(spanUnsafe(tracing)("request")): Unit
+    supervised {
+      fork(Flow.fromValues(1, 2, 3).spanStages(tracing)("process-numbers")(_.map(_ + 1)).runToList()).join()
+    }
+  }
+  println(result)
+```
+
+Optional parameters are `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
+stages, if any, and has to return the span. If it throws, the exception is added as suppressed to the one thrown by the
+stages; 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)
+    .spanStages(tracing)("process-numbers", beforeClose = s => s.span.setAttribute("failed", s.errorOpt.isDefined))(
+      _.map(_ + 1).filter(_ % 2 == 0)
+    )
+    .runToList()
+```
+
+`mapSpanned` is `map`, where each call of the function is run within a span. Unlike `spanStages`, it doesn't use a fork
+and a buffer, so the span covers only the call of the function:
+
+```scala mdoc:compile-only
+import ox.flow.Flow
+import ox.telemetry.tracing.*
+
+def perElement(tracing: OxTracing, flow: Flow[Int]): Unit =
+  flow.mapSpanned(tracing)("increment")(_ + 1).runDrain()
+```
+
+## 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/flow/tracing/OxFlowTracing.scala

@@ -0,0 +1,48 @@
+package ox.flow.tracing
+
+import io.opentelemetry.api.common.Attributes
+import io.opentelemetry.api.trace.{Span, SpanContext, SpanKind}
+import ox.channels.BufferCapacity
+import ox.flow.Flow
+import ox.telemetry.tracing.{OxTracing, SpanWithError}
+
+private[ox] trait OxFlowTracing:
+  self: OxTracing =>
+
+  /** Runs each element of `flow` through the stages added by `f` within its own span: from the element entering the stages, to the last
+    * stage finishing the processing of everything derived from it. Elements are processed one after another. `f` is called within the span
+    * of each element, so spans opened by its body are children of that span, and if it throws, the exception is recorded on it. Stages
+    * created by `f`, including async ones (e.g. `buffer`), are covered by the span. The span is ended when the stages complete, fail or are
+    * interrupted; errors are recorded in it as events.
+    *
+    * The span is separated from the stages placed after `spanStages` by an async boundary, so their processing time is counted only as far
+    * as the `BufferCapacity` allows.
+    *
+    * @param beforeClose
+    *   callback that will be executed right before closing a span
+    */
+  def spanStages[T, U](flow: Flow[T])(
+      spanName: String,
+      spanKind: SpanKind = SpanKind.INTERNAL,
+      attributes: Attributes = Attributes.empty(),
+      links: Seq[SpanContext] = Seq.empty,
+      beforeClose: SpanWithError => Span = _._1
+  )(f: Flow[T] => Flow[U])(using BufferCapacity): Flow[U] =
+    flow.flatMap: t =>
+      val spanned = Flow.usingEmitInline[U]: emit =>
+        self.byNameSpan(f(Flow.fromValues(t)).runToEmit(emit))(spanName, spanKind, attributes, links, beforeClose)
+      spanned.buffer()
+
+  /** As [[ox.flow.Flow.map]] but spans around execution of mapping function Span is made child of span which is current when executing a
+    * flow
+    */
+  def mapSpanned[T, U](flow: Flow[T])(
+      spanName: String,
+      spanKind: SpanKind = SpanKind.INTERNAL,
+      attributes: Attributes = Attributes.empty(),
+      links: Seq[SpanContext] = Seq.empty,
+      beforeClose: SpanWithError => Span = _._1
+  )(f: T => U): Flow[U] =
+    flow.map(t => self.byNameSpan(f(t))(spanName, spanKind, attributes, links, beforeClose))
+
+end OxFlowTracing

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

@@ -0,0 +1,28 @@
+package ox.flow.tracing
+
+import io.opentelemetry.api.common.Attributes
+import io.opentelemetry.api.trace.{Span, SpanContext, SpanKind}
+import ox.channels.BufferCapacity
+import ox.flow.Flow
+import ox.telemetry.tracing.SpanWithError
+
+extension [T](flow: Flow[T])
+  /** Runs each element through the stages added by `f` within its own span, see [[ox.flow.tracing.OxFlowTracing.spanStages]]. */
+  def spanStages[U](oxTracing: OxFlowTracing)(
+      spanName: String,
+      spanKind: SpanKind = SpanKind.INTERNAL,
+      attributes: Attributes = Attributes.empty(),
+      links: Seq[SpanContext] = Seq.empty,
+      beforeClose: SpanWithError => Span = _.span
+  )(f: Flow[T] => Flow[U])(using BufferCapacity): Flow[U] =
+    oxTracing.spanStages(flow)(spanName, spanKind, attributes, links, beforeClose)(f)
+
+  /** As [[ox.flow.Flow.map]], with each call of `f` within a span, see [[ox.flow.tracing.OxFlowTracing.mapSpanned]]. */
+  def mapSpanned[U](oxTracing: OxFlowTracing)(
+      spanName: String,
+      spanKind: SpanKind = SpanKind.INTERNAL,
+      attributes: Attributes = Attributes.empty(),
+      links: Seq[SpanContext] = Seq.empty,
+      beforeClose: SpanWithError => Span = _.span
+  )(f: T => U): Flow[U] = oxTracing.mapSpanned(flow)(spanName, spanKind, attributes, links, beforeClose)(f)
+end extension

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()