From dc9b74cee5b9cfef31f2f8e8c781a77b42334a77 Mon Sep 17 00:00:00 2001
From: Arman Bilge <armanbilge@gmail.com>
Date: Tue, 16 May 2023 19:29:49 +0000
Subject: [PATCH] POC thread-local iolocals

---
 .../scala/cats/effect/IOFiberConstants.scala  |  2 +
 .../java/cats/effect/IOFiberConstants.java    | 46 ++++++++++---------
 .../src/main/scala/cats/effect/IOFiber.scala  | 11 +++++
 .../scala/cats/effect/unsafe/IOLocals.scala   | 21 +++++++++
 4 files changed, 58 insertions(+), 22 deletions(-)
 create mode 100644 core/shared/src/main/scala/cats/effect/unsafe/IOLocals.scala

diff --git a/core/js-native/src/main/scala/cats/effect/IOFiberConstants.scala b/core/js-native/src/main/scala/cats/effect/IOFiberConstants.scala
index 782fefdfb20..771b99924f2 100644
--- 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
index 5951b447451..acf826b4235 100644
--- 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
index 984e70a03fb..4e36a94fb9a 100644
--- 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
new file mode 100644
index 00000000000..a0ea21a6326
--- /dev/null
+++ 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)
+}