Fiber locals

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

@@ -216,6 +216,7 @@ sealed abstract class IO[+A] private () extends IOPlatform[A] {
 
     val fiber = new IOFiber[A](
       0,
+      Map(),
       oc =>
         oc.fold(
           (),
@@ -665,6 +666,11 @@ object IO extends IOCompanionPlatform with IOLowPriorityImplicits {
     def tag = 20
   }
 
+  // Not part of the run loop. Only used in the implementation of IO#to.
+  private[effect] final case class UnmaskTo[F[_], +A](ioa: IO[A], poll: Poll[F]) extends IO[A] {
+    def tag = -1
+  }
+
   private[effect] final case class Blocking[+A](hint: Sync.Type, thunk: () => A) extends IO[A] {
     def tag = 21
   }
@@ -677,6 +683,14 @@ object IO extends IOCompanionPlatform with IOLowPriorityImplicits {
     def tag = 23
   }
 
+  private[effect] final case class GetLocal[A](index: Int) extends IO[Option[A]] {
+    def tag = 24
+  }
+
+  private[effect] final case class SetLocal[A](index: Int, value: A) extends IO[Unit] {
+    def tag = 25
+  }
+
   // INTERNAL, only created by the runloop itself as the terminal state of several operations
   private[effect] case object EndFiber extends IO[Nothing] {
     def tag = -1

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

@@ -26,7 +26,7 @@ import scala.concurrent.duration._
 import scala.util.control.NonFatal
 
 import java.util.concurrent.RejectedExecutionException
-import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicReference}
 import scala.util.control.NoStackTrace
 
 /*
@@ -61,6 +61,7 @@ import scala.util.control.NoStackTrace
  */
 private final class IOFiber[A](
     initMask: Int,
+    initLocalState: scala.collection.immutable.Map[Int, Any],
     cb: OutcomeIO[A] => Unit,
     startIO: IO[A],
     startEC: ExecutionContext,
@@ -101,6 +102,8 @@ private final class IOFiber[A](
   /* true when semantically blocking (ensures that we only unblock *once*) */
   private[this] val suspended: AtomicBoolean = new AtomicBoolean(true)
 
+  private[this] var localState: scala.collection.immutable.Map[Int, Any] = initLocalState
+
   @volatile
   private[this] var outcome: OutcomeIO[A] = _
 
@@ -595,6 +598,7 @@ private final class IOFiber[A](
           val ec = currentCtx
           val fiber = new IOFiber[Any](
             initMask2,
+            localState,
             null,
             cur.ioa,
             ec,
@@ -621,13 +625,15 @@ private final class IOFiber[A](
                   val ec = currentCtx
                   val fiberA = new IOFiber[Any](
                     initMask2,
+                    localState,
                     null,
                     cur.ioa,
                     ec,
                     runtime
                   )
                   val fiberB = new IOFiber[Any](
                     initMask2,
+                    localState,
                     null,
                     cur.iob,
                     ec,
@@ -712,13 +718,15 @@ private final class IOFiber[A](
                 val ec = currentCtx
                 val fiberA = new IOFiber[Any](
                   initMask2,
+                  localState,
                   null,
                   cur.ioa,
                   ec,
                   runtime
                 )
                 val fiberB = new IOFiber[Any](
                   initMask2,
+                  localState,
                   null,
                   cur.iob,
                   ec,
@@ -798,13 +806,15 @@ private final class IOFiber[A](
                 val ec = currentCtx
                 val fiberA = new IOFiber[Any](
                   initMask2,
+                  localState,
                   null,
                   cur.ioa,
                   ec,
                   runtime
                 )
                 val fiberB = new IOFiber[Any](
                   initMask2,
+                  localState,
                   null,
                   cur.iob,
                   ec,
@@ -873,6 +883,18 @@ private final class IOFiber[A](
             }.guarantee(IO.defer(finalizer.get()))
 
           runLoop(next, nextIteration)
+
+        case 24 =>
+          val cur = cur0.asInstanceOf[GetLocal[Any]]
+
+          val value = localState.get(cur.index)
+          runLoop(succeeded(value, 0), nextIteration)
+
+        case 25 =>
+          val cur = cur0.asInstanceOf[SetLocal[Any]]
+
+          localState = localState + (cur.index -> cur.value)
+          runLoop(succeeded((), 0), nextIteration)
       }
     }
   }
@@ -1320,9 +1342,14 @@ private final class IOFiber[A](
 }
 
 private object IOFiber {
+  private val localIndex = new AtomicInteger(0)
+
   /* prefetch */
   private val OutcomeCanceled = Outcome.Canceled()
   private[effect] val RightUnit = Right(())
+
+  private[effect] def nextLocalIndex(): Int =
+    localIndex.getAndIncrement()
 }
 
 private[effect] case object AsyncPropagateCancelation extends NoStackTrace

core/shared/src/main/scala/cats/effect/local/Local.scala

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2020 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.local
+
+import cats.effect.{IO, IOFiber}
+
+final class Local[A] private (index: Int, default: A) {
+
+  def get: IO[A] =
+    IO.GetLocal[A](index).map(_.getOrElse(default))
+
+  def set(value: A): IO[Unit] =
+    IO.SetLocal[A](index, value)
+
+  def update(f: A => A): IO[Unit] =
+    get.flatMap(a => set(f(a)))
+
+  def getAndSet(value: A): IO[A] =
+    get <* set(value)
+
+  def clear: IO[Unit] =
+    IO.SetLocal[A](index, default)
+
+}
+
+object Local {
+
+  def apply[A](default: A): IO[Local[A]] =
+    of(default)
+
+  def of[A](default: A): IO[Local[A]] =
+    IO {
+      val index = IOFiber.nextLocalIndex()
+      new Local(index, default)
+    }
+
+}

core/shared/src/main/scala/cats/effect/local/LocalRef.scala

@@ -0,0 +1,64 @@
+/*
+ * Copyright 2020 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.local
+
+import cats.effect.IO
+import cats.effect.kernel.Ref
+
+final class LocalRef[A] private (local: Local[Ref[IO, A]]) extends Ref[IO, A] {
+
+  override def get: IO[A] =
+    local.get.flatMap(_.get)
+
+  override def set(a: A): IO[Unit] =
+    local.get.flatMap(_.set(a))
+
+  override def access: IO[(A, A => IO[Boolean])] =
+    local.get.flatMap(_.access)
+
+  override def tryUpdate(f: A => A): IO[Boolean] =
+    local.get.flatMap(_.tryUpdate(f))
+
+  override def tryModify[B](f: A => (A, B)): IO[Option[B]] =
+    local.get.flatMap(_.tryModify(f))
+
+  override def update(f: A => A): IO[Unit] =
+    local.get.flatMap(_.update(f))
+
+  override def modify[B](f: A => (A, B)): IO[B] =
+    local.get.flatMap(_.modify(f))
+
+  override def tryModifyState[B](state: cats.data.State[A, B]): IO[Option[B]] =
+    local.get.flatMap(_.tryModifyState(state))
+
+  override def modifyState[B](state: cats.data.State[A, B]): IO[B] =
+    local.get.flatMap(_.modifyState(state))
+
+}
+
+object LocalRef {
+
+  def apply[A](default: A): IO[LocalRef[A]] =
+    of(default)
+
+  def of[A](default: A): IO[LocalRef[A]] =
+    for {
+      ref <- Ref.of[IO, A](default)
+      local <- Local.of[Ref[IO, A]](ref)
+    } yield new LocalRef(local)
+
+}

core/shared/src/test/scala/cats/effect/local/LocalRefSpec.scala

@@ -0,0 +1,87 @@
+/*
+ * Copyright 2020 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
+package effect
+package local
+
+class LocalRefSpec extends BaseSpec {
+
+  "LocalRef" should {
+    "return a default value" in ticked { implicit ticker =>
+      val io = LocalRef(0).flatMap(_.get)
+
+      io must completeAs(0)
+    }
+
+    "set and get a value" in ticked { implicit ticker =>
+      val io = for {
+        local <- LocalRef(0)
+        _ <- local.set(10)
+        value <- local.get
+      } yield value
+
+      io must completeAs(10)
+    }
+
+    "preserve locals across async boundaries" in ticked { implicit ticker =>
+      val io = for {
+        local <- LocalRef(0)
+        _ <- local.set(10)
+        _ <- IO.cede
+        value <- local.get
+      } yield value
+
+      io must completeAs(10)
+    }
+
+    "children fibers can read locals" in ticked { implicit ticker =>
+      val io = for {
+        local <- LocalRef(0)
+        _ <- local.set(10)
+        f <- local.get.start
+        value <- f.joinAndEmbedNever
+      } yield value
+
+      io must completeAs(10)
+    }
+
+    "child local manipulation is visible to parents" in ticked { implicit ticker =>
+      val io = for {
+        local <- LocalRef(0)
+        f <- local.set(20).start
+        _ <- f.join
+        value <- local.get
+      } yield value
+
+      io must completeAs(20)
+    }
+
+    "parent local manipulation is visible to children" in ticked { implicit ticker =>
+      val io = for {
+        local <- LocalRef(0)
+        d1 <- Deferred[IO, Unit]
+        f <- (d1.get *> local.get).start
+        _ <- local.set(10)
+        _ <- d1.complete(())
+        value <- f.joinAndEmbedNever
+      } yield value
+
+      io must completeAs(10)
+    }
+  }
+
+}