POC thread-local iolocals

typelevel · May 16, 2023 · dc9b74c · dc9b74c
1 parent 6ffcb8e
commit dc9b74c
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 22 deletions.
diff --git a/core/js-native/src/main/scala/cats/effect/IOFiberConstants.scala b/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 dumpBooleans = false
 }
diff --git a/core/jvm/src/main/java/cats/effect/IOFiberConstants.java b/core/jvm/src/main/java/cats/effect/IOFiberConstants.java
@@ -19,28 +19,30 @@
 // defined in Java since Scala doesn't let us define static fields
 final class IOFiberConstants {
 
-  static final int MaxStackDepth = 512;
+	static final int MaxStackDepth = 512;
 
-  // continuation ids (should all be inlined)
-  static final byte MapK = 0;
-  static final byte FlatMapK = 1;
-  static final byte CancelationLoopK = 2;
-  static final byte RunTerminusK = 3;
-  static final byte EvalOnK = 4;
-  static final byte HandleErrorWithK = 5;
-  static final byte OnCancelK = 6;
-  static final byte UncancelableK = 7;
-  static final byte UnmaskK = 8;
-  static final byte AttemptK = 9;
+	// continuation ids (should all be inlined)
+	static final byte MapK = 0;
+	static final byte FlatMapK = 1;
+	static final byte CancelationLoopK = 2;
+	static final byte RunTerminusK = 3;
+	static final byte EvalOnK = 4;
+	static final byte HandleErrorWithK = 5;
+	static final byte OnCancelK = 6;
+	static final byte UncancelableK = 7;
+	static final byte UnmaskK = 8;
+	static final byte AttemptK = 9;
 
-  // resume ids
-  static final byte ExecR = 0;
-  static final byte AsyncContinueSuccessfulR = 1;
-  static final byte AsyncContinueFailedR = 2;
-  static final byte AsyncContinueCanceledR = 3;
-  static final byte AsyncContinueCanceledWithFinalizerR = 4;
-  static final byte BlockingR = 5;
-  static final byte CedeR = 6;
-  static final byte AutoCedeR = 7;
-  static final byte DoneR = 8;
+	// resume ids
+	static final byte ExecR = 0;
+	static final byte AsyncContinueSuccessfulR = 1;
+	static final byte AsyncContinueFailedR = 2;
+	static final byte AsyncContinueCanceledR = 3;
+	static final byte AsyncContinueCanceledWithFinalizerR = 4;
+	static final byte BlockingR = 5;
+	static final byte CedeR = 6;
+	static final byte AutoCedeR = 7;
+	static final byte DoneR = 8;
+
+	static final boolean dumpLocals = Boolean.getBoolean("cats.effect.tracing.dumpLocals");
 }
diff --git a/core/shared/src/main/scala/cats/effect/IOFiber.scala b/core/shared/src/main/scala/cats/effect/IOFiber.scala
@@ -250,6 +250,12 @@ private final class IOFiber[A](
             pushTracingEvent(cur.event)
           }
 
+          var locals: IOLocals = null
+          if (dumpLocals) {
+            locals = new IOLocals(localState)
+            IOLocals.threadLocal.set(locals)
+          }
+
           var error: Throwable = null
           val r =
             try cur.thunk()
@@ -260,6 +266,11 @@ private final class IOFiber[A](
                 onFatalFailure(t)
             }
 
+          if (dumpLocals) {
+            localState = locals.getState()
+            IOLocals.threadLocal.set(null)
+          }
+
           val next =
             if (error == null) succeeded(r, 0)
             else failed(error, 0)

diff --git a/core/shared/src/main/scala/cats/effect/unsafe/IOLocals.scala b/core/shared/src/main/scala/cats/effect/unsafe/IOLocals.scala
@@ -0,0 +1,21 @@
+package cats.effect
+package unsafe
+
+// TODO handle defaults and lenses. all do-able, just needs refactoring ...
+final class IOLocals private[effect] (private[this] var state: IOLocalState) {
+
+  private[effect] def getState(): IOLocalState = state
+
+  def get[A](iol: IOLocal[A]): A = state(iol).asInstanceOf[A]
+
+  def set[A](iol: IOLocal[A], value: A): Unit = state += (iol -> value)
+
+  def reset[A](iol: IOLocal[A]): Unit = state -= iol
+
+  // TODO other ops from IOLocal
+
+}
+
+object IOLocals {
+  val threadLocal = ThreadLocal.withInitial[IOLocals](() => null)
+}