Refactoring of Cancellables

Author: odersky

tests/run/suspend-strawman-2/Async.scala

@@ -1,8 +1,6 @@
 package concurrent
 import java.util.concurrent.atomic.AtomicBoolean
 import scala.collection.mutable
-import fiberRuntime.suspend
-import fiberRuntime.boundary
 
 /** A context that allows to suspend waiting for asynchronous data sources
  */
@@ -11,13 +9,16 @@ trait Async extends Async.Config:
   /** Wait for completion of async source `src` and return the result */
   def await[T](src: Async.Source[T]): T
 
+  /** An Async of the same kind as this one, with a given cancellation group */
+  def withGroup(group: Cancellable.Group): Async
+
 object Async:
 
   /** The underlying configuration of an async block */
   trait Config:
 
-    /** The cancellable async source underlying this async computation */
-    def root: Cancellable
+    /** The group of cancellableup to which nested futures belong */
+    def group: Cancellable.Group
 
     /** The scheduler for runnables defined in this async computation */
     def scheduler: Scheduler
@@ -28,59 +29,13 @@ object Async:
      *  ignores cancellation requests
      */
     given fromScheduler(using s: Scheduler): Config with
-      def root = Cancellable.empty
+      def group = Cancellable.Unlinked
       def scheduler = s
 
   end Config
 
-  /** A possible implementation of Async. Defines an `await` method based
-   *  on a method to check for cancellation that needs to be implemented by
-   *  subclasses.
-   *
-   *   @param   root      the root of the Async's config
-   *   @param   scheduler the scheduler of the Async's config
-   *   @param   label     the label of the boundary that defines the representedd async block
-   */
-  abstract class Impl(val root: Cancellable, val scheduler: Scheduler)
-     (using label: boundary.Label[Unit]) extends Async:
-
-    protected def checkCancellation(): Unit
-
-    /** Await a source first by polling it, and, if that fails, by suspending
-     *  in a onComplete call.
-     */
-    def await[T](src: Source[T]): T =
-      checkCancellation()
-      src.poll().getOrElse:
-        try
-          var result: Option[T] = None // Not needed if we have full continuations
-          suspend[T, Unit]: k =>
-            src.onComplete: x =>
-              scheduler.schedule: () =>
-                result = Some(x)
-                k.resume()
-              true // signals to `src` that result `x` was consumed
-          result.get
-          /* With full continuations, the try block can be written more simply as follows:
-
-            suspend[T, Unit]: k =>
-              src.onComplete: x =>
-                scheduler.schedule: () =>
-                  k.resume(x)
-              true
-          */
-        finally checkCancellation()
-
-  end Impl
-
   /** An implementation of Async that blocks the running thread when waiting */
-  private class Blocking(val scheduler: Scheduler = Scheduler) extends Async:
-
-    def root = Cancellable.empty
-
-    protected def checkCancellation(): Unit = ()
-
-    private var hasResumed = false
+  private class Blocking(val scheduler: Scheduler, val group: Cancellable.Group) extends Async:
 
     def await[T](src: Source[T]): T =
       src.poll().getOrElse:
@@ -94,18 +49,26 @@ object Async:
           while result.isEmpty do wait()
           result.get
 
+    def withGroup(group: Cancellable.Group) = Blocking(scheduler, group)
+  end Blocking
+
+
   /** Execute asynchronous computation `body` on currently running thread.
    *  The thread will suspend when the computation waits.
    */
   def blocking[T](body: Async ?=> T, scheduler: Scheduler = Scheduler): T =
-    body(using Blocking())
+    body(using Blocking(scheduler, Cancellable.Unlinked))
 
   /** The currently executing Async context */
   inline def current(using async: Async): Async = async
 
   /** Await source result in currently executing Async context */
   inline def await[T](src: Source[T])(using async: Async): T = async.await(src)
 
