Deregister stale callbacks in Future.firstCompletedOf

Author: lrytz

library/src/scala/concurrent/Future.scala

@@ -126,6 +126,11 @@ trait Future[+T] extends Awaitable[T] {
    */
   def onComplete[U](f: Try[T] => U)(implicit executor: ExecutionContext): Unit
 
+  /** The same as [[onComplete]], but additionally returns a function which can be
+   *  invoked to unregister the callback function. Removing a callback from a long-lived
+   *  future can enable garbage collection of objects referenced by the closure.
+   */
+  private[concurrent] def onCompleteWithUnregister[U](f: Try[T] => U)(implicit executor: ExecutionContext): () => Unit
 
   /* Miscellaneous */
 
@@ -616,6 +621,7 @@ object Future {
     }
 
     override final def onComplete[U](f: Try[Nothing] => U)(implicit executor: ExecutionContext): Unit = ()
+    override private[concurrent] final def onCompleteWithUnregister[U](f: Try[Nothing] => U)(implicit executor: ExecutionContext): () => Unit = () => ()
     override final def isCompleted: Boolean = false
     override final def value: Option[Try[Nothing]] = None
     override final def failed: Future[Throwable] = this
@@ -732,15 +738,25 @@ object Future {
     if (!i.hasNext) Future.never
     else {
       val p = Promise[T]()
-      val firstCompleteHandler = new AtomicReference[Promise[T]](p) with (Try[T] => Unit) {
-        override final def apply(v1: Try[T]): Unit =  {
-          val r = getAndSet(null)
-          if (r ne null)
-            r tryComplete v1 // tryComplete is likely to be cheaper than complete
+      val firstCompleteHandler = new AtomicReference(List.empty[() => Unit]) with (Try[T] => Unit) {
+        final def apply(res: Try[T]): Unit =  {
+          val deregs = getAndSet(null)
+          if (deregs != null) {
+            p.tryComplete(res) // tryComplete is likely to be cheaper than complete
+            deregs.foreach(_.apply())
+          }
+        }
+      }
+      var completed = false
+      while (i.hasNext && !completed) {
+        val deregs = firstCompleteHandler.get
+        if (deregs == null) completed = true
+        else {
+          val d = i.next().onCompleteWithUnregister(firstCompleteHandler)
+          if (!firstCompleteHandler.compareAndSet(deregs, d :: deregs))
+            d.apply()
         }
       }
-      while(i.hasNext && firstCompleteHandler.get != null) // exit early if possible
-        i.next().onComplete(firstCompleteHandler)
       p.future
     }
   }

library/src/scala/concurrent/impl/Promise.scala

@@ -215,6 +215,12 @@ private[concurrent] object Promise {
     override final def onComplete[U](func: Try[T] => U)(implicit executor: ExecutionContext): Unit =
       dispatchOrAddCallbacks(get(), new Transformation[T, Unit](Xform_onComplete, func, executor))
 
+    override private[concurrent] final def onCompleteWithUnregister[U](func: Try[T] => U)(implicit executor: ExecutionContext): () => Unit = {
+      val t = new Transformation[T, Unit](Xform_onComplete, func, executor)
+      dispatchOrAddCallbacks(get(), t)
+      () => unregisterCallback(t)
+    }
+
     override final def failed: Future[Throwable] =
       if (!get().isInstanceOf[Success[_]]) super.failed
       else Future.failedFailureFuture // Cached instance in case of already known success
@@ -319,6 +325,15 @@ private[concurrent] object Promise {
         p.dispatchOrAddCallbacks(p.get(), callbacks)
       }
 
+    @tailrec private def unregisterCallback(t: Transformation[_, _]): Unit = {
+      val state = get()
+      if (state eq t) {
+        if (!compareAndSet(state, Noop)) unregisterCallback(t)
+      } else if (state.isInstanceOf[ManyCallbacks[_]]) {
+        if (!compareAndSet(state, removeCallback(state.asInstanceOf[ManyCallbacks[T]], t))) unregisterCallback(t)
+      }
+    }
+
     // IMPORTANT: Noop should never be passed in here, neither as left OR as right
     @tailrec private[this] final def concatCallbacks(left: Callbacks[T], right: Callbacks[T]): Callbacks[T] =
       if (left.isInstanceOf[Transformation[T,_]]) new ManyCallbacks[T](left.asInstanceOf[Transformation[T,_]], right)
@@ -327,6 +342,20 @@ private[concurrent] object Promise {
         concatCallbacks(m.rest, new ManyCallbacks(m.first, right))
       }
 
+    @tailrec private[this] final def removeCallback(cs: Callbacks[T], t: Transformation[_, _], result: Callbacks[T] = null): AnyRef =
+      if (cs eq t) {
+        if (result == null) Noop
+        else result
+      }
+      else if (cs.isInstanceOf[ManyCallbacks[_]]) {
+        val m = cs.asInstanceOf[ManyCallbacks[T]]
+        if (m.first eq t) {
+          if (result == null) m.rest
+          else concatCallbacks(m.rest, result)
+        }
+        else removeCallback(m.rest, t, if (result == null) m.first else new ManyCallbacks(m.first, result))
+      } else cs
+
     // IMPORTANT: Noop should not be passed in here, `callbacks` cannot be null
     @tailrec
     private[this] final def submitWithValue(callbacks: Callbacks[T], resolved: Try[T]): Unit =