Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fiber locals #1393

Merged
merged 24 commits into from
Mar 26, 2021
Merged
Show file tree
Hide file tree
Changes from 17 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
6 changes: 6 additions & 0 deletions core/shared/src/main/scala/cats/effect/IO.scala
Expand Up @@ -224,6 +224,7 @@ sealed abstract class IO[+A] private () extends IOPlatform[A] {

val fiber = new IOFiber[A](
0,
Map(),
oc =>
oc.fold(
canceled,
Expand Down Expand Up @@ -691,6 +692,11 @@ object IO extends IOCompanionPlatform with IOLowPriorityImplicits {
def tag = 20
}

private[effect] final case class Local[+A](f: IOLocalState => (IOLocalState, A))
extends IO[A] {
def tag = 21
}

// 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
Expand Down
18 changes: 17 additions & 1 deletion core/shared/src/main/scala/cats/effect/IOFiber.scala
Expand Up @@ -26,7 +26,7 @@ import scala.concurrent.duration._
import scala.util.control.NonFatal

import java.util.concurrent.RejectedExecutionException
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
import scala.util.control.NoStackTrace

/*
Expand Down Expand Up @@ -67,6 +67,7 @@ import scala.util.control.NoStackTrace
*/
private final class IOFiber[A](
initMask: Int,
initLocalState: IOLocalState,
cb: OutcomeIO[A] => Unit,
startIO: IO[A],
startEC: ExecutionContext,
Expand Down Expand Up @@ -106,6 +107,8 @@ private final class IOFiber[A](

private[this] val callbacks = new CallbackStack[A](cb)

private[this] var localState: IOLocalState = initLocalState

@volatile
private[this] var outcome: OutcomeIO[A] = _

Expand Down Expand Up @@ -594,6 +597,7 @@ private final class IOFiber[A](
val ec = currentCtx
val fiber = new IOFiber[Any](
initMask2,
localState,
null,
cur.ioa,
ec,
Expand Down Expand Up @@ -658,6 +662,13 @@ private final class IOFiber[A](
} else {
runLoop(interruptibleImpl(cur, runtime.blocking), nextIteration)
}

case 21 =>
val cur = cur0.asInstanceOf[Local[Any]]

val (nextLocalState, value) = cur.f(localState)
localState = nextLocalState
runLoop(succeeded(value, 0), nextIteration)
}
}
}
Expand Down Expand Up @@ -1125,9 +1136,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
58 changes: 58 additions & 0 deletions core/shared/src/main/scala/cats/effect/Local.scala
@@ -0,0 +1,58 @@
/*
* Copyright 2020-2021 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

final class Local[A] private (index: Int, default: A) {

def get: IO[A] =
IO.Local(state => (state, state.get(index).map(_.asInstanceOf[A]).getOrElse(default)))

def set(value: A): IO[Unit] =
IO.Local(state => (state + (index -> value), ()))

def clear: IO[Unit] =
IO.Local(state => (state - index, ()))

def update(f: A => A): IO[Unit] =
get.flatMap(a => set(f(a)))

def modify[B](f: A => (A, B)): IO[B] =
get.flatMap { a =>
val (a2, b) = f(a)
set(a2).as(b)
}

def getAndSet(value: A): IO[A] =
get <* set(value)

def getAndClear: IO[A] =
get <* clear

}

object Local {

def apply[A](default: A): IO[Local[A]] =
RaasAhsan marked this conversation as resolved.
Show resolved Hide resolved
of(default)

def of[A](default: A): IO[Local[A]] =
IO {
val index = IOFiber.nextLocalIndex()
new Local(index, default)
}

}
61 changes: 61 additions & 0 deletions core/shared/src/main/scala/cats/effect/LocalRef.scala
@@ -0,0 +1,61 @@
/*
* Copyright 2020-2021 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

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]] =
RaasAhsan marked this conversation as resolved.
Show resolved Hide resolved
for {
ref <- Ref.of[IO, A](default)
local <- Local.of[Ref[IO, A]](ref)
} yield new LocalRef(local)

}
2 changes: 2 additions & 0 deletions core/shared/src/main/scala/cats/effect/package.scala
Expand Up @@ -76,4 +76,6 @@ package object effect {

type Ref[F[_], A] = cekernel.Ref[F, A]
val Ref = cekernel.Ref

private[effect] type IOLocalState = scala.collection.immutable.Map[Int, Any]
RaasAhsan marked this conversation as resolved.
Show resolved Hide resolved
}
37 changes: 37 additions & 0 deletions core/shared/src/test/scala/cats/effect/TestContextSpec.scala
@@ -0,0 +1,37 @@
/*
* Copyright 2020-2021 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

// import cats.implicits._
// import scala.concurrent.duration._

// class TestContextSpec extends BaseSpec {
RaasAhsan marked this conversation as resolved.
Show resolved Hide resolved
// def io: IO[Unit] =
// for {
// fiber <- IO.unit.foreverM.start
// _ <- IO(println("A"))
// _ <- IO.sleep(100.millis)
// _ <- IO(println("B"))
// _ <- fiber.cancel
// } yield ()

// // "tc works" in real {
// // io.as(success)
// // }

// // "tc deadlocks" in ticked { implicit ticker => io must completeAs(()) }
// }
86 changes: 86 additions & 0 deletions tests/shared/src/test/scala/cats/effect/LocalRefSpec.scala
@@ -0,0 +1,86 @@
/*
* Copyright 2020-2021 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

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.joinWithNever
} 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.joinWithNever
} yield value

io must completeAs(10)
}
}

}