+  def group[T](body: Async ?=> T)(using async: Async): T =
+    val newGroup = Cancellable.Group().link()
+    body(using async.withGroup(newGroup))
+
   /** A function `T => Boolean` whose lineage is recorded by its implementing
    *  classes. The Listener function accepts values of type `T` and returns
    *  `true` iff the value was consumed by an async block.

tests/run/suspend-strawman-2/Cancellable.scala

@@ -1,21 +1,62 @@
 package concurrent
+import scala.collection.mutable
 
 /** A trait for cancellable entities that can be grouped */
 trait Cancellable:
 
+  private var group: Cancellable.Group = Cancellable.Unlinked
+
   /** Issue a cancel request */
   def cancel(): Unit
 
-  /** Add a given child to this Cancellable, so that the child will be cancelled
-   *  when the Cancellable itself is cancelled.
+  /** Add this cancellable to the given group after removing
+   *  it from the previous group in which it was.
+   */
+  def link(group: Cancellable.Group): this.type =
+    this.group.drop(this)
+    this.group = group
+    this.group.add(this)
+    this
+
+  /** Link this cancellable to the cancellable group of the
+   *  current async context.
    */
-  def addChild(child: Cancellable): Unit
+  def link()(using ac: Async): this.type =
+    link(ac.group)
+
+  /** Unlink this cancellable from its group. */
+  def unlink(): this.type =
+    link(Cancellable.Unlinked)
 
 object Cancellable:
 
-  /** A cancelled entity that ignores all `cancel` and `addChild` requests */
-  object empty extends Cancellable:
-    def cancel() = ()
-    def addChild(child: Cancellable) = ()
+  /** A group of cancellable members */
+  class Group extends Cancellable:
+    private var members: mutable.Set[Cancellable] = mutable.Set()
+
+    /** Cancel all members and clear the members set */
+    def cancel() =
+      members.toArray.foreach(_.cancel())
+      members.clear()
+
+    /** Add given member to the members set */
+    def add(member: Cancellable): Unit = synchronized:
+      members += member
+
+    /** Remove given member from the members set if it is an element */
+    def drop(member: Cancellable): Unit = synchronized:
+      members -= member
+  end Group
+
+  /** A sentinal group of cancellables that are in fact not linked
+   *  to any real group. `cancel`, `add`, and `drop` do nothing when
+   *  called on this group.
+   */
+  object Unlinked extends Group:
+    override def cancel() = ()
+    override def add(member: Cancellable): Unit = ()
+    override def drop(member: Cancellable): Unit = ()
+  end Unlinked
 
 end Cancellable
+

tests/run/suspend-strawman-2/futures.scala

@@ -1,6 +1,7 @@
 package concurrent
 
 import scala.collection.mutable, mutable.ListBuffer
+import fiberRuntime.suspend
 import fiberRuntime.boundary
 import scala.compiletime.uninitialized
 import scala.util.{Try, Success, Failure}
@@ -17,15 +18,10 @@ trait Future[+T] extends Async.OriginalSource[Try[T]], Cancellable:
   def value(using async: Async): T
 
   /** Eventually stop computation of this future and fail with
-   *  a `Cancellation` exception. Also cancel all children.
+   *  a `Cancellation` exception.
    */
   def cancel(): Unit
 
