Thread local fallback weak bag

benchmarks/src/main/scala/cats/effect/benchmarks/ThreadLocalBenchmark.scala

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2020-2022 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.benchmarks
+
+import cats.effect.unsafe.FiberMonitor
+
+import org.openjdk.jmh.annotations._
+
+import scala.concurrent.{Await, ExecutionContext, Future}
+import scala.concurrent.duration._
+
+import java.util.concurrent.TimeUnit
+
+/**
+ * To run the benchmark from within sbt:
+ *
+ * jmh:run -i 10 -wi 10 -f 2 -t 1 cats.effect.benchmarks.ThreadLocalBenchmark
+ *
+ * Which means "10 iterations", "10 warm-up iterations", "2 forks", "1 thread". Please note that
+ * benchmarks should be usually executed at least in 10 iterations (as a rule of thumb), but
+ * more is better.
+ */
+@State(Scope.Thread)
+@BenchmarkMode(Array(Mode.Throughput))
+@OutputTimeUnit(TimeUnit.SECONDS)
+class ThreadLocalBenchmark {
+
+  final implicit val executionContext: ExecutionContext = ExecutionContext.global
+
+  @Param(Array("2000"))
+  var size: Int = _
+
+  @Benchmark
+  def contention() = {
+    val monitor = new FiberMonitor(null)
+
+    def future(): Future[Unit] = Future {
+      monitor.monitorSuspended(null)
+      ()
+    }
+
+    Await.result(Future.sequence(List.fill(size)(future())), Duration.Inf)
+  }
+}

build.sbt

@@ -451,7 +451,10 @@ lazy val core = crossProject(JSPlatform, JVMPlatform)
       // introduced by #2769, Simplify the transfer of WorkerThread data structures when blocking
       // changes to `cats.effect.unsafe` package private code
       ProblemFilters.exclude[MissingClassProblem]("cats.effect.unsafe.WorkerThread$"),
-      ProblemFilters.exclude[MissingClassProblem]("cats.effect.unsafe.WorkerThread$Data")
+      ProblemFilters.exclude[MissingClassProblem]("cats.effect.unsafe.WorkerThread$Data"),
+      // introduced by #2844, Thread local fallback weak bag
+      // changes to `cats.effect.unsafe` package private code
+      ProblemFilters.exclude[MissingClassProblem]("cats.effect.unsafe.SynchronizedWeakBag")
     ) ++ {
       if (isDotty.value) {
         // Scala 3 specific exclusions

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

@@ -18,11 +18,12 @@ package cats.effect
 package unsafe
 
 import cats.effect.tracing.TracingConstants
+import cats.effect.unsafe.ref.WeakReference
 
 import scala.collection.mutable
 import scala.concurrent.ExecutionContext
 
-import java.util.concurrent.ThreadLocalRandom
+import java.util.concurrent.ConcurrentLinkedQueue
 
 /**
  * A slightly more involved implementation of an unordered bag used for tracking asynchronously
@@ -49,16 +50,8 @@ private[effect] final class FiberMonitor(
     private[this] val compute: WorkStealingThreadPool
 ) extends FiberMonitorShared {
 
-  private[this] val size: Int = Runtime.getRuntime().availableProcessors() << 2
-  private[this] val bags: Array[SynchronizedWeakBag[IOFiber[_]]] = new Array(size)
-
-  {
-    var i = 0
-    while (i < size) {
-      bags(i) = new SynchronizedWeakBag()
-      i += 1
-    }
-  }
+  private[this] final val Bags = FiberMonitor.Bags
+  private[this] final val BagReferences = FiberMonitor.BagReferences
 
   /**
    * Registers a suspended fiber.
@@ -145,18 +138,21 @@ private[effect] final class FiberMonitor(
     else ()
 
   private[this] def monitorFallback(fiber: IOFiber[_]): WeakBag.Handle = {
-    val rnd = ThreadLocalRandom.current()
-    val idx = rnd.nextInt(size)
-    bags(idx).insert(fiber)
+    val bag = Bags.get()
+    val handle = bag.insert(fiber)
+    bag.synchronizationPoint.lazySet(true)
+    handle
   }
 
   private[this] def foreignFibers(): Set[IOFiber[_]] = {
     val foreign = mutable.Set.empty[IOFiber[_]]
 
-    var i = 0
-    while (i < size) {
-      foreign ++= bags(i).toSet
-      i += 1
+    BagReferences.iterator().forEachRemaining { bagRef =>
+      val bag = bagRef.get()
+      if (bag ne null) {
+        val _ = bag.synchronizationPoint.get()
+        foreign ++= bag.toSet
+      }
     }
 
     foreign.toSet
@@ -172,4 +168,15 @@ private[effect] object FiberMonitor {
       new FiberMonitor(null)
     }
   }
+
+  private[FiberMonitor] final val Bags: ThreadLocal[WeakBag[IOFiber[_]]] =
+    ThreadLocal.withInitial { () =>
+      val bag = new WeakBag[IOFiber[_]]()
+      BagReferences.offer(new WeakReference(bag))
+      bag
+    }
+
+  private[FiberMonitor] final val BagReferences
+      : ConcurrentLinkedQueue[WeakReference[WeakBag[IOFiber[_]]]] =
+    new ConcurrentLinkedQueue()
 }

core/shared/src/main/scala/cats/effect/unsafe/WeakBag.scala

@@ -21,6 +21,8 @@ import cats.effect.unsafe.ref.{ReferenceQueue, WeakReference}
 import scala.annotation.tailrec
 import scala.collection.mutable
 
+import java.util.concurrent.atomic.AtomicBoolean
+
 private final class WeakBag[A <: AnyRef] {
   import WeakBag._
 
@@ -29,7 +31,8 @@ private final class WeakBag[A <: AnyRef] {
   private[this] val queue: ReferenceQueue[A] = new ReferenceQueue()
   private[this] var capacity: Int = 256
   private[this] var table: Array[Entry[A]] = new Array(capacity)
-  private[this] var index = 0
+  private[this] var index: Int = 0
+  private[unsafe] val synchronizationPoint: AtomicBoolean = new AtomicBoolean(true)
 
   @tailrec
   def insert(a: A): Handle = {