Introduce ThreadHerd instead of StructuredConcurrencyScope, broaden Java compatibility

Author: adamw

core/src/main/scala/ox/Ox.scala

@@ -7,6 +7,7 @@ import java.util.concurrent.StructuredTaskScope
 import java.util.concurrent.atomic.AtomicReference
 import scala.annotation.implicitNotFound
 import ox.channels.BufferCapacity
+import ox.internal.ThreadHerd
 
 /** Capability granted by an [[unsupervised]] concurrency scope (as well as, via subtyping, by [[supervised]] and [[supervisedError]]).
   *
@@ -22,7 +23,7 @@ import ox.channels.BufferCapacity
   "This operation must be run within a `supervised`, `supervisedError` or `unsupervised` block. Alternatively, you must require that the enclosing method is run within a scope, by adding a `using OxUnsupervised` parameter list."
 )
 trait OxUnsupervised:
-  private[ox] def scope: StructuredTaskScope[Any]
+  private[ox] def herd: ThreadHerd
   private[ox] def finalizers: AtomicReference[List[() => Unit]]
   private[ox] def supervisor: Supervisor[Nothing]
   private[ox] def addFinalizer(f: () => Unit): Unit = finalizers.updateAndGet(f :: _).discard
@@ -68,16 +69,16 @@ private[ox] trait RunInScope:
   "This operation must be run within a `supervisedError` block. Alternatively, you must require that the enclosing method is run within a scope, by adding a `using OxError[E, F]` parameter list, using the desired error mode type parameters."
 )
 case class OxError[E, F[_]](
-    private[ox] val scope: StructuredTaskScope[Any],
+    private[ox] val herd: ThreadHerd,
     private[ox] val finalizers: AtomicReference[List[() => Unit]],
     private[ox] val supervisor: Supervisor[E],
     private[ox] val errorMode: ErrorMode[E, F]
 ) extends Ox:
   override private[ox] def asNoErrorMode: OxError[Nothing, [T] =>> T] =
     if errorMode == NoErrorMode then this.asInstanceOf[OxError[Nothing, [T] =>> T]]
-    else OxError(scope, finalizers, supervisor, NoErrorMode)
+    else OxError(herd, finalizers, supervisor, NoErrorMode)
 end OxError
 
 object OxError:
   def apply[E, F[_]](s: Supervisor[E], em: ErrorMode[E, F]): OxError[E, F] =
-    OxError(DoNothingScope[Any](oxThreadFactory), new AtomicReference(Nil), s, em)
+    OxError(ThreadHerd(oxThreadFactory), new AtomicReference(Nil), s, em)

core/src/main/scala/ox/fork.scala

@@ -34,15 +34,16 @@ def forkError[E, F[_], T](using OxError[E, F])(f: => F[T]): Fork[T] =
   val oxError = summon[OxError[E, F]]
   // the separate result future is needed to wait for the result, as there's no .join on individual tasks (only whole scopes can be joined)
   val result = new CompletableFuture[T]()
-  oxError.scope.fork { () =>
+
+  oxError.herd.startThread:
     val supervisor = oxError.supervisor
     try
       val resultOrError = f
       val errorMode = oxError.errorMode
       if errorMode.isError(resultOrError) then
         // result is never completed, the supervisor should end the scope
         supervisor.forkAppError(errorMode.getError(resultOrError))
-      else result.complete(errorMode.getT(resultOrError))
+      else result.complete(errorMode.getT(resultOrError)).discard
     catch
       case e: Throwable =>
         // we notify the supervisor first, so that if this is the first failing fork in the scope, the supervisor will
@@ -51,7 +52,7 @@ def forkError[E, F[_], T](using OxError[E, F])(f: => F[T]): Fork[T] =
         // completing the result; any joins will end up being interrupted
         if !supervisor.forkException(e) then result.completeExceptionally(e).discard
     end try
-  }
+
   new ForkUsingResult(result) {}
 end forkError
 
@@ -81,7 +82,8 @@ def forkUserError[E, F[_], T](using OxError[E, F])(f: => F[T]): Fork[T] =
   val oxError = summon[OxError[E, F]]
   val result = new CompletableFuture[T]()
   oxError.supervisor.forkStarts()
-  oxError.scope.fork { () =>
+
+  oxError.herd.startThread:
     val supervisor = oxError.supervisor.asInstanceOf[DefaultSupervisor[E]]
     try
       val resultOrError = f
@@ -96,7 +98,7 @@ def forkUserError[E, F[_], T](using OxError[E, F])(f: => F[T]): Fork[T] =
       case e: Throwable =>
         if !supervisor.forkException(e) then result.completeExceptionally(e).discard
     end try
-  }
+
   new ForkUsingResult(result) {}
 end forkUserError
 
@@ -112,11 +114,13 @@ end forkUserError
   */
 def forkUnsupervised[T](f: => T)(using OxUnsupervised): UnsupervisedFork[T] =
   val result = new CompletableFuture[T]()
