scala · SethTisue · Feb 15, 2019 · Jan 9, 2019 · Feb 4, 2019 · Feb 4, 2019
diff --git a/.gitignore b/.gitignore
@@ -50,5 +50,7 @@
 /project/project/target/
 /project/project/project/target/
 /test/macro-annot/target/
+/test/files/target/
+/test/target/
 /build-sbt/
 local.sbt
@@ -12,10 +12,10 @@
 
 package scala.concurrent
 
-import java.util.ArrayDeque
 import java.util.concurrent.Executor
-import scala.annotation.{ switch, tailrec }
+import java.util.Objects
 import scala.util.control.NonFatal
+import scala.annotation.{switch, tailrec}
 
 /**
  * Marker trait to indicate that a Runnable is Batchable by BatchingExecutors
@@ -26,6 +26,17 @@ trait Batchable {
 
 private[concurrent] object BatchingExecutorStatics {
   final val emptyBatchArray: Array[Runnable] = new Array[Runnable](0)
+
+  // Max number of Runnables executed nested before starting to batch (to prevent stack exhaustion)
+  final val syncPreBatchDepth = 16
+
+  // Max number of Runnables processed in one go (to prevent starvation of other tasks on the pool)
+  final val runLimit = 1024
+
+  final object MissingParentBlockContext extends BlockContext {
+    override def blockOn[T](thunk: => T)(implicit permission: CanAwait): T =
+      try thunk finally throw new IllegalStateException("BUG in BatchingExecutor.Batch: parentBlockContext is null")
+  }
 }
 
 /**
@@ -38,154 +49,219 @@ private[concurrent] object BatchingExecutorStatics {
  * thread which may improve CPU affinity. However,
  * if tasks passed to the Executor are blocking
  * or expensive, this optimization can prevent work-stealing
- * and make performance worse. Also, some ExecutionContext
- * may be fast enough natively that this optimization just
- * adds overhead.
- * The default ExecutionContext.global is already batching
- * or fast enough not to benefit from it; while
- * `fromExecutor` and `fromExecutorService` do NOT add
- * this optimization since they don't know whether the underlying
- * executor will benefit from it.
+ * and make performance worse.
  * A batching executor can create deadlocks if code does
  * not use `scala.concurrent.blocking` when it should,
  * because tasks created within other tasks will block
  * on the outer task completing.
  * This executor may run tasks in any order, including LIFO order.
  * There are no ordering guarantees.
  *
- * WARNING: The underlying Executor's execute-method must not execute the submitted Runnable
- * in the calling thread synchronously. It must enqueue/handoff the Runnable.
+ * WARNING: Only use *EITHER* `submitAsyncBatched` OR `submitSyncBatched`!!
+ *
+ * When you implement this trait for async executors like thread pools,
+ * you're going to need to implement it something like the following:
+ *
+ * {{{
+ *  final override def submitAsync(runnable: Runnable): Unit =
+ *    super[SuperClass].execute(runnable) // To prevent reentrancy into `execute`
+ *
+ *  final override def execute(runnable: Runnable): Unit =
+ *    if (runnable.isInstanceOf[Batchable]) // Or other logic
+ *      submitAsyncBatched(runnable)
+ *    else
+ *      submitAsync(runnable)
+ *
+ *  final override def reportFailure(cause: Throwable): Unit = …
+ *  }}}
+ *
+ *  And if you want to implement if for a sync, trampolining, executor you're
+ *  going to implement it something like this:
+ *
+ * {{{
+ *  final override def submitAsync(runnable: Runnable): Unit = ()
+ *
+ *  final override def execute(runnable: Runnable): Unit =
+ *    submitSyncBatched(runnable) // You typically will want to batch everything
+ *
+ *  final override def reportFailure(cause: Throwable): Unit =
+ *    ExecutionContext.defaultReporter(cause) // Or choose something more fitting
+ * }}}
+ *
  */