-  /** If this future has not yet completed, add `child` so that it will
-   *  be cancelled together with this future in case the future is cancelled.
-   */
-  def addChild(child: Cancellable): Unit
-
 object Future:
 
   /**  A future that is completed explicitly by calling its
@@ -37,14 +33,9 @@ object Future:
   private class CoreFuture[+T] extends Future[T]:
 
     @volatile protected var hasCompleted: Boolean = false
+    protected var cancelRequest = false
     private var result: Try[T] = uninitialized // guaranteed to be set if hasCompleted = true
     private val waiting: mutable.Set[Try[T] => Boolean] = mutable.Set()
-    private val children: mutable.Set[Cancellable] = mutable.Set()
-
-    private def extract[T](s: mutable.Set[T]): List[T] = synchronized:
-      val xs = s.toList
-      s.clear()
-      xs
 
     // Async.Source method implementations
 
@@ -60,16 +51,7 @@ object Future:
     // Cancellable method implementations
 
     def cancel(): Unit =
-      val othersToCancel = synchronized:
-        if hasCompleted then Nil
-        else
-          result = Failure(new CancellationException())
-          hasCompleted = true
-          extract(children)
-      othersToCancel.foreach(_.cancel())
-
-    def addChild(child: Cancellable): Unit = synchronized:
-      if !hasCompleted then children += this
+      cancelRequest = true
 
     // Future method implementations
 
@@ -86,30 +68,68 @@ object Future:
      *     the type with which the future was created since `Promise` is invariant.
      */
     private[Future] def complete(result: Try[T] @uncheckedVariance): Unit =
-      if !hasCompleted then
-        this.result = result
-        hasCompleted = true
-      for listener <- extract(waiting) do listener(result)
+      val toNotify = synchronized:
+        if hasCompleted then Nil
+        else
+          this.result = result
+          hasCompleted = true
+          val ws = waiting.toList
+          waiting.clear()
+          ws
+      for listener <- toNotify do listener(result)
 
   end CoreFuture
 
   /** A future that is completed by evaluating `body` as a separate
    *  asynchronous operation in the given `scheduler`
    */
-  private class RunnableFuture[+T](body: Async ?=> T)(using scheduler: Scheduler)
+  private class RunnableFuture[+T](body: Async ?=> T)(using ac: Async.Config)
   extends CoreFuture[T]:
 
     /** a handler for Async */
     private def async(body: Async ?=> Unit): Unit =
+      class FutureAsync(val scheduler: Scheduler, val group: Cancellable.Group) extends Async:
+
+        def checkCancellation() =
+          if cancelRequest then throw CancellationException()
+
+        /** Await a source first by polling it, and, if that fails, by suspending
+         *  in a onComplete call.
+         */
+        def await[T](src: Async.Source[T]): T =
+          checkCancellation()
+          src.poll().getOrElse:
+            try
+              var result: Option[T] = None // Not needed if we have full continuations
+              suspend[T, Unit]: k =>
+                src.onComplete: x =>
+                  scheduler.schedule: () =>
+                    result = Some(x)
+                    k.resume()
+                  true // signals to `src` that result `x` was consumed
+              result.get
+              /* With full continuations, the try block can be written more simply as follows:
+
+                suspend[T, Unit]: k =>
+                  src.onComplete: x =>
+                    scheduler.schedule: () =>
+                      k.resume(x)
+                  true
+              */
+            finally checkCancellation()
+
+        def withGroup(group: Cancellable.Group) = FutureAsync(scheduler, group)
+
       boundary [Unit]:
-        given Async = new Async.Impl(this, scheduler):
-          def checkCancellation() =
-            if hasCompleted then throw new CancellationException()
-        body
+        body(using FutureAsync(ac.scheduler, ac.group))
     end async
 
-    scheduler.schedule: () =>
-      async(complete(Try(body)))
+    ac.scheduler.schedule: () =>
+      async:
+        link()
+        Async.group:
+          complete(Try(body))
+        unlink()
 
   end RunnableFuture
 
@@ -119,9 +139,7 @@ object Future:
    *  children of that context's root.
    */
   def apply[T](body: Async ?=> T)(using ac: Async.Config): Future[T] =
-    val f = RunnableFuture(body)(using ac.scheduler)
-    ac.root.addChild(f)
-    f
+    RunnableFuture(body)
 
   extension [T](f1: Future[T])
 

tests/run/suspend-strawman-2/scheduler.scala

@@ -5,6 +5,5 @@ trait Scheduler:
   def schedule(task: Runnable): Unit = task.run()
 
 object Scheduler extends Scheduler:
-  given fromAsyncConfig(using ac: Async.Config): Scheduler = ac.scheduler
 end Scheduler