-  summon[OxUnsupervised].scope.fork { () =>
-    try result.complete(f)
-    catch case e: Throwable => result.completeExceptionally(e)
-  }
+
+  summon[OxUnsupervised].herd.startThread:
+    try result.complete(f).discard
+    catch case e: Throwable => result.completeExceptionally(e).discard
+
   new ForkUsingResult(result) with UnsupervisedFork[T] {}
+end forkUnsupervised
 
 /** For each thunk in the given sequence, starts a fork using [[fork]]. All forks are guaranteed to complete before the enclosing
   * [[supervised]] or [[unsupervised]] block completes.
@@ -149,21 +153,27 @@ def forkCancellable[T](f: => T)(using OxUnsupervised): CancellableFork[T] =
   // interrupt signal
   val done = new Semaphore(0)
   val ox = summon[OxUnsupervised]
-  ox.scope.fork { () =>
-    val nestedOx = OxError(NoOpSupervisor, NoErrorMode)
-    scopedWithCapability(nestedOx) {
-      nestedOx.scope.fork { () =>
-        // "else" means that the fork is already cancelled, so doing nothing in that case
-        if !started.getAndSet(true) then
-          try result.complete(f).discard
-          catch case e: Throwable => result.completeExceptionally(e).discard
-
-        done.release() // the nested scope can now finish
+
+  ox.herd.startThread:
+    try
+      val nestedOx = OxError(NoOpSupervisor, NoErrorMode)
+      scopedWithCapability(nestedOx) {
+        nestedOx.herd.startThread {
+          // "else" means that the fork is already cancelled, so doing nothing in that case
+          if !started.getAndSet(true) then
+            try result.complete(f).discard
+            catch case e: Throwable => result.completeExceptionally(e).discard
+
+          done.release() // the nested scope can now finish
+        }.discard
+
+        done.acquire()
       }
+    catch
+      // if this thread was interrupted, any context is already captured as part of the thread started in the nested scope
+      // hence, ignoring this exception, so that it's not logged by the uncaught exception handler
+      case e: InterruptedException =>
 
-      done.acquire()
-    }
-  }
   new ForkUsingResult(result) with CancellableFork[T]:
     override def cancel(): Either[Throwable, T] =
       cancelNow()

core/src/main/scala/ox/internal/ThreadHerd.scala

@@ -0,0 +1,81 @@
+package ox.internal
+
+import ox.discard
+
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.ThreadFactory
+import java.util.concurrent.atomic.AtomicBoolean
+
+/** Until the structured concurrency JEP is stable in an LTS release, a replacement for `StructuredTaskScope`. Downside: does not integrate
+  * with scoped values. Upside: works with any Java 21+.
+  *
+  * Naming: `ThreadFlock` is part of the Structured Concurrency JEP; `ThreadGroup` is already part of the JDK. Hence, using "herd."
+  */
+private[ox] class ThreadHerd(threadFactory: ThreadFactory):
+  private val herdOwner = Thread.currentThread()
+
+  private val shutdownInProgress = new AtomicBoolean(false)
+  private val threads = ConcurrentHashMap.newKeySet[Thread]()
+
+  def startThread(t: => Unit): Unit =
+    assertNotShuttingDown()
+    assertOnOwnerOrHerdThread()
+
+    val thread = threadFactory.newThread(() =>
+      try t
+      finally threads.remove(Thread.currentThread()).discard
+    )
+    threads.add(thread)
+
+    thread.start()
+  end startThread
+
+  /** Interrupts all thread in this heard, and waits for them to complete. If this process is interrupted, the exception is captured and
+    * rethrown only after all threads from the herd have completed.
+    */
+  def interruptAllAndJoinUntilCompleted(): Unit =
+    assertOnOwnerThread()
+
+    shutdownInProgress.set(true)
+
+    /*
+    If a startThread() is in progress, after the no-shutdown-assertion, the thread that it creates will be discovered
+    in the loop below, as the thread that is starting the new thread itself has to be joined; and this can only happen
+    after the new thread is added to the `threads` set.
+     */
+
+    var interruptedThreads = Set.empty[Thread] // making sure we interrupt each thread only once
+    def interruptIfNeeded(t: Thread): Unit = if !interruptedThreads.contains(t) then
+      t.interrupt()
+      interruptedThreads += t
+
+    var interruptedException: InterruptedException = null
+
+    while !threads.isEmpty do
+      threads.forEach(interruptIfNeeded)
+      threads.forEach: thread =>
+        interruptIfNeeded(thread) // double-checking, as this thread might have been added after the previous forEach
+        try thread.join()
+        catch
+          case e: InterruptedException => // if the thread is not yet done, we'll join it again in the next iteration
+            if interruptedException == null then interruptedException = e
+            else interruptedException.addSuppressed(e)
+    end while
+
+    if interruptedException != null then throw interruptedException
+  end interruptAllAndJoinUntilCompleted
+
+  private def assertNotShuttingDown(): Unit =
+    if shutdownInProgress.get() then
+      throw new IllegalStateException("Scope is shutting down, cannot start new threads or join existing ones.")
+
+  private def assertOnOwnerOrHerdThread(): Unit =
+    val current = Thread.currentThread()
+    if current != herdOwner && !threads.contains(current) then
+      throw new IllegalStateException("Forks can only be started from threads that are part of the scope.")
+
+  private def assertOnOwnerThread(): Unit =
+    val current = Thread.currentThread()
+    if current != herdOwner then
+      throw new IllegalStateException("This operation can only be performed from the thread that created the scope.")
+end ThreadHerd

