Profiling fix

Author: rochala

build.mill

@@ -70,6 +70,10 @@ object sls extends CommonScalaModule {
 
   def moduleDeps: Seq[JavaModule] = Seq(slsSmithy, profilingRuntime)
 
+  override def forkArgs: T[Seq[String]] = Seq(
+    "-Dcats.effect.trackFiberContext=true"
+  )
+
   def ivyDeps = Agg(
     ivy"org.typelevel::cats-effect:3.6.3",
     ivy"co.fs2::fs2-io:3.13.0-M2",
@@ -79,15 +83,14 @@ object sls extends CommonScalaModule {
     ivy"ch.qos.logback:logback-classic:1.4.14",
     ivy"com.lihaoyi::os-lib:0.11.4",
     ivy"org.polyvariant.smithy4s-bsp::bsp4s:0.5.0",
-    ivy"org.scalameta:mtags-interfaces:1.5.1",
+    ivy"org.scalameta:mtags-interfaces:1.6.1-SNAPSHOT",
     ivy"com.evolution::scache:5.1.2",
     ivy"org.typelevel::cats-parse:1.1.0",
     ivy"io.get-coursier:interface:1.0.28",
   )
 
   def javacOptions = Seq(
     "-Dotel.service.name=simple-language-server",
-    "-Dcats.effect.trackFiberContext=true"
   )
 
   object test extends ScalaTests {

profilingRuntime/src/org/scala/abusers/profiling/runtime/LocalSpan.scala

@@ -1,20 +1,34 @@
-package org.scala.abusers.profiling.runtime
+package cats.effect.org.scala.abusers.profiling.runtime
 
 import scala.concurrent.ExecutionContext
 import io.pyroscope.vendor.one.profiler.AsyncProfiler
+import org.typelevel.otel4s.sdk.context.Context
+import org.typelevel.otel4s.sdk.trace.SdkContextKeys
+import org.typelevel.otel4s.sdk.context.TraceContext
 
 object ProfilingExecutionContext {
 
-  def wrapExecutionContext(ec: ExecutionContext, profiler: AsyncProfiler, localSpan: LocalSpan): ExecutionContext = {
+  def wrapExecutionContext(ec: ExecutionContext, profiler: AsyncProfiler, localContext: ThreadLocal[Context]): ExecutionContext = {
     new ExecutionContext {
       def execute(runnable: Runnable): Unit = {
-        val spanOpt = localSpan.get()
 
         ec.execute(new Runnable {
+          val ctx = localContext.get()
           def run(): Unit = {
 
+            val spanOpt = ctx
+              .get(SdkContextKeys.SpanContextKey)
+              .filter(_.isValid)
+              .map { ctx =>
+                TraceContext(
+                  ctx.traceId,
+                  ctx.spanId,
+                  ctx.isSampled
+                )
+              }
+
             spanOpt.foreach { span =>
-              profiler.setTracingContext(span.spanId, 0)
+                profiler.setTracingContext(span.spanId.toLong(signed = false), 0)
             }
 
             try {

profilingRuntime/src/org/scala/abusers/profiling/runtime/ProfilingExecutionContext.scala

@@ -19,11 +19,18 @@ import cats.effect.unsafe.IORuntimeBuilder
 import org.typelevel.otel4s.experimental.metrics._
 import org.typelevel.otel4s.trace.Tracer
 import org.typelevel.otel4s.metrics.Meter
+import org.typelevel.otel4s.context.LocalProvider
+import org.typelevel.otel4s.sdk.context.Context
+import cats.effect.org.scala.abusers.profiling.runtime.ProfilingExecutionContext
 
+object ProfilingIOAppSettings {
+  lazy val isEnabled: Boolean = sys.env.get("SLS_PROFILING").contains("true") || sys.props.get("sls.profiling").contains("true")
+}
 
 trait ProfilingIOApp extends IOApp {
   private lazy val profiler = PyroscopeAsyncProfiler.getAsyncProfiler()
-  private lazy val localSpan0: IOLocal[Option[TraceSpan]] = IOLocal[Option[TraceSpan]](None)
+
+  private given localCtx: IOLocal[Context] = IOLocal[Context](Context.root)
     .syncStep(100)
     .flatMap(
       _.leftMap(_ =>
@@ -34,28 +41,29 @@ trait ProfilingIOApp extends IOApp {
     )
     .unsafeRunSync()
 
+  private val threadLocal = localCtx.unsafeThreadLocal()
 
-  private val localSpan = new LocalSpan(() => localSpan0)
-  private val _runtime = {
+  private lazy val _runtime =
     IORuntimeBuilder()
-      .transformCompute { ec => ProfilingExecutionContext.wrapExecutionContext(ec, profiler, localSpan) }
-      .transformBlocking { ec => ProfilingExecutionContext.wrapExecutionContext(ec, profiler, localSpan) }
+      .transformCompute { ec => ProfilingExecutionContext.wrapExecutionContext(ec, profiler, threadLocal) }
+      .transformBlocking { ec => ProfilingExecutionContext.wrapExecutionContext(ec, profiler, threadLocal) }
       .build()
-  }
 
-  override protected def runtime: IORuntime = _runtime
+  override protected def runtime: IORuntime = if ProfilingIOAppSettings.isEnabled then _runtime else null /* Null is the default value */
+
+  def program(using meter: Meter[IO], tracer: Tracer[IO]): Resource[IO, Unit]
+  def applicationName: String
 
   override def run(args: List[String]): IO[ExitCode] =
     OpenTelemetrySdk
-      .autoConfigured[IO]( // register OTLP exporters configurer
+      .autoConfigured[IO](
         _.addExportersConfigurer(OtlpExportersAutoConfigure[IO])
-          .addTracerProviderCustomizer((builder, _) =>
-              builder.addSpanProcessor(new ProfilingSpanProcessor(localSpan)
-        ))
+         .addTracerProviderCustomizer((builder, _) => builder.addSpanProcessor(new ProfilingSpanProcessor))
       )
-      .flatTap(_ => registerPyroscope()) // register Pyroscope agent
+      .flatTap(_ => registerPyroscope())
       .use { autoConfigured =>
         val sdk = autoConfigured.sdk
+
         val app = for {
           given Meter[IO] <- sdk.meterProvider.get(applicationName).toResource
           given Tracer[IO] <- sdk.tracerProvider.get(applicationName).toResource
@@ -66,9 +74,6 @@ trait ProfilingIOApp extends IOApp {
         app.useForever
       }.as(ExitCode.Success)
 
-  def program(using meter: Meter[IO], tracer: Tracer[IO]): Resource[IO, Unit]
-  def applicationName: String
-
   private def registerPyroscope(): Resource[IO, Unit] = {
     val acquire = IO.delay {
       PyroscopeAgent.start(

profilingRuntime/src/org/scala/abusers/profiling/runtime/ProfilingIOApp.scala

@@ -0,0 +1,17 @@
+package org.scala.abusers.profiling.runtime
+
+
+import org.typelevel.otel4s.trace.SpanOps
+import org.typelevel.otel4s.trace.Span
+import cats.effect.IO
+
+object ProfilingOps {
+  extension (ops: SpanOps[IO]) {
+    def profilingUse[A](f: Span[IO] => IO[A]): IO[A] =
+      if ProfilingIOAppSettings.isEnabled then
+        ops.use(span => (IO.cede *> f(span)).guarantee(IO.cede))
+      else ops.use(f)
+
+    inline def profilingSurround[A](fa: IO[A]): IO[A] = profilingUse(_ => fa)
+  }
+}

profilingRuntime/src/org/scala/abusers/profiling/runtime/ProfilingOps.scala

@@ -9,23 +9,19 @@ import org.typelevel.otel4s.sdk.trace.SpanRef
 import org.typelevel.otel4s.sdk.trace.data.SpanData
 import org.typelevel.otel4s.Attribute
 
-class ProfilingSpanProcessor(localSpan: LocalSpan) extends SpanProcessor[IO] {
+class ProfilingSpanProcessor extends SpanProcessor[IO] {
 
   override def name: String = "ProfilingSpanProcessor"
 
   override def onStart: OnStart[IO] = new OnStart[IO] {
     override def apply(parentContext: Option[SpanContext], span: SpanRef[IO]): IO[Unit] = {
       val spanIdStr = span.context.spanIdHex
-      val spanId = java.lang.Long.parseUnsignedLong(spanIdStr, 16)
-      for {
-        _ <- span.addAttributes(Seq(Attribute("pyroscope.profile.id", spanIdStr)))
-        _ <- IO(localSpan.set(Some(TraceSpan(spanId))))
-      } yield ()
+      span.addAttributes(Seq(Attribute("pyroscope.profile.id", spanIdStr)))
     }
   }
 
   override def onEnd: OnEnd[IO] = new OnEnd[IO] {
-    override def apply(span: SpanData): IO[Unit] = IO(localSpan.set(None))
+    override def apply(span: SpanData): IO[Unit] = IO.unit
   }
 
   override def forceFlush: IO[Unit] = IO.unit

profilingRuntime/src/org/scala/abusers/profiling/runtime/ProfilingSpanProcessor.scala

@@ -2,24 +2,19 @@ package org.scala.abusers.sls
 
 import cats.effect.IO
 import cats.effect.std.Semaphore
-import org.typelevel.otel4s.trace.Tracer
 
 sealed trait SynchronizedState
 
 trait ComputationQueue {
   protected given SynchronizedState = new SynchronizedState {}
-  def synchronously[A](computation: SynchronizedState ?=> IO[A])(using Tracer[IO]): IO[A]
+  def synchronously[A](computation: SynchronizedState ?=> IO[A]): IO[A]
 }
 
 class ComputationQueueImpl(semaphore: Semaphore[IO]) extends ComputationQueue {
-  def synchronously[A](computation: SynchronizedState ?=> IO[A])(using Tracer[IO]): IO[A] = {
-      semaphore.permit.use { _ =>
-        (computation)
-      }
-  }
+  def synchronously[A](computation: SynchronizedState ?=> IO[A]): IO[A] =
+    semaphore.permit.surround(computation)
 }
 
 object ComputationQueue {
-  def instance: IO[ComputationQueue] =
-    Semaphore[IO](1).map(ComputationQueueImpl.apply)
+  def instance: IO[ComputationQueue] = Semaphore[IO](1).map(ComputationQueueImpl.apply)
 }

sls/src/org/scala/abusers/sls/ComputationQueue.scala

@@ -35,6 +35,7 @@ import LoggingUtils.*
 import ScalaBuildTargetInformation.*
 import org.typelevel.otel4s.trace.Tracer
 import org.typelevel.otel4s.metrics.Meter
+import org.scala.abusers.profiling.runtime.ProfilingOps.*
 
 
 class ServerImpl(
@@ -85,23 +86,23 @@ class ServerImpl(
     computationQueue.synchronously { bspStateManager.importBuild }
 
   def textDocumentCompletionOp(params: lsp.CompletionParams): IO[lsp.TextDocumentCompletionOpOutput] =
-    Tracer[IO].span("completion").surround(handleCompletion(params))
+    Tracer[IO].span("completion").profilingSurround(handleCompletion(params))
   def textDocumentDefinitionOp(params: lsp.DefinitionParams): IO[lsp.TextDocumentDefinitionOpOutput] =
-    Tracer[IO].span("go-to-defintion").surround(handleDefinition(params))
+    Tracer[IO].span("go-to-defintion").profilingSurround(handleDefinition(params))
   def textDocumentDidChange(params: lsp.DidChangeTextDocumentParams): IO[Unit]        =
-    Tracer[IO].span("did-change").surround(handleDidChange(params))
+    Tracer[IO].span("did-change").profilingSurround(handleDidChange(params))
   def textDocumentDidClose(params: lsp.DidCloseTextDocumentParams): IO[Unit]          =
-    Tracer[IO].span("did-close").surround(handleDidClose(params))
+    Tracer[IO].span("did-close").profilingSurround(handleDidClose(params))
   def textDocumentDidOpen(params: lsp.DidOpenTextDocumentParams): IO[Unit]            =
-    Tracer[IO].span("did-open").surround(handleDidOpen(params))
+    Tracer[IO].span("did-open").profilingSurround(handleDidOpen(params))
   def textDocumentDidSave(params: lsp.DidSaveTextDocumentParams): IO[Unit]            =
-    Tracer[IO].span("did-save").surround(handleDidSave(params))
+    Tracer[IO].span("did-save").profilingSurround(handleDidSave(params))
   def textDocumentHoverOp(params: lsp.HoverParams): IO[lsp.TextDocumentHoverOpOutput] =
-    Tracer[IO].span("hover").surround(handleHover(params))
+    Tracer[IO].span("hover").profilingSurround(handleHover(params))
   def textDocumentInlayHintOp(params: lsp.InlayHintParams): IO[lsp.TextDocumentInlayHintOpOutput] =
-    Tracer[IO].span("inlay-hints").surround(handleInlayHints(params))
+    Tracer[IO].span("inlay-hints").profilingSurround(handleInlayHints(params))
   def textDocumentSignatureHelpOp(params: lsp.SignatureHelpParams): IO[lsp.TextDocumentSignatureHelpOpOutput] =
-    Tracer[IO].span("signature-help").surround(handleSignatureHelp(params))
+    Tracer[IO].span("signature-help").profilingSurround(handleSignatureHelp(params))
 
   // // TODO: goto type definition with container types
   def handleCompletion(params: lsp.CompletionParams) =

sls/src/org/scala/abusers/sls/ServerImpl.scala

@@ -1,6 +1,5 @@
 package org.scala.abusers.sls
 
-import cats.effect.*
 import cats.syntax.all.*
 import jsonrpclib.fs2.*
 import jsonrpclib.smithy4sinterop.ClientStub
@@ -9,6 +8,8 @@ import org.scala.abusers.pc.IOCancelTokens
 import org.scala.abusers.pc.PresentationCompilerProvider
 import org.typelevel.otel4s.trace.Tracer
 import org.typelevel.otel4s.metrics.Meter
+import org.scala.abusers.profiling.runtime.ProfilingIOApp
+import cats.effect.*
 
 case class BuildServer(
     generic: bsp.BuildServer[IO],
@@ -41,7 +42,7 @@ object LSPCancelRequest {
     )
 }
 
-object SimpleScalaServer extends org.scala.abusers.profiling.runtime.ProfilingIOApp {
+object SimpleScalaServer extends ProfilingIOApp {
   import jsonrpclib.smithy4sinterop.ServerEndpoints
 
   override def applicationName: String = "simple-language-server"
@@ -70,7 +71,7 @@ object SimpleScalaServer extends org.scala.abusers.profiling.runtime.ProfilingIO
         )
         .compile
         .drain
-        .background
+        .toResource
     } yield ()
 
   private def server(lspClient: SlsLanguageClient[IO])(using Tracer[IO], Meter[IO]): Resource[IO, ServerImpl] =