IOLocal propagation for unsafe access

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 dumpLocals = false
 }

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

@@ -46,4 +46,5 @@ private[unsafe] sealed abstract class WorkerThread 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

@@ -17,7 +17,7 @@
 package cats.effect;
 
 // defined in Java since Scala doesn't let us define static fields
-final class IOFiberConstants {
+public final class IOFiberConstants {
 
   static final int MaxStackDepth = 512;
 
@@ -43,4 +43,6 @@ final class IOFiberConstants {
   static final byte CedeR = 6;
   static final byte AutoCedeR = 7;
   static final byte DoneR = 8;
+
+  public static final boolean dumpLocals = Boolean.getBoolean("cats.effect.tracing.dumpLocals");
 }

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

@@ -100,6 +100,8 @@ private final class WorkerThread(
   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

@@ -1009,7 +1009,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.dumpLocals) unsafe.IOLocals.getState else Map.empty,
       oc =>
         oc.fold(
           {

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

@@ -250,6 +250,10 @@ private final class IOFiber[A](
             pushTracingEvent(cur.event)
           }
 
+          if (dumpLocals) {
+            IOLocals.setState(localState)
+          }
+
           var error: Throwable = null
           val r =
             try cur.thunk()
@@ -260,6 +264,10 @@ private final class IOFiber[A](
                 onFatalFailure(t)
             }
 
+          if (dumpLocals) {
+            localState = IOLocals.getAndClearState()
+          }
+
           val next =
             if (error == null) succeeded(r, 0)
             else failed(error, 0)
@@ -324,6 +332,10 @@ private final class IOFiber[A](
                 pushTracingEvent(delay.event)
               }
 
+              if (dumpLocals) {
+                IOLocals.setState(localState)
+              }
+
               // this code is inlined in order to avoid two `try` blocks
               var error: Throwable = null
               val result =
@@ -335,6 +347,10 @@ private final class IOFiber[A](
                     onFatalFailure(t)
                 }
 
+              if (dumpLocals) {
+                localState = IOLocals.getAndClearState()
+              }
+
               val nextIO = if (error == null) succeeded(result, 0) else failed(error, 0)
               runLoop(nextIO, nextCancelation - 1, nextAutoCede)
 
@@ -391,6 +407,10 @@ private final class IOFiber[A](
                 pushTracingEvent(delay.event)
               }
 
+              if (dumpLocals) {
+                IOLocals.setState(localState)
+              }
+
               // this code is inlined in order to avoid two `try` blocks
               val result =
                 try f(delay.thunk())
@@ -401,6 +421,10 @@ private final class IOFiber[A](
                     onFatalFailure(t)
                 }
 
+              if (dumpLocals) {
+                localState = IOLocals.getAndClearState()
+              }
+
               runLoop(result, nextCancelation - 1, nextAutoCede)
 
             case 3 =>
@@ -446,6 +470,10 @@ private final class IOFiber[A](
                 pushTracingEvent(delay.event)
               }
 
+              if (dumpLocals) {
+                IOLocals.setState(localState)
+              }
+
               // this code is inlined in order to avoid two `try` blocks
               var error: Throwable = null
               val result =
@@ -460,6 +488,10 @@ private final class IOFiber[A](
                     onFatalFailure(t)
                 }
 
+              if (dumpLocals) {
+                localState = IOLocals.getAndClearState()
+              }
+
               val next =
                 if (error == null) succeeded(Right(result), 0) else succeeded(Left(error), 0)
               runLoop(next, nextCancelation - 1, nextAutoCede)
@@ -965,6 +997,10 @@ private final class IOFiber[A](
             if (ec.isInstanceOf[WorkStealingThreadPool]) {
               val wstp = ec.asInstanceOf[WorkStealingThreadPool]
               if (wstp.canExecuteBlockingCode()) {
+                if (dumpLocals) {
+                  IOLocals.setState(localState)
+                }
+
                 var error: Throwable = null
                 val r =
                   try {
@@ -976,6 +1012,10 @@ private final class IOFiber[A](
                       onFatalFailure(t)
                   }
 
+                if (dumpLocals) {
+                  localState = IOLocals.getAndClearState()
+                }
+
                 val next = if (error eq null) succeeded(r, 0) else failed(error, 0)
                 runLoop(next, nextCancelation, nextAutoCede)
               } else {
@@ -1378,6 +1418,11 @@ private final class IOFiber[A](
     var error: Throwable = null
     val cur = resumeIO.asInstanceOf[Blocking[Any]]
     resumeIO = null
+
+    if (dumpLocals) {
+      IOLocals.setState(localState)
+    }
+
     val r =
       try cur.thunk()
       catch {
@@ -1387,6 +1432,10 @@ private final class IOFiber[A](
           onFatalFailure(t)
       }
 
+    if (dumpLocals) {
+      localState = IOLocals.getAndClearState()
+    }
+
     if (isStackTracing) {
       // Remove the reference to the fiber monitor handle
       objectState.pop().asInstanceOf[WeakBag.Handle].deregister()

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

@@ -136,35 +136,49 @@ import cats.data.AndThen
  * @tparam A
  *   the type of the local value
  */
-sealed trait IOLocal[A] { self =>
+sealed trait IOLocal[A] {
+
+  protected[effect] def getOrDefault(state: IOLocalState): A
+
+  protected[effect] def set(state: IOLocalState, a: A): IOLocalState
+
+  protected[effect] def reset(state: IOLocalState): IOLocalState
 
   /**
    * Returns the current value.
    */
-  def get: IO[A]
+  final def get: IO[A] =
+    IO.Local(state => (state, getOrDefault(state)))
 
   /**
    * Sets the current value to `value`.
    */
-  def set(value: A): IO[Unit]
+  final def set(value: A): IO[Unit] =
+    IO.Local(state => (set(state, value), ()))
 
   /**
    * Replaces the current value with the initial value.
    */
-  def reset: IO[Unit]
+  final def reset: IO[Unit] =
+    IO.Local(state => (reset(state), ()))
 
   /**
    * Modifies the current value using the given update function.
    */
-  def update(f: A => A): IO[Unit]
+  final def update(f: A => A): IO[Unit] =
+    IO.Local(state => (set(state, f(getOrDefault(state))), ()))
 
   /**
    * Like [[update]] but allows the update function to return an output value of type `B`.
    *
    * @see
    *   [[update]]
    */
-  def modify[B](f: A => (A, B)): IO[B]
+  final def modify[B](f: A => (A, B)): IO[B] =
+    IO.Local { state =>
+      val (a2, b) = f(getOrDefault(state))
+      (set(state, a2), b)
+    }
 
   /**
    * Replaces the current value with `value`, returning the previous value.
@@ -176,7 +190,8 @@ sealed trait IOLocal[A] { self =>
    * @see
    *   [[set]]
    */
-  def getAndSet(value: A): IO[A]
+  final def getAndSet(value: A): IO[A] =
+    IO.Local(state => (set(state, value), getOrDefault(state)))
 
   /**
    * Replaces the current value with the initial value, returning the previous value.
@@ -188,7 +203,8 @@ sealed trait IOLocal[A] { self =>
    * @see
    *   [[reset]]
    */
-  def getAndReset: IO[A]
+  final def getAndReset: IO[A] =
+    IO.Local(state => (reset(state), getOrDefault(state)))
 
   /**
    * Creates a lens to a value of some type `B` from current value and two functions: getter and
@@ -220,28 +236,7 @@ sealed trait IOLocal[A] { self =>
    *   } yield ()
    * }}}
    */
-  final def lens[B](get: A => B)(set: A => B => A): IOLocal[B] = {
-    import IOLocal.IOLocalLens
-
-    self match {
-      case lens: IOLocalLens[aa, A] =>
-        // We process already created lens separately so
-        // we wont pay additional `.get.flatMap` price for every call of
-        // `set`, `update` or `modify` of resulting lens.
-        // After all, our getters and setters are pure,
-        // so `AndThen` allows us to safely compose them and
-        // proxy calls to the 'original' `IOLocal` independent of
-        // current nesting level.
-
-        val getter = lens.getter.andThen(get)
-        val setter = lens.setter.compose((p: (aa, B)) => (p._1, set(lens.getter(p._1))(p._2)))
-        new IOLocalLens(lens.underlying, getter, setter)
-      case _ =>
-        val getter = AndThen(get)
-        val setter = AndThen((p: (A, B)) => set(p._1)(p._2))
-        new IOLocalLens(self, getter, setter)
-    }
-  }
+  def lens[B](get: A => B)(set: A => B => A): IOLocal[B]
 
 }
 
@@ -260,63 +255,49 @@ object IOLocal {
    */
   def apply[A](default: A): IO[IOLocal[A]] = IO(new IOLocalImpl(default))
 
-  private[IOLocal] final class IOLocalImpl[A](default: A) extends IOLocal[A] { self =>
-    private[this] def getOrDefault(state: IOLocalState): A =
-      state.getOrElse(self, default).asInstanceOf[A]
-
-    def get: IO[A] =
-      IO.Local(state => (state, getOrDefault(state)))
+  private[effect] abstract class AbstractIOLocal[A] extends IOLocal[A] {}
 
-    def set(value: A): IO[Unit] =
-      IO.Local(state => (state.updated(self, value), ()))
+  private final class IOLocalImpl[A](default: A) extends IOLocal[A] {
 
-    def reset: IO[Unit] =
-      IO.Local(state => (state - self, ()))
+    def getOrDefault(state: IOLocalState): A =
+      state.getOrElse(this, default).asInstanceOf[A]
 
-    def update(f: A => A): IO[Unit] =
-      IO.Local(state => (state.updated(self, f(getOrDefault(state))), ()))
+    def set(state: IOLocalState, a: A): IOLocalState = state.updated(this, a)
 
-    def modify[B](f: A => (A, B)): IO[B] =
-      IO.Local { state =>
-        val (a2, b) = f(getOrDefault(state))
-        (state.updated(self, a2), b)
-      }
+    def reset(state: IOLocalState): IOLocalState = state - this
 
-    def getAndSet(value: A): IO[A] =
-      IO.Local(state => (state.updated(self, value), getOrDefault(state)))
-
-    def getAndReset: IO[A] =
-      IO.Local(state => (state - self, getOrDefault(state)))
+    def lens[B](get: A => B)(set: A => B => A): IOLocal[B] =
+      new IOLocal.IOLocalLens(this, get, (ab: (A, B)) => set(ab._1)(ab._2))
   }
 
-  private[IOLocal] final class IOLocalLens[S, A](
-      val underlying: IOLocal[S],
-      val getter: AndThen[S, A],
-      val setter: AndThen[(S, A), S])
+  private final class IOLocalLens[S, A](
+      underlying: IOLocal[S],
+      getter: S => A,
+      setter: ((S, A)) => S)
       extends IOLocal[A] {
-    def get: IO[A] =
-      underlying.get.map(getter(_))
-
-    def set(value: A): IO[Unit] =
-      underlying.get.flatMap(s => underlying.set(setter(s -> value)))
 
-    def reset: IO[Unit] =
-      underlying.reset
+    def getOrDefault(state: IOLocalState): A =
+      getter(underlying.getOrDefault(state))
 
-    def update(f: A => A): IO[Unit] =
-      underlying.get.flatMap(s => underlying.set(setter(s -> f(getter(s)))))
+    def set(state: IOLocalState, a: A): IOLocalState =
+      underlying.set(state, setter((underlying.getOrDefault(state), a)))
 
-    def modify[B](f: A => (A, B)): IO[B] =
-      underlying.get.flatMap { s =>
-        val (a2, b) = f(getter(s))
-        underlying.set(setter(s -> a2)).as(b)
-      }
+    def reset(state: IOLocalState): IOLocalState = underlying.reset(state)
 
-    def getAndSet(value: A): IO[A] =
-      underlying.get.flatMap(s => underlying.set(setter(s -> value)).as(getter(s)))
+    def lens[B](get: A => B)(set: A => B => A): IOLocal[B] = {
+      // We process already created lens separately so
+      // we wont pay additional `.get.flatMap` price for every call of
+      // `set`, `update` or `modify` of resulting lens.
+      // After all, our getters and setters are pure,
+      // so `AndThen` allows us to safely compose them and
+      // proxy calls to the 'original' `IOLocal` independent of
+      // current nesting level.
 
-    def getAndReset: IO[A] =
-      underlying.get.flatMap(s => underlying.reset.as(getter(s)))
+      val getter = AndThen(this.getter).andThen(get)
+      val setter =
+        AndThen(this.setter).compose((sb: (S, B)) => (sb._1, set(this.getter(sb._1))(sb._2)))
+      new IOLocalLens(underlying, getter, setter)
+    }
   }
 
 }

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

@@ -80,6 +80,9 @@ package object effect {
   val Ref = cekernel.Ref
 
   private[effect] type IOLocalState = scala.collection.immutable.Map[IOLocal[_], Any]
+  private[effect] object IOLocalState {
+    val empty: IOLocalState = scala.collection.immutable.Map.empty
+  }
 
   private[effect] type ByteStack = ByteStack.T
 }