Replace synchronized(this) with lock-free initState in LazyList

JDK 23+ allocates `ObjectMonitor` instances in native memory eagerly in
the single-threaded case under nested recursive synchronized blocks.

This can lead to huge memory consumption while evaluating LazyLists.

This commit replaces synchronized with a lock-free coordination.

Author: lrytz

library-js/src/scala/collection/immutable/LazyListBase.scala

@@ -0,0 +1,51 @@
+/*
+ * Scala (https://www.scala-lang.org)
+ *
+ * Copyright EPFL and Lightbend, Inc. dba Akka
+ *
+ * Licensed under Apache License 2.0
+ * (http://www.apache.org/licenses/LICENSE-2.0).
+ *
+ * See the NOTICE file distributed with this work for
+ * additional information regarding copyright ownership.
+ */
+
+package scala.collection.immutable
+
+import scala.language.`2.13`
+
+/**
+ * Base class for [[LazyList]] to split out code that uses concurrency utilities that are not available on Scala.js.
+ */
+abstract class LazyListBase[+A] private[immutable] (initialTail: AnyRef | Null) extends AbstractSeq[A] with Serializable {
+  /** See [[LazyList._head]] for the possible states of this field. */
+  @volatile private var _tail: AnyRef | Null /* () => LazyList[A] | Thread | InRace | LazyList[A] | Null */ = initialTail
+
+  private[immutable] def rawTail: AnyRef | Null = _tail
+
+  private[immutable] def setRawTail(value: AnyRef): Unit = _tail = value
+
+  private[immutable] def makeTailUpdater: LazyListBase.TailUpdater = LazyListBase.TailUpdater()
+}
+
+private[immutable] object LazyListBase {
+  final class TailUpdater {
+    @inline def compareAndSet(ll: LazyListBase[?], expected: AnyRef, value: AnyRef): Boolean =
+      if (ll._tail eq expected) { ll._tail = value; true } else false
+
+    @inline def getAndSet(ll: LazyListBase[?], value: AnyRef | Null): AnyRef | Null = {
+      val old = ll._tail
+      ll._tail = value
+      old
+    }
+  }
+
+  def isCurrentThread(t: Thread): Boolean = true
+
+  def InRace(t: Thread): InRace = throw new Exception("unreachable")
+
+  final class InRace private[LazyListBase] (val owner: Thread) {
+    def await(): Unit = ()
+    def countDown(): Unit = ()
+  }
+}

library/src/scala/collection/immutable/LazyList.scala

@@ -17,9 +17,9 @@ package immutable
 import scala.language.`2.13`
 import java.io.{ObjectInputStream, ObjectOutputStream}
 import java.lang.{StringBuilder => JStringBuilder}
-
 import scala.annotation.tailrec
 import scala.collection.generic.SerializeEnd
+import scala.collection.immutable.LazyListBase.InRace
 import scala.collection.mutable.{Builder, ReusableBuilder, StringBuilder}
 import scala.language.implicitConversions
 import scala.runtime.Statics
@@ -264,7 +264,7 @@ import scala.runtime.Statics
  */
 @SerialVersionUID(4L)
 final class LazyList[+A] private (lazyState: AnyRef /* EmptyMarker.type | () => LazyList[A] */)
-  extends AbstractSeq[A]
+  extends LazyListBase[A](if (lazyState eq LazyList.EmptyMarker) null else lazyState)
     with LinearSeq[A]
     with LinearSeqOps[A, LazyList, LazyList[A]]
     with IterableFactoryDefaults[A, LazyList]
