IOLocal propagation for unsafe access

build.sbt

@@ -650,7 +650,15 @@ lazy val core = crossProject(JSPlatform, JVMPlatform, NativePlatform)
         "cats.effect.unsafe.IORuntimeBuilder.this"),
       // introduced by #3695, which enabled fiber dumps on native
       ProblemFilters.exclude[MissingClassProblem](
-        "cats.effect.unsafe.FiberMonitorCompanionPlatform")
+        "cats.effect.unsafe.FiberMonitorCompanionPlatform"),
+      // introduced by #3636, IOLocal propagation
+      // IOLocal is a sealed trait
+      ProblemFilters.exclude[ReversedMissingMethodProblem]("cats.effect.IOLocal.getOrDefault"),
+      ProblemFilters.exclude[ReversedMissingMethodProblem]("cats.effect.IOLocal.set"),
+      ProblemFilters.exclude[ReversedMissingMethodProblem]("cats.effect.IOLocal.reset"),
+      ProblemFilters.exclude[ReversedMissingMethodProblem]("cats.effect.IOLocal.lens"),
+      // this filter is particulary terrible, because it can also mask real issues :(
+      ProblemFilters.exclude[DirectMissingMethodProblem]("cats.effect.IOLocal.lens")
     ) ++ {
       if (tlIsScala3.value) {
         // Scala 3 specific exclusions
@@ -890,7 +898,11 @@ lazy val tests: CrossProject = crossProject(JSPlatform, JVMPlatform, NativePlatf
     scalacOptions ~= { _.filterNot(_.startsWith("-P:scalajs:mapSourceURI")) }
   )
   .jvmSettings(
-    Test / fork := true
+    Test / fork := true,
+    Test / javaOptions ++= Seq(
+      "-Dcats.effect.ioLocalPropagation=true",
+      s"-Dsbt.classpath=${(Test / fullClasspath).value.map(_.data.getAbsolutePath).mkString(File.pathSeparator)}"
+    )
   )
   .nativeSettings(
     Compile / mainClass := Some("catseffect.examples.NativeRunner")

core/js-native/src/main/scala/cats/effect/IOFiberConstants.scala

@@ -43,4 +43,6 @@ private object IOFiberConstants {
   final val CedeR = 6
   final val AutoCedeR = 7
   final val DoneR = 8
+
+  final val ioLocalPropagation = false
 }

core/js-native/src/main/scala/cats/effect/unsafe/IOLocalsConstants.scala

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2020-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 cats.effect
+
+private object IOLocalsConstants {
+  final val ioLocalPropagation = false
+}

core/js-native/src/main/scala/cats/effect/unsafe/WorkStealingThreadPool.scala

@@ -46,4 +46,5 @@ private[unsafe] sealed abstract class WorkerThread[P] private () extends Thread
   private[unsafe] def isOwnedBy(threadPool: WorkStealingThreadPool[_]): Boolean
   private[unsafe] def monitor(fiber: Runnable): WeakBag.Handle
   private[unsafe] def index: Int
+  private[unsafe] var ioLocalState: IOLocalState
 }

core/jvm/src/main/java/cats/effect/IOFiberConstants.java

@@ -43,4 +43,6 @@ final class IOFiberConstants {
   static final byte CedeR = 6;
   static final byte AutoCedeR = 7;
   static final byte DoneR = 8;
+
+  static final boolean ioLocalPropagation = Boolean.getBoolean("cats.effect.ioLocalPropagation");
 }

core/jvm/src/main/java/cats/effect/unsafe/IOLocalsConstants.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2020-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 cats.effect.unsafe;
+
+// defined in Java since Scala doesn't let us define static fields
+final class IOLocalsConstants {
+  static final boolean ioLocalPropagation = Boolean.getBoolean("cats.effect.ioLocalPropagation");
+}

core/jvm/src/main/scala/cats/effect/IOPlatform.scala

@@ -71,7 +71,7 @@ abstract private[effect] class IOPlatform[+A] extends Serializable { self: IO[A]
       implicit runtime: unsafe.IORuntime): Option[A] = {
     val queue = new ArrayBlockingQueue[Either[Throwable, A]](1)
 
-    unsafeRunAsync { r =>
+    val fiber = unsafeRunAsyncImpl { r =>
       queue.offer(r)
       ()
     }
@@ -82,6 +82,9 @@ abstract private[effect] class IOPlatform[+A] extends Serializable { self: IO[A]
     } catch {
       case _: InterruptedException =>
         None
+    } finally {
+      if (IOFiberConstants.ioLocalPropagation)
+        unsafe.IOLocals.setState(fiber.getLocalState())
     }
   }
 

core/jvm/src/main/scala/cats/effect/unsafe/WorkerThread.scala

@@ -107,6 +107,8 @@ private final class WorkerThread[P](
   private val indexTransfer: LinkedTransferQueue[Integer] = new LinkedTransferQueue()
   private[this] val runtimeBlockingExpiration: Duration = pool.runtimeBlockingExpiration
 
+  private[unsafe] var ioLocalState: IOLocalState = IOLocalState.empty
+
   val nameIndex: Int = pool.blockedWorkerThreadNamingIndex.getAndIncrement()
 
   // Constructor code.

core/shared/src/main/scala/cats/effect/IO.scala

@@ -912,6 +912,12 @@ sealed abstract class IO[+A] private () extends IOPlatform[A] {
    */
   def unsafeRunAsync(cb: Either[Throwable, A] => Unit)(
       implicit runtime: unsafe.IORuntime): Unit = {
+    unsafeRunAsyncImpl(cb)
+    ()
+  }
+
+  private[effect] def unsafeRunAsyncImpl(cb: Either[Throwable, A] => Unit)(
+      implicit runtime: unsafe.IORuntime): IOFiber[A @uncheckedVariance] =
     unsafeRunFiber(
       cb(Left(new CancellationException("The fiber was canceled"))),
       t => {
@@ -922,8 +928,6 @@ sealed abstract class IO[+A] private () extends IOPlatform[A] {
       },
       a => cb(Right(a))
     )
-    ()
-  }
 
   def unsafeRunAsyncOutcome(cb: Outcome[Id, Throwable, A @uncheckedVariance] => Unit)(
       implicit runtime: unsafe.IORuntime): Unit = {
@@ -1026,7 +1030,7 @@ sealed abstract class IO[+A] private () extends IOPlatform[A] {
       implicit runtime: unsafe.IORuntime): IOFiber[A @uncheckedVariance] = {
 
     val fiber = new IOFiber[A](
-      Map.empty,
+      if (IOFiberConstants.ioLocalPropagation) unsafe.IOLocals.getState else Map.empty,
       oc =>
         oc.fold(
           {

core/shared/src/main/scala/cats/effect/IOFiber.scala

@@ -105,6 +105,8 @@ private final class IOFiber[A](
   @volatile
   private[this] var outcome: OutcomeIO[A] = _
 
+  def getLocalState(): IOLocalState = localState
+
   override def run(): Unit = {
     // insert a read barrier after every async boundary
     readBarrier()
@@ -250,6 +252,10 @@ private final class IOFiber[A](
             pushTracingEvent(cur.event)
           }
 
+          if (ioLocalPropagation) {
+            IOLocals.setState(localState)
+          }
+
           var error: Throwable = null
           val r =
             try cur.thunk()
@@ -260,6 +266,10 @@ private final class IOFiber[A](
                 onFatalFailure(t)
             }
 
+          if (ioLocalPropagation) {
+            localState = IOLocals.getAndClearState()
+          }
+
           val next =
             if (error == null) succeeded(r, 0)
             else failed(error, 0)
@@ -324,6 +334,10 @@ private final class IOFiber[A](
                 pushTracingEvent(delay.event)
               }
 
+              if (ioLocalPropagation) {
+                IOLocals.setState(localState)
+              }
+
               // this code is inlined in order to avoid two `try` blocks
               var error: Throwable = null
               val result =
@@ -335,6 +349,10 @@ private final class IOFiber[A](
                     onFatalFailure(t)
                 }
 
+              if (ioLocalPropagation) {
+                localState = IOLocals.getAndClearState()
+              }
+
               val nextIO = if (error == null) succeeded(result, 0) else failed(error, 0)
               runLoop(nextIO, nextCancelation - 1, nextAutoCede)
 
@@ -391,6 +409,10 @@ private final class IOFiber[A](
                 pushTracingEvent(delay.event)
               }
 
+              if (ioLocalPropagation) {
+                IOLocals.setState(localState)
+              }
+
               // this code is inlined in order to avoid two `try` blocks
               val result =
                 try f(delay.thunk())
@@ -401,6 +423,10 @@ private final class IOFiber[A](
                     onFatalFailure(t)
                 }
 
+              if (ioLocalPropagation) {
+                localState = IOLocals.getAndClearState()
+              }
+
               runLoop(result, nextCancelation - 1, nextAutoCede)
 
             case 3 =>
@@ -446,6 +472,10 @@ private final class IOFiber[A](
                 pushTracingEvent(delay.event)
               }
 
+              if (ioLocalPropagation) {
+                IOLocals.setState(localState)
+              }
+
               // this code is inlined in order to avoid two `try` blocks
               var error: Throwable = null
               val result =
@@ -460,6 +490,10 @@ private final class IOFiber[A](
                     onFatalFailure(t)
                 }
 
+              if (ioLocalPropagation) {
+                localState = IOLocals.getAndClearState()
+              }
+
               val next =
                 if (error == null) succeeded(Right(result), 0) else succeeded(Left(error), 0)
               runLoop(next, nextCancelation - 1, nextAutoCede)
@@ -973,6 +1007,10 @@ private final class IOFiber[A](
             if (ec.isInstanceOf[WorkStealingThreadPool[_]]) {
               val wstp = ec.asInstanceOf[WorkStealingThreadPool[_]]
               if (wstp.canExecuteBlockingCode()) {
+                if (ioLocalPropagation) {
+                  IOLocals.setState(localState)
+                }
+
                 var error: Throwable = null
                 val r =
                   try {
@@ -984,6 +1022,10 @@ private final class IOFiber[A](
                       onFatalFailure(t)
                   }
 
+                if (ioLocalPropagation) {
+                  localState = IOLocals.getAndClearState()
+                }
+
                 val next = if (error eq null) succeeded(r, 0) else failed(error, 0)
                 runLoop(next, nextCancelation, nextAutoCede)
               } else {
@@ -1390,6 +1432,11 @@ private final class IOFiber[A](
     var error: Throwable = null
     val cur = resumeIO.asInstanceOf[Blocking[Any]]
     resumeIO = null
+
+    if (ioLocalPropagation) {
+      IOLocals.setState(localState)
+    }
+
     val r =
       try cur.thunk()
       catch {
@@ -1399,6 +1446,10 @@ private final class IOFiber[A](
           onFatalFailure(t)
       }
 
+    if (ioLocalPropagation) {
+      localState = IOLocals.getAndClearState()
+    }
+
     if (isStackTracing) {
       // Remove the reference to the fiber monitor handle
       objectState.pop().asInstanceOf[WeakBag.Handle].deregister()