core/src/main/scala/ox/unsupervised.scala

@@ -1,10 +1,5 @@
 package ox
 
-import java.util.concurrent.StructuredTaskScope
-import java.util.concurrent.ThreadFactory
-
-private class DoNothingScope[T](threadFactory: ThreadFactory) extends StructuredTaskScope[T](null, threadFactory) {}
-
 /** Starts a new concurrency scope, which allows starting forks in the given code block `f`. Forks can be started using
   * [[forkUnsupervised]], and [[forkCancellable]]. All forks are guaranteed to complete before this scope completes.
   *
@@ -30,7 +25,7 @@ private[ox] def scopedWithCapability[T](capability: Ox)(f: Ox ?=> T): T =
     es.tail.foreach(e.addSuppressed)
     throw e
 
-  val scope = capability.scope
+  val herd = capability.herd
   val finalizers = capability.finalizers
   def runFinalizers(result: Either[Throwable, T]): T =
     val fs = finalizers.get
@@ -51,17 +46,10 @@ private[ox] def scopedWithCapability[T](capability: Ox)(f: Ox ?=> T): T =
     end if
   end runFinalizers
 
-  // if this is inlined manually (in the program's text), it gets incorrectly indented
-  inline def runFAndJoinScope =
-    try f(using capability)
-    finally
-      scope.shutdown()
-      scope.join().discard
-
   try
     val t =
-      try runFAndJoinScope
-      finally scope.close() // join might have been interrupted, hence the finally
+      try f(using capability)
+      finally herd.interruptAllAndJoinUntilCompleted()
 
     // running the finalizers only once we are sure that all child threads have been terminated, so that no new
     // finalizers are added, and none are lost

core/src/test/scala/ox/resilience/RateLimiterTest.scala

@@ -205,6 +205,7 @@ class RateLimiterTest extends AnyFlatSpec with Matchers with EitherValues with T
         forkUserDiscard:
           result2 = rateLimiter.runOrDrop(operation)
         forkUserDiscard:
+          sleep(500.millis)
           result3 = rateLimiter.runBlocking(operation)
         forkUserDiscard:
           // Wait for next window for 3rd operation to start, take number of operations running