- private[concurrent] trait BatchingExecutor extends Executor {
-  private[this] final val _tasksLocal = new ThreadLocal[Batch]()
-
-  private[this] final class Batch extends Runnable with BlockContext with (BlockContext => Throwable) {
-    private[this] final var parentBlockContext: BlockContext = _
-    private[this] final var first: Runnable = _
-    private[this] final var size: Int = _
-    private[this] final var other: Array[Runnable] = BatchingExecutorStatics.emptyBatchArray
-
-    def this(r: Runnable) = {
-      this()
-      first = r
-      size = 1
-    }
+private[concurrent] trait BatchingExecutor extends Executor {
+  private[this] final val _tasksLocal = new ThreadLocal[AnyRef]()
 
-    private def this(first: Runnable, other: Array[Runnable], size: Int) = {
-      this()
-      this.first = first
-      this.other = other
-      this.size = size
-    }
+  /*
+   * Batch implements a LIFO queue (stack) and is used as a trampolining Runnable.
+   * In order to conserve allocations, the first element in the batch is stored "unboxed" in
+   * the `first` field. Subsequent Runnables are stored in the array called `other`.
+  */
+  private[this] sealed abstract class AbstractBatch protected (protected final var first: Runnable, protected final var other: Array[Runnable], protected final var size: Int) {
 
-    private[this] final def cloneAndClear(): Batch = {
-      val newBatch = new Batch(first, other, size)
-      this.first = null
-      this.other = BatchingExecutorStatics.emptyBatchArray
-      this.size = 0
-      newBatch
-    }
+    private[this] final def ensureCapacity(curSize: Int): Array[Runnable] = {
+      val curOther = this.other
+      val curLen = curOther.length
+      if (curSize <= curLen) curOther
+      else {
+        val newLen = if (curLen == 0) 4 else curLen << 1
 
-    private[this] final def grow(): Unit = {
-      val len = other.length
-      other =
-        if (len == 0) new Array[Runnable](4)
-        else {
-          val newOther = new Array[Runnable](len << 1)
-          System.arraycopy(other, 0, newOther, 0, len)
-          newOther
-        }
+        if (newLen <= curLen) throw new StackOverflowError("Space limit of asynchronous stack reached: " + curLen)
+        val newOther = new Array[Runnable](newLen)
+        System.arraycopy(curOther, 0, newOther, 0, curLen)
+        this.other = newOther
+        newOther
+      }
     }
 
     final def push(r: Runnable): Unit = {
-      val sz = size
-      if(sz > 0) {
-        if (sz > other.length)
-          grow()
-        other(sz - 1) = r
-      } else first = r
-      size = sz + 1
+      val sz = this.size
+      if(sz == 0)
+        this.first = r
+      else
+        ensureCapacity(sz)(sz - 1) = r
+      this.size = sz + 1
     }
 
-    final def pop(): Runnable =
-      (size: @switch) match {
-        case 0 => null
-        case 1 =>
-          val ret = first
-          first = null
-          size = 0
-          ret
-        case n =>
-          val ret = other(n - 2)
-          other(n - 2) = null
-          size = n - 1
-          ret
-      }
+    @tailrec protected final def runN(n: Int): Unit =
+      if (n > 0)
+        (this.size: @switch) match {
+          case 0 =>
+          case 1 =>
+            val next = this.first
+            this.first = null
+            this.size = 0
+            next.run()
+            runN(n - 1)
+          case sz =>
+            val o = this.other
+            val next = o(sz - 2)
+            o(sz - 2) = null
+            this.size = sz - 1
+            next.run()
+            runN(n - 1)
+          }
+  }
+
+  private[this] final class AsyncBatch private(_first: Runnable, _other: Array[Runnable], _size: Int) extends AbstractBatch(_first, _other, _size) with Runnable with BlockContext with (BlockContext => Throwable) {
+    private[this] final var parentBlockContext: BlockContext = BatchingExecutorStatics.MissingParentBlockContext
+
+    final def this(runnable: Runnable) = this(runnable, BatchingExecutorStatics.emptyBatchArray, 1)
 
-    // this method runs in the delegate ExecutionContext's thread
     override final def run(): Unit = {
-      //This invariant needs to hold: require(_tasksLocal.get eq null)
-      _tasksLocal.set(this)
-      val failure = BlockContext.usingBlockContext(this)(this)
-      _tasksLocal.remove()
-      if (failure ne null)
-        throw handleRunFailure(failure)
+      _tasksLocal.set(this) // This is later cleared in `apply` or `runWithoutResubmit`
+
+      val f = resubmit(BlockContext.usingBlockContext(this)(this))
+
+      if (f != null)
+        throw f
     }
 
-    override final def apply(prevBlockContext: BlockContext): Throwable = {
+    /* LOGIC FOR ASYNCHRONOUS BATCHES */
+    override final def apply(prevBlockContext: BlockContext): Throwable = try {
       parentBlockContext = prevBlockContext
-      var failure: Throwable = null
-      try {
-        var r = pop()
-        while(r ne null) {
-          r.run()
-          r = pop()
-        }
-      } catch {
-        case t: Throwable => failure = t
-      }
-      parentBlockContext = null
-      failure
+      runN(BatchingExecutorStatics.runLimit)
+      null
+    } catch {
+      case t: Throwable => t // We are handling exceptions on the outside of this method
+    } finally {
+      parentBlockContext = BatchingExecutorStatics.MissingParentBlockContext
+      _tasksLocal.remove()
     }
 
-    private[this] final def handleRunFailure(cause: Throwable): Throwable =
-      if (size > 0 && (NonFatal(cause) || cause.isInstanceOf[InterruptedException])) {
-        try { unbatchedExecute(this); cause } catch {
+    /* Attempts to resubmit this Batch to the underlying ExecutionContext,
+     * this only happens for Batches where `resubmitOnBlock` is `true`.
+     * Only attempt to resubmit when there are `Runnables` left to process.
+     * Note that `cause` can be `null`.
+     */
+    private[this] final def resubmit(cause: Throwable): Throwable =
+      if (this.size > 0) {
+        try { submitForExecution(this); cause } catch {
           case inner: Throwable =>
             if (NonFatal(inner)) {
               val e = new ExecutionException("Non-fatal error occurred and resubmission failed, see suppressed exception.", cause)
               e.addSuppressed(inner)
               e
             } else inner
         }
-      } else cause
+      } else cause // TODO: consider if NonFatals should simply be `reportFailure`:ed rather than rethrown
 
-    override def blockOn[T](thunk: => T)(implicit permission: CanAwait): T = {
-      val pbc = parentBlockContext
-      if(size > 0) // if we know there will be blocking, we don't want to keep tasks queued up because it could deadlock.
-        unbatchedExecute(cloneAndClear())
+    private[this] final def cloneAndClear(): AsyncBatch = {
+      val newBatch = new AsyncBatch(this.first, this.other, this.size)
+      this.first = null
+      this.parentBlockContext = BatchingExecutorStatics.MissingParentBlockContext
+      this.other = BatchingExecutorStatics.emptyBatchArray
+      this.size = 0
+      newBatch
+    }
 
-      if (pbc ne null) pbc.blockOn(thunk) // now delegate the blocking to the previous BC
-      else {
-        try thunk finally throw new IllegalStateException("BUG in BatchingExecutor.Batch: parentBlockContext is null")
+    override final def blockOn[T](thunk: => T)(implicit permission: CanAwait): T = {
+      val pbc = parentBlockContext // Store this for later since `cloneAndClear()` will reset it
+
+      // If we know there will be blocking, we don't want to keep tasks queued up because it could deadlock.
+      if(this.size > 0)
+        submitForExecution(cloneAndClear()) // If this throws then we have bigger problems
+
+      pbc.blockOn(thunk) // Now delegate the blocking to the previous BC
+    }
+  }
+
+  private[this] final class SyncBatch(runnable: Runnable) extends AbstractBatch(runnable, BatchingExecutorStatics.emptyBatchArray, 1) with Runnable {
+    @tailrec override final def run(): Unit = {
+      try runN(BatchingExecutorStatics.runLimit) catch {
+        case ie: InterruptedException =>
+          reportFailure(ie) // TODO: Handle InterruptedException differently?
+        case f if NonFatal(f) =>
+          reportFailure(f)
       }
+
+      if (this.size > 0)
+        run()
     }
   }
 
-  protected def unbatchedExecute(r: Runnable): Unit
+  /** MUST throw a NullPointerException when `runnable` is null
+   * When implementing a sync BatchingExecutor, it is RECOMMENDED
+   * to implement this method as `runnable.run()`
+  */
+  protected def submitForExecution(runnable: Runnable): Unit
+
+  /** Reports that an asynchronous computation failed.
+   *  See `ExecutionContext.reportFailure(throwable: Throwable)`
+  */
+  protected def reportFailure(throwable: Throwable): Unit
 
-  private[this] final def batchedExecute(runnable: Runnable): Unit = {
+  /**
+   * WARNING: Never use both `submitAsyncBatched` and `submitSyncBatched` in the same
+   * implementation of `BatchingExecutor`
+   */
+  protected final def submitAsyncBatched(runnable: Runnable): Unit = {
     val b = _tasksLocal.get
-    if (b ne null) b.push(runnable)
-    else unbatchedExecute(new Batch(runnable))
+    if (b.isInstanceOf[AsyncBatch]) b.asInstanceOf[AsyncBatch].push(runnable)
+    else submitForExecution(new AsyncBatch(runnable))
   }
 
-  override def execute(runnable: Runnable): Unit =
-    if(batchable(runnable)) batchedExecute(runnable)
-    else unbatchedExecute(runnable)
-
-  /** Override this to define which runnables will be batched.
-    * By default it tests the Runnable for being an instance of [Batchable].
-   **/
-  protected def batchable(runnable: Runnable): Boolean = runnable.isInstanceOf[Batchable]
+  /**
+   * WARNING: Never use both `submitAsyncBatched` and `submitSyncBatched` in the same
+   * implementation of `BatchingExecutor`
+   */
+  protected final def submitSyncBatched(runnable: Runnable): Unit = {
+    Objects.requireNonNull(runnable, "runnable is null")
+    val tl = _tasksLocal
+    val b = tl.get
+    if (b.isInstanceOf[SyncBatch]) b.asInstanceOf[SyncBatch].push(runnable)
+    else {
+      val i = if (b ne null) b.asInstanceOf[java.lang.Integer].intValue else 0
+      if (i < BatchingExecutorStatics.syncPreBatchDepth) {
+        tl.set(java.lang.Integer.valueOf(i + 1))
+        try submitForExecution(runnable) // User code so needs to be try-finally guarded here
+        finally tl.set(b)
+      } else {
+        val batch = new SyncBatch(runnable)
+        tl.set(batch)
+        submitForExecution(batch)
+        tl.set(b) // Batch only throws fatals so no need for try-finally here
+      }
+    }
+  }
 }
@@ -88,11 +88,10 @@ object BlockContext {
    **/
   final def withBlockContext[T](blockContext: BlockContext)(body: => T): T = {
     val old = contextLocal.get // can be null
-    try {
+    if (old eq blockContext) body
+    else {
       contextLocal.set(blockContext)
-      body
-    } finally {
-      contextLocal.set(old)
+      try body finally contextLocal.set(old)
     }
   }
 
@@ -102,7 +101,10 @@ object BlockContext {
    **/
   final def usingBlockContext[I, T](blockContext: BlockContext)(f: BlockContext => T): T = {
     val old = contextLocal.get // can be null
-    contextLocal.set(blockContext)
-    try f(prefer(old)) finally contextLocal.set(old)
+    if (old eq blockContext) f(prefer(old))
+    else {
+      contextLocal.set(blockContext)
+      try f(prefer(old)) finally contextLocal.set(old)
+    }
   }
 }