Leaking fiber monitors when running in custom IORuntime using JDK 21 virtual threads

[jacum]

import cats.effect.std.Dispatcher
import cats.effect.unsafe.{IORuntime, IORuntimeBuilder}
import cats.effect.{ExitCode, IO, IOApp}
import cats.implicits.toTraverseOps
import com.sun.management.HotSpotDiagnosticMXBean

import java.io.File
import java.lang.management.ManagementFactory
import java.util.UUID
import java.util.concurrent.Executors
import scala.concurrent.ExecutionContext

object CatsEffectLeakDumper extends IOApp {

  override protected def runtime: IORuntime = {
    val compute          = Executors.newVirtualThreadPerTaskExecutor()
    val computeCtx = ExecutionContext.fromExecutor(compute);
    IORuntimeBuilder
      .apply()
      .setCompute(computeCtx, () => compute.shutdown())
      .build()
  }

  def dumpHeap(fileName: String, live: Boolean = false): Unit = {
    val server = ManagementFactory.getPlatformMBeanServer
    val hotspotMBean = ManagementFactory.newPlatformMXBeanProxy(server, "com.sun.management:type=HotSpotDiagnostic", classOf[HotSpotDiagnosticMXBean])
    if (new File(fileName).exists()) {
      val _ = new File(fileName).delete
    }
    hotspotMBean.dumpHeap(fileName, live)
  }
  override def run(args: List[String]): IO[ExitCode] = Dispatcher.parallel[IO].flatMap { d =>
    IO(
      Runtime.getRuntime.addShutdownHook(new Thread(() => dumpHeap("heapdump.hprof", live = true)))
    ) >> {
      for {
        _ <- (1 to 1000000).toList.traverse(_ => IO.delay {
          d.unsafeRunSync(IO(UUID.randomUUID()))
        })
      } yield ()
    } toResource
  }.use(_ => IO.pure(ExitCode.Success))

}

Run with -Dcats.effect.tracing.mode=FULL on JDK21.

Observe heapdump.hprof -> biggest objects -> cats.effect.unsafe.FiberMonitor -- many MBs
cc @djspiewak

[durban]

It seems the leak is due to the nodes of the BagReferences ConcurrentLinkedQueue in FiberMonitor, and the (empty) weakref objects in them. As far as I can see, these are never cleaned up, so the same would happen with normal threads too (except those are probably not created in an unbounded number).

[armanbilge]

Hmm, that's actually really annoying. I think we'll want to rethink that thread-local strategy for virtual threads since there's a 1-to-1 fiber-to-virtual-thread mapping. Ideally we'd be using the carrier thread's local storage, not the virtual thread's for the bag. But that's probably not possible.

so the same would happen with normal threads too (except those are probably not created in an unbounded number).

But we should probably still fix this anyway.

[durban]

Ideally we'd be using the carrier thread's local storage, not the virtual thread's for the bag. But that's probably not possible.

Well... https://github.com/openjdk/jdk21u/blob/060c4f7589e7f13febd402f4dac3320f4c032b08/src/java.base/share/classes/jdk/internal/misc/CarrierThreadLocal.java#L35

(I'm not seriously suggesting we use that, just found it interesting.)

[armanbilge]

Thanks, makes sense that we are not the only ones needing that 😇

I'm not seriously suggesting we use that

even if we wanted to these new JVMs make it annoying/impossible to access their internals

[djspiewak]

You know, I'm actually very surprised that thread locals aren't just cleaned up with the thread itself. Arman is right that we really want the carrier thread here, not the virtual thread, but even with the virtual thread this doesn't feel like a thing that should be leaking (just a thing that would be very inefficient).

[armanbilge]

It seems the leak is due to the nodes of the BagReferences ConcurrentLinkedQueue in FiberMonitor, and the (empty) weakref objects in them.

@djspiewak as Daniel U pointed out, the leak is actually our fault:

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

Lines 52 to 58 in 0207637

	private[this] final val BagReferences =
	new ConcurrentLinkedQueue[WeakReference[WeakBag[Runnable]]]
	private[this] final val Bags = ThreadLocal.withInitial { () =>
	val bag = new WeakBag[Runnable]()
	BagReferences.offer(new WeakReference(bag))
	bag
	}

We are lacking a "packing" mechanism for BagReferences.

[djspiewak]

Right but that's a weak reference, so it shouldn't prevent the cleanup of bag.

[armanbilge]

Right, but the ConcurrentLinkedQueue will keep growing more and more nodes, pointing to GCed bags.

[djspiewak]

The solution is probably to do something annoying with ReferenceQueue and periodically check that queue for size. Once the queue gets big enough, we go through and purge out the collected refs in Bags. This would have to be amortized into monitorSuspended or similar, creating a moderate overhead. Probably not enough overhead to worry about gating it to virtual threads, so this would also work to resolve the similar issue with blocking.

We wouldn't need JDK 21 for this (since it's not loom specific), so we could probably resolve it in 3.5.x

[armanbilge]

I've proposed a fix for the general issue in #3964.

But from a performance standpoint I still think that we need a different solution for virtual threads. The current mechanism adds a lot of overhead:

1. a single point of contention for registering every thread
2. allocating an entire WeakBag (intended to be long-lived and hold references to many fibers over its lifetime) for a virtual thread, even though we will only place a single fiber in there
3. nested layers of weak references, which we know that the GC just loves /s
   Specifically we hold a weak reference to a bag which itself hold a weak reference to the fiber.

[armanbilge]

Fixed the leak in #3964 and opened a follow-up.

• More efficient monitoring of fibers on virtual threads #4014