@@ -276,56 +276,86 @@ final class LazyList[+A] private (lazyState: AnyRef /* EmptyMarker.type | () =>
   private def this(head: A, tail: LazyList[A]) = {
     this(LazyList.EmptyMarker)
     _head = head
-    _tail = tail
+    setRawTail(tail)
   }
 
-  // used to synchronize lazy state evaluation
-  // after initialization (`_head ne Uninitialized`)
+  // `_head` and `_tail` are used to synchronize lazy state evaluation.
+  //
+  // initially, `_head` is `Uninitialized`. after initialization, `_head` holds:
   //   - `null` if this is an empty lazy list
-  //   - `head: A` otherwise (can be `null`, `_tail == null` is used to test emptiness)
+  //   - `head: A` otherwise (can be the `null` value, `_tail == null` is used to test emptiness)
+  //
+  // `_tail` (declared in `LazyListBase`) can hold the following values:
+  //   - when `_head eq Uninitialized`
+  //     - `lazyState: () => LazyList[A]`
+  //     - while evaluating `lazyState`: the evaluating `Thread`
+  //     - if multiple threads attempt initialization: an `InRace` instance
+  //   - when `_head ne Uninitialized`
+  //     - `null` if this is an empty lazy list
+  //     - `tail: LazyList[A]` otherwise
   @volatile private var _head: Any /* Uninitialized | A */ =
     if (lazyState eq EmptyMarker) null else Uninitialized
 
-  // when `_head eq Uninitialized`
-  //   - `lazySate: () => LazyList[A]`
-  //   - MidEvaluation while evaluating lazyState
-  // when `_head ne Uninitialized`
-  //   - `null` if this is an empty lazy list
-  //   - `tail: LazyList[A]` otherwise
-  private var _tail: AnyRef | Null /* () => LazyList[A] | MidEvaluation.type | LazyList[A] | Null */ =
-    if (lazyState eq EmptyMarker) null else lazyState
-
   private def rawHead: Any = _head
-  private def rawTail: AnyRef | Null = _tail
 
   @inline private def isEvaluated: Boolean = _head.asInstanceOf[AnyRef] ne Uninitialized
 
-  private def initState(): Unit = synchronized {
-    if (!isEvaluated) {
-      // if it's already mid-evaluation, we're stuck in an infinite
-      // self-referential loop (also it's empty)
-      if (_tail eq MidEvaluation)
-        throw new RuntimeException(
-          "LazyList evaluation depends on its own result (self-reference); see docs for more info")
-
-      val fun = _tail.asInstanceOf[() => LazyList[A]]
-      _tail = MidEvaluation
-      val l =
-        // `fun` returns a LazyList that represents the state (head/tail) of `this`. We call `l.evaluated` to ensure
-        // `l` is initialized, to prevent races when reading `rawTail` / `rawHead` below.
-        // Often, lazy lists are created with `newLL(eagerCons(...))` so `l` is already initialized, but `newLL` also
-        // accepts non-evaluated lazy lists.
-        try fun().evaluated
-        // restore `fun` in finally so we can try again later if an exception was thrown (similar to lazy val)
-        finally _tail = fun
-      _tail = l.rawTail
-      _head = l.rawHead
+  private def initState(): Unit = {
+    def selfRef(): Nothing =
+      // if it's already mid-evaluation, we're stuck in an infinite self-referential loop (also it's empty)
+      throw new RuntimeException(
+        "LazyList evaluation depends on its own result (self-reference); see docs for more info")
+
+    while (!isEvaluated) {
+      rawTail match {
+        case t: Thread =>
+          if (LazyListBase.isCurrentThread(t)) selfRef()
+          val ir = InRace(t)
+          if (_tailUpdater.compareAndSet(this, t, ir))
+            ir.await()
+          // loop on lost CAS
+
+        case ir: InRace =>
+          if (LazyListBase.isCurrentThread(ir.owner)) selfRef()
+          ir.await()
+
+        case fun: Function0[_] =>
+          // use the current thread as marker that `fun` is being evaluated.
+          // this way, there is no allocation in the common case where there's no race.
+          // if multiple threads attempt to initialize a LazyList, an `InRace` instance is created to coordinate.
+          if (_tailUpdater.compareAndSet(this, fun, Thread.currentThread)) {
+            var ex: Throwable | Null = null
+            // `fun` returns a LazyList that represents the state (head/tail) of `this`. We call `evaluated` to ensure
+            // the result is initialized, to prevent races when reading `rawTail` / `rawHead` below.
+            // Often, lazy lists are created with `newLL(eagerCons(...))` so `l` is already initialized, but `newLL`
+            // also accepts non-evaluated lazy lists.
+            val l = try fun().asInstanceOf[LazyList[A]].evaluated catch {
+              case t: Throwable =>
+                ex = t
+                null
+            }
+            // update `_tail` before `_head`, because `_head` is used to test `isEvaluated`
+            val newTail = if (ex == null) l.nn.rawTail else fun
+            val sentinel = _tailUpdater.getAndSet(this, newTail)
+            if (ex == null) _head = l.nn.rawHead
+            sentinel match {
+              case ir: InRace => ir.countDown()
+              case _ =>
+            }
+            if (ex != null) throw ex.nn
+          }
+          // loop on lost CAS
+
+        case _ =>
+          // loop when _tail is a LazyList but _head is still `Uninitialized`
+          // could call `Thread.onSpinWait()` on JDK 9+
+      }
     }
   }
 
   @tailrec private def evaluated: LazyList[A] =
     if (isEvaluated) {
-      if (_tail == null) Empty
+      if (rawTail == null) Empty
       else this
     } else {
       initState()
@@ -1160,11 +1190,13 @@ object LazyList extends SeqFactory[LazyList] {
   // def kount(): Unit = k += 1
 
   private object Uninitialized extends Serializable
-  private object MidEvaluation
   private object EmptyMarker
 
   private val Empty: LazyList[Nothing] = new LazyList(EmptyMarker)
 
+  // lazy val to break cycle (Predef -> scala.package -> val LazyList -> makeTailUpdater -> Predef.classOf)
+  private lazy val _tailUpdater: LazyListBase.TailUpdater = Empty.makeTailUpdater
+
   /** Creates a new LazyList.
    *
    *  @tparam A the element type of the lazy list

library/src/scala/collection/immutable/LazyListBase.scala

@@ -0,0 +1,66 @@
+/*
+ * Scala (https://www.scala-lang.org)
+ *
+ * Copyright EPFL and Lightbend, Inc. dba Akka
+ *
+ * Licensed under Apache License 2.0
+ * (http://www.apache.org/licenses/LICENSE-2.0).
+ *
+ * See the NOTICE file distributed with this work for
+ * additional information regarding copyright ownership.
+ */
+
+package scala.collection.immutable
+
+import scala.language.`2.13`
+
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.atomic.AtomicReferenceFieldUpdater
+
+/**
+ * Base class for [[LazyList]] to split out code that uses concurrency utilities that are not available
+ * on Scala.js. This way, Scala.js does not need to override all of LazyList.
+ *
+ * This class cannot be a trait because `AtomicReferenceFieldUpdater.newUpdater` checks if the caller
+ * class has access to the corresponding field. So it needs to be called in the class where the field is
+ * declared (fields are always private in Scala).
+ */
+abstract class LazyListBase[+A] private[immutable] (initialTail: AnyRef | Null) extends AbstractSeq[A] with Serializable {
+  /** See [[LazyList._head]] for the possible states of this field. */
+  @volatile private var _tail: AnyRef | Null /* () => LazyList[A] | Thread | InRace | LazyList[A] | Null */ = initialTail
+
+  private[immutable] def rawTail: AnyRef | Null = _tail
+
+  private[immutable] def setRawTail(value: AnyRef): Unit = _tail = value
+
+  @noinline private[immutable] def makeTailUpdater: LazyListBase.TailUpdater =
+    new LazyListBase.TailUpdater(AtomicReferenceFieldUpdater.newUpdater(classOf[LazyListBase[?]], classOf[AnyRef], "_tail"))
+}
+
+private[immutable] object LazyListBase {
+  final class TailUpdater(u: AtomicReferenceFieldUpdater[LazyListBase[?], AnyRef]) {
+    def compareAndSet(ll: LazyListBase[?], expected: AnyRef, value: AnyRef): Boolean = u.compareAndSet(ll, expected, value)
+    def getAndSet(ll: LazyListBase[?], value: AnyRef | Null): AnyRef | Null = u.getAndSet(ll, value)
+  }
+
+  // this utility is constant `true` on Scala.js -> enables DCE in LazyList
+  def isCurrentThread(t: Thread): Boolean = t eq Thread.currentThread
+  // also for Scala.js
+  def InRace(t: Thread): InRace = new InRace(t)
+
+  final class InRace private[LazyListBase] (val owner: Thread) {
+    private val done: CountDownLatch = new CountDownLatch(1)
+
+    def await(): Unit = {
+      var interrupted = false
+      while (done.getCount > 0) {
+        try done.await() catch {
+          case _: InterruptedException => interrupted = true
+        }
+      }
+      if (interrupted) Thread.currentThread().interrupt()
+    }
+
+    def countDown(): Unit = done.countDown()
+  }
+}

library/test/scala/collection/immutable/LazyListTest.scala

@@ -8,6 +8,10 @@ import scala.annotation.unused
 import scala.collection.immutable.LazyListTest.sd
 import scala.collection.mutable.{Builder, ListBuffer}
 import tools.AssertUtil
+
+import java.io.NotSerializableException
+import java.util.concurrent.atomic.AtomicInteger
+import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch}
 import scala.util.Try
 
 class LazyListTest {
@@ -52,6 +56,57 @@ class LazyListTest {
     assertTrue(ll.toList == List(1, 2))
   }
 
+  @Test def racySerialization(): Unit = {
+    import sd._
+    val ll = 1 #:: { Thread.sleep(500); 2} #:: LazyList.empty
+    new Thread(() => println(ll.toList)).start()
+    Thread.sleep(200)
+    AssertUtil.assertThrows[NotSerializableException](serialize(ll), msg => msg.contains("java.lang.Thread") || msg.contains("InRace"))
+  }
+
+  @Test def forceSameCellManyThreads(): Unit = {
+    val N = 8
+    val trials = 5
+    for (trial <- 1 to trials) {
+      val counts = Array.fill(20)(new AtomicInteger(0))
+      val ll = LazyList.tabulate(20) { i =>
+        counts(i).incrementAndGet()
+        i * 10
+      }
+      val go = new CountDownLatch(1)
+      val done = new CountDownLatch(N)
+      val results = new ConcurrentLinkedQueue[List[Int]]
+      for (_ <- 1 to N)
+        new Thread(() => {
+          go.await()
+          results.add(ll.toList)
+          done.countDown()
+        }).start()
+      go.countDown()
+      done.await()
+
+      val expected = (0 until 20).map(_ * 10).toList
+      val it = results.iterator
+      var seen = 0
+      while (it.hasNext) { assertEquals(s"trial $trial", expected, it.next()); seen += 1 }
+      assertEquals(s"trial $trial: result count", N, seen)
+      for (i <- 0 until 20)
+        assertEquals(s"trial $trial: element $i evaluated multiple times", 1, counts(i).get())
+    }
+  }
+
+  @Test def initStateExceptionRecovery(): Unit = {
+    var n = 0
+    val ll = 1 #:: {
+      n += 1
+      if (n == 1) throw new RuntimeException("first attempt throws") else 2
+    } #:: LazyList.empty
+
+    AssertUtil.assertThrows[RuntimeException](ll.toList, _.contains("first attempt throws"))
+    assertEquals(List(1, 2), ll.toList)
+    assertEquals(2, n)
+  }
+
   @Test def storeNull(): Unit = {
     val l = "1" #:: null #:: "2" #:: LazyList.empty
     assert(l.toList == List("1", null, "2"))

project/MiMaFilters.scala

@@ -24,8 +24,15 @@ object MiMaFilters {
 
       // Breaking changes since last reference version
       Build.mimaPreviousDottyVersion -> Seq(
+
+        // scala/scala3#26100
+        ProblemFilters.exclude[MissingTypesProblem]("scala.collection.immutable.LazyList"),
+        ProblemFilters.exclude[MissingTypesProblem]("scala.collection.immutable.LazyList$MidEvaluation$"),
+        ProblemFilters.exclude[DirectMissingMethodProblem]("scala.collection.immutable.LazyList.<clinit>"),
+
         ProblemFilters.exclude[DirectMissingMethodProblem]("scala.None.orNull"),
         ProblemFilters.exclude[MissingTypesProblem]("scala.util.control.NonLocalReturns$ReturnThrowable"),
+
         // THIS IS FINE, IT SHOULD HAVE BEEN THIS WAY
         ProblemFilters.exclude[MissingTypesProblem]("scala.Function1$"),
         ProblemFilters.exclude[MissingTypesProblem]("scala.Function1$UnliftOps$"),