From 06ca01f9b904ab25ced0506b204210a8555232fb Mon Sep 17 00:00:00 2001
From: Nicholas Marcott <481161+bmarcott@users.noreply.github.com>
Date: Wed, 27 Nov 2019 01:58:27 -0600
Subject: [PATCH 1/4] utilize scheduable interface for delay scheduling

---
 .../org/apache/spark/scheduler/Pool.scala     | 24 ++++++++
 .../apache/spark/scheduler/Schedulable.scala  |  1 +
 .../spark/scheduler/TaskSchedulerImpl.scala   | 12 +++-
 .../spark/scheduler/TaskSetManager.scala      |  8 ++-
 .../apache/spark/scheduler/PoolSuite.scala    | 61 +++++++++++++++++++
 .../scheduler/TaskSchedulerImplSuite.scala    | 61 +++++++++++++++++++
 6 files changed, 163 insertions(+), 4 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/Pool.scala b/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
index 80805df256a15..2e25b585af8d2 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
@@ -119,4 +119,28 @@ private[spark] class Pool(
       parent.decreaseRunningTasks(taskNum)
     }
   }
+
+  override def updateAvailableSlots(numSlots: Float): Unit = {
+    schedulingMode match {
+      case SchedulingMode.FAIR =>
+        val usableWeights = schedulableQueue.asScala
+          .map(s => if (s.getSortedTaskSetQueue.nonEmpty) (s, s.weight) else (s, 0))
+        val totalWeights = usableWeights.map(_._2).sum
+        usableWeights.foreach({case (schedulable, usableWeight) =>
+          schedulable.updateAvailableSlots(
+            Math.max(numSlots * usableWeight / totalWeights, schedulable.minShare))
+        })
+      case SchedulingMode.FIFO =>
+        val sortedSchedulableQueue =
+          schedulableQueue.asScala.toSeq.sortWith(taskSetSchedulingAlgorithm.comparator)
+        var isFirst = true
+        for (schedulable <- sortedSchedulableQueue) {
+          schedulable.updateAvailableSlots(if (isFirst) numSlots else 0)
+          isFirst = false
+        }
+      case _ =>
+        val msg = s"Unsupported scheduling mode: $schedulingMode. Use FAIR or FIFO instead."
+        throw new IllegalArgumentException(msg)
+    }
+  }
 }
diff --git a/core/src/main/scala/org/apache/spark/scheduler/Schedulable.scala b/core/src/main/scala/org/apache/spark/scheduler/Schedulable.scala
index b6f88ed0a93aa..ed4b2fa4deaf0 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/Schedulable.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/Schedulable.scala
@@ -39,6 +39,7 @@ private[spark] trait Schedulable {
   def stageId: Int
   def name: String
 
+  def updateAvailableSlots(numSlots: Float): Unit
   def addSchedulable(schedulable: Schedulable): Unit
   def removeSchedulable(schedulable: Schedulable): Unit
   def getSchedulableByName(name: String): Schedulable
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
index f25a36c7af22a..7b840fe8ed1a3 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
@@ -36,7 +36,7 @@ import org.apache.spark.rpc.RpcEndpoint
 import org.apache.spark.scheduler.SchedulingMode.SchedulingMode
 import org.apache.spark.scheduler.TaskLocality.TaskLocality
 import org.apache.spark.storage.BlockManagerId
-import org.apache.spark.util.{AccumulatorV2, SystemClock, ThreadUtils, Utils}
+import org.apache.spark.util.{AccumulatorV2, Clock, SystemClock, ThreadUtils, Utils}
 
 /**
  * Schedules tasks for multiple types of clusters by acting through a SchedulerBackend.
@@ -61,7 +61,8 @@ import org.apache.spark.util.{AccumulatorV2, SystemClock, ThreadUtils, Utils}
 private[spark] class TaskSchedulerImpl(
     val sc: SparkContext,
     val maxTaskFailures: Int,
-    isLocal: Boolean = false)
+    isLocal: Boolean = false,
+    clock: Clock = new SystemClock)
   extends TaskScheduler with Logging {
 
   import TaskSchedulerImpl._
@@ -125,10 +126,10 @@ private[spark] class TaskSchedulerImpl(
 
   protected val hostsByRack = new HashMap[String, HashSet[String]]
 
+  protected  val executorIdToCores = new HashMap[String, Int]
   protected val executorIdToHost = new HashMap[String, String]
 
   private val abortTimer = new Timer(true)
-  private val clock = new SystemClock
   // Exposed for testing
   val unschedulableTaskSetToExpiryTime = new HashMap[TaskSetManager, Long]
 
@@ -403,6 +404,7 @@ private[spark] class TaskSchedulerImpl(
       if (!executorIdToRunningTaskIds.contains(o.executorId)) {
         hostToExecutors(o.host) += o.executorId
         executorAdded(o.executorId, o.host)
+        executorIdToCores(o.executorId) = o.cores
         executorIdToHost(o.executorId) = o.host
         executorIdToRunningTaskIds(o.executorId) = HashSet[Long]()
         newExecAvail = true
@@ -439,6 +441,9 @@ private[spark] class TaskSchedulerImpl(
       }
     }
 
+    val availableSlots = executorIdToCores.values.map(c => c / CPUS_PER_TASK).sum
+    rootPool.updateAvailableSlots(availableSlots)
+
     // Take each TaskSet in our scheduling order, and then offer it each node in increasing order
     // of locality levels so that it gets a chance to launch local tasks on all of them.
     // NOTE: the preferredLocality order: PROCESS_LOCAL, NODE_LOCAL, NO_PREF, RACK_LOCAL, ANY
@@ -808,6 +813,7 @@ private[spark] class TaskSchedulerImpl(
       // happen below in the rootPool.executorLost() call.
       taskIds.foreach(cleanupTaskState)
     }
+    executorIdToCores.remove(executorId)
 
     val host = executorIdToHost(executorId)
     val execs = hostToExecutors.getOrElse(host, new HashSet)
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
index 5c0bc497dd1b3..d75b18044d649 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
@@ -75,6 +75,7 @@ private[spark] class TaskSetManager(
     .map { case (t, idx) => t.partitionId -> idx }.toMap
   val numTasks = tasks.length
   val copiesRunning = new Array[Int](numTasks)
+  var numAvailableSlot: Float = 0
 
   val speculationEnabled = conf.get(SPECULATION_ENABLED)
   // Quantile of tasks at which to start speculation
@@ -412,7 +413,8 @@ private[spark] class TaskSetManager(
         taskAttempts(index) = info :: taskAttempts(index)
         // Update our locality level for delay scheduling
         // NO_PREF will not affect the variables related to delay scheduling
-        if (maxLocality != TaskLocality.NO_PREF) {
+        val utilizingAllSlots = runningTasks + 1 >= numAvailableSlot
+        if (maxLocality != TaskLocality.NO_PREF && utilizingAllSlots) {
           currentLocalityIndex = getLocalityIndex(taskLocality)
           lastLaunchTime = curTime
         }
@@ -1054,6 +1056,10 @@ private[spark] class TaskSetManager(
   def executorAdded(): Unit = {
     recomputeLocality()
   }
+
+  override def updateAvailableSlots(numSlots: Float): Unit = {
+    numAvailableSlot = numSlots
+  }
 }
 
 private[spark] object TaskSetManager {
diff --git a/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala
index b953add9d58cb..80913c3057f71 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala
@@ -50,6 +50,67 @@ class PoolSuite extends SparkFunSuite with LocalSparkContext {
     nextTaskSetToSchedule.get.addRunningTask(taskId)
     assert(nextTaskSetToSchedule.get.stageId === expectedStageId)
   }
+  test("validate FIFO slot distributions") {
+    sc = new SparkContext(LOCAL, APP_NAME)
+    val taskScheduler = new TaskSchedulerImpl(sc)
+
+    val rootPool = new Pool("", FIFO, 0, 0)
+    val schedulableBuilder = new FIFOSchedulableBuilder(rootPool)
+
+    val taskSetManager0 = createTaskSetManager(0, 2, taskScheduler)
+    val taskSetManager1 = createTaskSetManager(1, 2, taskScheduler)
+    val taskSetManager2 = createTaskSetManager(2, 2, taskScheduler)
+    schedulableBuilder.addTaskSetManager(taskSetManager0, null)
+    schedulableBuilder.addTaskSetManager(taskSetManager1, null)
+    schedulableBuilder.addTaskSetManager(taskSetManager2, null)
+
+    rootPool.updateAvailableSlots(9)
+    assert(taskSetManager0.numAvailableSlot === 9)
+    assert(taskSetManager1.numAvailableSlot === 0)
+    assert(taskSetManager2.numAvailableSlot === 0)
+  }
+
+  test("validate FAIR slot distributions") {
+    val xmlPath = getClass.getClassLoader.getResource("fairscheduler-with-valid-data.xml").getFile()
+    val conf = new SparkConf().set(SCHEDULER_ALLOCATION_FILE, xmlPath)
+    sc = new SparkContext(LOCAL, APP_NAME, conf)
+    val taskScheduler = new TaskSchedulerImpl(sc)
+
+    val rootPool = new Pool("", FAIR, 0, 0)
+    val schedulableBuilder = new FairSchedulableBuilder(rootPool, sc.conf)
+    schedulableBuilder.buildPools()
+
+    val taskSetManager0 = createTaskSetManager(0, 2, taskScheduler)
+    val taskSetManager1 = createTaskSetManager(1, 2, taskScheduler)
+    sc.setLocalProperty(SparkContext.SPARK_SCHEDULER_POOL, "pool1")
+    var localProperties = sc.getLocalProperties
+    schedulableBuilder.addTaskSetManager(taskSetManager0, localProperties)
+    schedulableBuilder.addTaskSetManager(taskSetManager1, localProperties)
+
+    sc.setLocalProperty(SparkContext.SPARK_SCHEDULER_POOL, "pool2")
+    localProperties = sc.getLocalProperties
+    val taskSetManager2 = createTaskSetManager(2, 2, taskScheduler)
+    val taskSetManager3 = createTaskSetManager(3, 2, taskScheduler)
+    schedulableBuilder.addTaskSetManager(taskSetManager2, localProperties)
+    schedulableBuilder.addTaskSetManager(taskSetManager3, localProperties)
+
+
+    rootPool.updateAvailableSlots(12)
+    // pool 1 weight is 1, pool 2 weight is 2, so they get 4 and 8 slots respectively
+    // since pool 1 is FIFO TSM 0 gets all 4 slots of pool 1
+    // since pool 2 is FAIR TSM 3 and 4 each get half (4) of pool 2's 8 slots
+    assert(taskSetManager0.numAvailableSlot === 4)
+    assert(taskSetManager1.numAvailableSlot === 0)
+    assert(taskSetManager2.numAvailableSlot === 4)
+    assert(taskSetManager3.numAvailableSlot === 4)
+
+    rootPool.updateAvailableSlots(6)
+    // similar to above but pool 1 has 3 slots due to having minShare of 3
+    assert(taskSetManager0.numAvailableSlot === 3)
+    assert(taskSetManager1.numAvailableSlot === 0)
+    assert(taskSetManager2.numAvailableSlot === 2)
+    assert(taskSetManager3.numAvailableSlot === 2)
+  }
 
   test("FIFO Scheduler Test") {
     sc = new SparkContext(LOCAL, APP_NAME)
diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
index e7ecf847ff4f4..1b0af8e4b4200 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
@@ -195,6 +195,66 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext with B
     assert(!failedTaskSet)
   }
 
+  test("executors should not sit idle for too long") {
+    val LOCALITY_WAIT_MS = 3000
+    val clock = new ManualClock
+    val conf = new SparkConf()
+    sc = new SparkContext("local", "TaskSchedulerImplSuite", conf)
+    // import org.slf4j.{Logger, LoggerFactory}
+    import org.apache.log4j.{Logger, Level}
+    Logger.getLogger("org.apache.spark.scheduler.TaskSetManager").setLevel(Level.DEBUG)
+    val taskScheduler = new TaskSchedulerImpl(sc,
+      sc.conf.get(config.TASK_MAX_FAILURES),
+      clock = clock) {
+      override def createTaskSetManager(taskSet: TaskSet, maxTaskFailures: Int): TaskSetManager = {
+        new TaskSetManager(this, taskSet, maxTaskFailures, blacklistTrackerOpt, clock)
+      }
+      override def shuffleOffers(offers: IndexedSeq[WorkerOffer]): IndexedSeq[WorkerOffer] = {
+        // Don't shuffle the offers around for this test.  Instead, we'll just pass in all
+        // the permutations we care about directly.
+        offers
+      }
+    }
+    // Need to initialize a DAGScheduler for the taskScheduler to use for callbacks.
+    new DAGScheduler(sc, taskScheduler) {
+      override def taskStarted(task: Task[_], taskInfo: TaskInfo): Unit = {}
+
+      override def executorAdded(execId: String, host: String): Unit = {}
+    }
+    taskScheduler.initialize(new FakeSchedulerBackend)
+    val taskSet = FakeTask.createTaskSet(5, 1, 1,
+      Seq(TaskLocation("host1", "exec1")),
+      Seq(TaskLocation("host1", "exec1")),
+      Seq(TaskLocation("host1", "exec1")),
+      Seq(TaskLocation("host1", "exec1")),
+      Seq(TaskLocation("host1", "exec1"))
+    )
+    taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+    taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+    taskScheduler.submitTasks(taskSet)
+
+    // First offer host2, exec2: no task should be chosen due to bad data locality
+    assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec2", "host2", 1)))
+      .flatten.isEmpty)
+    val task0 = taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+      .flatten.head
+    assert(task0.index === 0)
+    taskScheduler.statusUpdate(task0.taskId, TaskState.FINISHED, ByteBuffer.allocate(0))
+
+    clock.advance(LOCALITY_WAIT_MS * 2)
+    val task1 = taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+      .flatten.head
+    assert(task1.index === 1)
+    val task2 = taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec2", "host2", 1)))
+      .flatten.head
+    assert(task2.index === 2)
+    taskScheduler.statusUpdate(task1.taskId, TaskState.FINISHED, ByteBuffer.allocate(0))
+    assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host2", 1)))
+      .flatten.head.index === 3)
+    assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec2", "host2", 1)))
+      .flatten.isEmpty)
+  }
+
   test("Scheduler does not crash when tasks are not serializable") {
     val taskCpus = 2
     val taskScheduler = setupSchedulerWithMaster(
@@ -910,6 +970,7 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext with B
   test("SPARK-16106 locality levels updated if executor added to existing host") {
     val taskScheduler = setupScheduler()
 
+    taskScheduler.resourceOffers(IndexedSeq(new WorkerOffer("executor0", "host0", 1)))
     taskScheduler.submitTasks(FakeTask.createTaskSet(2, stageId = 0, stageAttemptId = 0,
       (0 until 2).map { _ => Seq(TaskLocation("host0", "executor2")) }: _*
     ))

From 624c41750fec72bcde7a3dfa0fba96a1a26f2ed2 Mon Sep 17 00:00:00 2001
From: Nicholas Marcott <481161+bmarcott@users.noreply.github.com>
Date: Mon, 9 Dec 2019 11:09:09 +0900
Subject: [PATCH 2/4] update total number of slots incrementally address
 formatting and doc comments

---
 .../org/apache/spark/scheduler/Pool.scala     |  5 +++--
 .../spark/scheduler/TaskSchedulerImpl.scala   | 10 +++++----
 .../spark/scheduler/TaskSetManager.scala      | 22 +++++++++++--------
 .../apache/spark/scheduler/PoolSuite.scala    |  1 +
 .../scheduler/TaskSchedulerImplSuite.scala    |  5 +----
 5 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/Pool.scala b/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
index 2e25b585af8d2..52d8de90a4e36 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
@@ -126,10 +126,11 @@ private[spark] class Pool(
         val usableWeights = schedulableQueue.asScala
           .map(s => if (s.getSortedTaskSetQueue.nonEmpty) (s, s.weight) else (s, 0))
         val totalWeights = usableWeights.map(_._2).sum
-        usableWeights.foreach({case (schedulable, usableWeight) =>
+        usableWeights.foreach { case (schedulable, usableWeight) =>
           schedulable.updateAvailableSlots(
             Math.max(numSlots * usableWeight / totalWeights, schedulable.minShare))
-        })
+        }
+      // for FIFO, allocate all slots to the first taskset in the queue
       case SchedulingMode.FIFO =>
         val sortedSchedulableQueue =
           schedulableQueue.asScala.toSeq.sortWith(taskSetSchedulingAlgorithm.comparator)
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
index 7b840fe8ed1a3..3efc885e3f058 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
@@ -126,7 +126,9 @@ private[spark] class TaskSchedulerImpl(
 
   protected val hostsByRack = new HashMap[String, HashSet[String]]
 
-  protected  val executorIdToCores = new HashMap[String, Int]
+  protected val executorIdToCores = new HashMap[String, Int]
+  protected var totalSlots = 0
+
   protected val executorIdToHost = new HashMap[String, String]
 
   private val abortTimer = new Timer(true)
@@ -405,6 +407,7 @@ private[spark] class TaskSchedulerImpl(
         hostToExecutors(o.host) += o.executorId
         executorAdded(o.executorId, o.host)
         executorIdToCores(o.executorId) = o.cores
+        totalSlots += o.cores / CPUS_PER_TASK
         executorIdToHost(o.executorId) = o.host
         executorIdToRunningTaskIds(o.executorId) = HashSet[Long]()
         newExecAvail = true
@@ -441,8 +444,7 @@ private[spark] class TaskSchedulerImpl(
       }
     }
 
-    val availableSlots = executorIdToCores.values.map(c => c / CPUS_PER_TASK).sum
-    rootPool.updateAvailableSlots(availableSlots)
+    rootPool.updateAvailableSlots(totalSlots)
 
     // Take each TaskSet in our scheduling order, and then offer it each node in increasing order
     // of locality levels so that it gets a chance to launch local tasks on all of them.
@@ -813,7 +815,7 @@ private[spark] class TaskSchedulerImpl(
       // happen below in the rootPool.executorLost() call.
       taskIds.foreach(cleanupTaskState)
     }
-    executorIdToCores.remove(executorId)
+    executorIdToCores.remove(executorId).foreach(totalSlots -= _ / CPUS_PER_TASK)
 
     val host = executorIdToHost(executorId)
     val execs = hostToExecutors.getOrElse(host, new HashSet)
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
index d75b18044d649..04d9c04abeea4 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
@@ -201,10 +201,13 @@ private[spark] class TaskSetManager(
   private[scheduler] var localityWaits = myLocalityLevels.map(getLocalityWait)
 
   // Delay scheduling variables: we keep track of our current locality level and the time we
-  // last launched a task at that level, and move up a level when localityWaits[curLevel] expires.
+  // last launched a task at that level and were also fully utilizing all slots
+  // Move up a level when localityWaits[curLevel] expires.
   // We then move down if we manage to launch a "more local" task.
   private var currentLocalityIndex = 0 // Index of our current locality level in validLocalityLevels
-  private var lastLaunchTime = clock.getTimeMillis()  // Time we last launched a task at this level
+  // Time we last launched a task at this level
+  // and number of running tasks was greater than or equal to numAvailableSlot
+  private var lastFullUtilizationTime = clock.getTimeMillis()
 
   override def schedulableQueue: ConcurrentLinkedQueue[Schedulable] = null
 
@@ -411,12 +414,13 @@ private[spark] class TaskSetManager(
           execId, host, taskLocality, speculative)
         taskInfos(taskId) = info
         taskAttempts(index) = info :: taskAttempts(index)
-        // Update our locality level for delay scheduling
+        // Update our locality level for delay scheduling if all allocated slots are being utilized
         // NO_PREF will not affect the variables related to delay scheduling
+        // Adding 1 to runningTasks, since runningTasks doesn't include this task yet
         val utilizingAllSlots = runningTasks + 1 >= numAvailableSlot
         if (maxLocality != TaskLocality.NO_PREF && utilizingAllSlots) {
           currentLocalityIndex = getLocalityIndex(taskLocality)
-          lastLaunchTime = curTime
+          lastFullUtilizationTime = curTime
         }
         // Serialize and return the task
         val serializedTask: ByteBuffer = try {
@@ -536,14 +540,14 @@ private[spark] class TaskSetManager(
         // This is a performance optimization: if there are no more tasks that can
         // be scheduled at a particular locality level, there is no point in waiting
         // for the locality wait timeout (SPARK-4939).
-        lastLaunchTime = curTime
+        lastFullUtilizationTime = curTime
         logDebug(s"No tasks for locality level ${myLocalityLevels(currentLocalityIndex)}, " +
           s"so moving to locality level ${myLocalityLevels(currentLocalityIndex + 1)}")
         currentLocalityIndex += 1
-      } else if (curTime - lastLaunchTime >= localityWaits(currentLocalityIndex)) {
-        // Jump to the next locality level, and reset lastLaunchTime so that the next locality
-        // wait timer doesn't immediately expire
-        lastLaunchTime += localityWaits(currentLocalityIndex)
+      } else if (curTime - lastFullUtilizationTime >= localityWaits(currentLocalityIndex)) {
+        // Jump to the next locality level, and reset lastFullUtilizationTime
+        // so that the next locality wait timer doesn't immediately expire
+        lastFullUtilizationTime += localityWaits(currentLocalityIndex)
         logDebug(s"Moving to ${myLocalityLevels(currentLocalityIndex + 1)} after waiting for " +
           s"${localityWaits(currentLocalityIndex)}ms")
         currentLocalityIndex += 1
diff --git a/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala
index 80913c3057f71..2b807e82c2bdd 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala
@@ -50,6 +50,7 @@ class PoolSuite extends SparkFunSuite with LocalSparkContext {
     nextTaskSetToSchedule.get.addRunningTask(taskId)
     assert(nextTaskSetToSchedule.get.stageId === expectedStageId)
   }
+
   test("validate FIFO slot distributions") {
     sc = new SparkContext(LOCAL, APP_NAME)
     val taskScheduler = new TaskSchedulerImpl(sc)
diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
index 1b0af8e4b4200..3ebe1478930f7 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
@@ -200,9 +200,6 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext with B
     val clock = new ManualClock
     val conf = new SparkConf()
     sc = new SparkContext("local", "TaskSchedulerImplSuite", conf)
-    // import org.slf4j.{Logger, LoggerFactory}
-    import org.apache.log4j.{Logger, Level}
-    Logger.getLogger("org.apache.spark.scheduler.TaskSetManager").setLevel(Level.DEBUG)
     val taskScheduler = new TaskSchedulerImpl(sc,
       sc.conf.get(config.TASK_MAX_FAILURES),
       clock = clock) {
@@ -249,7 +246,7 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext with B
       .flatten.head
     assert(task2.index === 2)
     taskScheduler.statusUpdate(task1.taskId, TaskState.FINISHED, ByteBuffer.allocate(0))
-    assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host2", 1)))
+    assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
       .flatten.head.index === 3)
     assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec2", "host2", 1)))
       .flatten.isEmpty)

From 168ab306c9df7568a1075ef3314a47dbec26ed95 Mon Sep 17 00:00:00 2001
From: Nicholas Marcott <481161+bmarcott@users.noreply.github.com>
Date: Mon, 16 Dec 2019 22:02:05 -0800
Subject: [PATCH 3/4] only allocate as many slots as remaining tasks add more
 tests update comments

---
 .../org/apache/spark/scheduler/Pool.scala     |  59 ++++++--
 .../spark/scheduler/TaskSchedulerImpl.scala   |   1 +
 .../spark/scheduler/TaskSetManager.scala      |   3 +-
 .../apache/spark/scheduler/PoolSuite.scala    |  80 +++++++++--
 .../scheduler/TaskSchedulerImplSuite.scala    | 133 ++++++++++++++++--
 5 files changed, 239 insertions(+), 37 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/Pool.scala b/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
index 52d8de90a4e36..e25fad3f95805 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
@@ -120,24 +120,61 @@ private[spark] class Pool(
     }
   }
 
+  // Update the number of slots considered available for each TaskSetManager whose ancestor
+  // in the tree is this pool
+  // For FAIR scheduling, slots are distributed among pools based on weights and minshare.
+  //   If a pool requires fewer slots than are available to it, the leftover slots are redistributed
+  //   to the remaining pools using the remaining pools' weights.
+  // For FIFO scheduling, the schedulable queue is iterated over in FIFO order,
+  //   giving each schedulable the remaining slots,
+  //   up to the number of remaining tasks for that schedulable.
   override def updateAvailableSlots(numSlots: Float): Unit = {
     schedulingMode match {
       case SchedulingMode.FAIR =>
-        val usableWeights = schedulableQueue.asScala
-          .map(s => if (s.getSortedTaskSetQueue.nonEmpty) (s, s.weight) else (s, 0))
-        val totalWeights = usableWeights.map(_._2).sum
-        usableWeights.foreach { case (schedulable, usableWeight) =>
-          schedulable.updateAvailableSlots(
-            Math.max(numSlots * usableWeight / totalWeights, schedulable.minShare))
-        }
-      // for FIFO, allocate all slots to the first taskset in the queue
+        val queueCopy = new ConcurrentLinkedQueue[Schedulable](schedulableQueue)
+        var shouldRedistribute = false
+        var totalWeights = schedulableQueue.asScala.map(_.weight).sum
+        var totalSlots = numSlots
+        do {
+          shouldRedistribute = false
+          var nextWeights = totalWeights
+          var nextSlots = totalSlots
+          queueCopy.forEach(schedulable => {
+            val numTasksRemaining = schedulable.getSortedTaskSetQueue
+              .map(tsm => tsm.tasks.length - tsm.tasksSuccessful).sum
+            val allocatedSlots = Math.max(
+              totalSlots * schedulable.weight / totalWeights,
+              schedulable.minShare)
+            if (numTasksRemaining < allocatedSlots) {
+              schedulable.updateAvailableSlots(numTasksRemaining)
+              nextWeights -= schedulable.weight
+              nextSlots -= numTasksRemaining
+              shouldRedistribute = true
+              queueCopy.remove(schedulable)
+            }
+          })
+          totalWeights = nextWeights
+          totalSlots = nextSlots
+        } while (shouldRedistribute)
+
+        // All schedulables remaining have more remaining tasks than their share of slots,
+        // so no need to recalculate remaining tasks. Just give them their total share of slots.
+        queueCopy.forEach(schedulable => {
+          val allocatedSlots = Math.max(
+            totalSlots * schedulable.weight / totalWeights,
+            schedulable.minShare)
+          schedulable.updateAvailableSlots(allocatedSlots)
+        })
       case SchedulingMode.FIFO =>
         val sortedSchedulableQueue =
           schedulableQueue.asScala.toSeq.sortWith(taskSetSchedulingAlgorithm.comparator)
-        var isFirst = true
+        var remainingSlots = numSlots
         for (schedulable <- sortedSchedulableQueue) {
-          schedulable.updateAvailableSlots(if (isFirst) numSlots else 0)
-          isFirst = false
+          val numTasksRemaining = schedulable.getSortedTaskSetQueue
+            .map(tsm => tsm.tasks.length - tsm.tasksSuccessful).sum
+          val toAssign = Math.min(numTasksRemaining, remainingSlots)
+          schedulable.updateAvailableSlots(toAssign)
+          remainingSlots -= toAssign
         }
       case _ =>
         val msg = s"Unsupported scheduling mode: $schedulingMode. Use FAIR or FIFO instead."
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
index 3efc885e3f058..9074500733bec 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
@@ -406,6 +406,7 @@ private[spark] class TaskSchedulerImpl(
       if (!executorIdToRunningTaskIds.contains(o.executorId)) {
         hostToExecutors(o.host) += o.executorId
         executorAdded(o.executorId, o.host)
+        // Assumes the first offer will include all cores (free cores == all cores)
         executorIdToCores(o.executorId) = o.cores
         totalSlots += o.cores / CPUS_PER_TASK
         executorIdToHost(o.executorId) = o.host
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
index 04d9c04abeea4..c46224c6eef09 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
@@ -201,7 +201,8 @@ private[spark] class TaskSetManager(
   private[scheduler] var localityWaits = myLocalityLevels.map(getLocalityWait)
 
   // Delay scheduling variables: we keep track of our current locality level and the time we
-  // last launched a task at that level and were also fully utilizing all slots
+  // last launched a task at that level and were also fully utilizing all "available" slots.
+  // See Pool.updateAvailableSlots() for how the number of available slots is calculated.
   // Move up a level when localityWaits[curLevel] expires.
   // We then move down if we manage to launch a "more local" task.
   private var currentLocalityIndex = 0 // Index of our current locality level in validLocalityLevels
diff --git a/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala
index 2b807e82c2bdd..4693e9d742356 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala
@@ -58,17 +58,36 @@ class PoolSuite extends SparkFunSuite with LocalSparkContext {
     val rootPool = new Pool("", FIFO, 0, 0)
     val schedulableBuilder = new FIFOSchedulableBuilder(rootPool)
 
-    val taskSetManager0 = createTaskSetManager(0, 2, taskScheduler)
-    val taskSetManager1 = createTaskSetManager(1, 2, taskScheduler)
+    val taskSetManager0 = createTaskSetManager(0, 5, taskScheduler)
+    val taskSetManager1 = createTaskSetManager(1, 3, taskScheduler)
     val taskSetManager2 = createTaskSetManager(2, 2, taskScheduler)
     schedulableBuilder.addTaskSetManager(taskSetManager0, null)
     schedulableBuilder.addTaskSetManager(taskSetManager1, null)
     schedulableBuilder.addTaskSetManager(taskSetManager2, null)
 
     rootPool.updateAvailableSlots(9)
-    assert(taskSetManager0.numAvailableSlot === 9)
-    assert(taskSetManager1.numAvailableSlot === 0)
-    assert(taskSetManager2.numAvailableSlot === 0)
+    // first task set only gets 5 slots as there are only 5 tasks
+    assert(taskSetManager0.numAvailableSlot === 5)
+    assert(taskSetManager1.numAvailableSlot === 3)
+    // last task set gets only 1 slot since there are 9 total and 8 have already been taken
+    assert(taskSetManager2.numAvailableSlot === 1)
+
+    rootPool.removeSchedulable(taskSetManager1)
+    rootPool.updateAvailableSlots(6)
+    assert(taskSetManager0.numAvailableSlot === 5)
+    // TSM 2 gets only 1 slot since there are 6 total and 5 have been taken by TSM 0
+    assert(taskSetManager2.numAvailableSlot === 1)
+
+    val taskSetManager3 = createTaskSetManager(3, 3, taskScheduler)
+    schedulableBuilder.addTaskSetManager(taskSetManager3, null)
+    taskSetManager0.tasksSuccessful += 1
+    taskSetManager0.addRunningTask(1)
+    rootPool.updateAvailableSlots(8)
+    // Gets only 4 slots since it has only 4 tasks remaining (1 complete)
+    assert(taskSetManager0.numAvailableSlot === 4)
+    assert(taskSetManager2.numAvailableSlot === 2)
+    // New TSM 3 gets single leftover slot:  8 - 4 (TSM 0) - 2 (TSM 2) = 2
+    assert(taskSetManager3.numAvailableSlot === 2)
   }
 
   test("validate FAIR slot distributions") {
@@ -81,7 +100,9 @@ class PoolSuite extends SparkFunSuite with LocalSparkContext {
     val schedulableBuilder = new FairSchedulableBuilder(rootPool, sc.conf)
     schedulableBuilder.buildPools()
 
-    val taskSetManager0 = createTaskSetManager(0, 2, taskScheduler)
+    val taskSetManager0 = createTaskSetManager(0, 4, taskScheduler)
+    taskSetManager0.tasksSuccessful += 1
+    taskSetManager0.addRunningTask(1)
     val taskSetManager1 = createTaskSetManager(1, 2, taskScheduler)
     sc.setLocalProperty(SparkContext.SPARK_SCHEDULER_POOL, "pool1")
     var localProperties = sc.getLocalProperties
@@ -90,27 +111,60 @@ class PoolSuite extends SparkFunSuite with LocalSparkContext {
 
     sc.setLocalProperty(SparkContext.SPARK_SCHEDULER_POOL, "pool2")
     localProperties = sc.getLocalProperties
-    val taskSetManager2 = createTaskSetManager(2, 2, taskScheduler)
+    val taskSetManager2 = createTaskSetManager(2, 9, taskScheduler)
     val taskSetManager3 = createTaskSetManager(3, 2, taskScheduler)
+    taskSetManager3.tasksSuccessful += 1
     schedulableBuilder.addTaskSetManager(taskSetManager2, localProperties)
     schedulableBuilder.addTaskSetManager(taskSetManager3, localProperties)
 
 
     rootPool.updateAvailableSlots(12)
     // pool 1 weight is 1, pool 2 weight is 2, so they get 4 and 8 slots respectively
-    // since pool 1 is FIFO TSM 0 gets all 4 slots of pool 1
-    // since pool 2 is FAIR TSM 3 and 4 each get half (4) of pool 2's 8 slots
-    assert(taskSetManager0.numAvailableSlot === 4)
-    assert(taskSetManager1.numAvailableSlot === 0)
-    assert(taskSetManager2.numAvailableSlot === 4)
-    assert(taskSetManager3.numAvailableSlot === 4)
+    // TSM 0 has 4 tasks total, but 1 is complete, so it needs only 3 slots.
+    assert(taskSetManager0.numAvailableSlot === 3)
+    // TSM 1 takes the remaining slot from pool 1.
+    assert(taskSetManager1.numAvailableSlot === 1)
+    // Since pool 2 is FAIR TSM 2 and 3 each could get half (4) of pool 2's 8 slots
+    // TSM 3 has only 1 incomplete task, so it takes only 1 out of 4 possible slots
+    assert(taskSetManager3.numAvailableSlot === 1)
+    // TSM 2 takes its original 4 slots + the 3 remaining from TSM 3
+    assert(taskSetManager2.numAvailableSlot === 7)
+
 
     rootPool.updateAvailableSlots(6)
     // similar to above but pool 1 has 3 slots due to having minShare of 3
     assert(taskSetManager0.numAvailableSlot === 3)
     assert(taskSetManager1.numAvailableSlot === 0)
+    assert(taskSetManager2.numAvailableSlot === 3)
+    assert(taskSetManager3.numAvailableSlot === 1)
+
+    taskSetManager3.tasksSuccessful -= 1
+    // same as above, but task 3 should use 2 slots since it has 0 complete tasks
+    rootPool.updateAvailableSlots(6)
+    assert(taskSetManager0.numAvailableSlot === 3)
+    assert(taskSetManager1.numAvailableSlot === 0)
     assert(taskSetManager2.numAvailableSlot === 2)
     assert(taskSetManager3.numAvailableSlot === 2)
+
+    rootPool.getSchedulableByName("pool2").removeSchedulable(taskSetManager2)
+    rootPool.updateAvailableSlots(6)
+    // 4 slots go to pool 2, but TSM 3 only takes 2, remaining 2 can go to pool 1
+    // Pool one has 2 slots + 2 from pool 2 = 4
+    // TSM 0 takes 3 and TSM 1 takes 1
+    assert(taskSetManager0.numAvailableSlot === 3)
+    assert(taskSetManager1.numAvailableSlot === 1)
+    assert(taskSetManager3.numAvailableSlot === 2)
+
+
+    // TSM 4 goes to pool 2
+    val taskSetManager4 = createTaskSetManager(4, 1, taskScheduler)
+    schedulableBuilder.addTaskSetManager(taskSetManager4, localProperties)
+    rootPool.updateAvailableSlots(6)
+    // Same as above but new TSM 4 in pool 2 takes 1 slot, giving only 1 to pool 1
+    assert(taskSetManager0.numAvailableSlot === 3)
+    assert(taskSetManager1.numAvailableSlot === 0)
+    assert(taskSetManager3.numAvailableSlot === 2)
+    assert(taskSetManager4.numAvailableSlot === 1)
   }
 
   test("FIFO Scheduler Test") {
diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
index 3ebe1478930f7..61cd597ed413b 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
@@ -195,7 +195,8 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext with B
     assert(!failedTaskSet)
   }
 
-  test("executors should not sit idle for too long") {
+  test("SPARK-18886 Delay scheduling should not delay some executors indefinitely " +
+    "if one task is scheduled before delay timeout") {
     val LOCALITY_WAIT_MS = 3000
     val clock = new ManualClock
     val conf = new SparkConf()
@@ -218,8 +219,13 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext with B
 
       override def executorAdded(execId: String, host: String): Unit = {}
     }
+    val valueSer = SparkEnv.get.serializer.newInstance()
     taskScheduler.initialize(new FakeSchedulerBackend)
-    val taskSet = FakeTask.createTaskSet(5, 1, 1,
+    val taskSet = FakeTask.createTaskSet(9, 1, 1,
+      Seq(TaskLocation("host1", "exec1")),
+      Seq(TaskLocation("host1", "exec1")),
+      Seq(TaskLocation("host1", "exec1")),
+      Seq(TaskLocation("host1", "exec1")),
       Seq(TaskLocation("host1", "exec1")),
       Seq(TaskLocation("host1", "exec1")),
       Seq(TaskLocation("host1", "exec1")),
@@ -233,23 +239,126 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext with B
     // First offer host2, exec2: no task should be chosen due to bad data locality
     assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec2", "host2", 1)))
       .flatten.isEmpty)
-    val task0 = taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
-      .flatten.head
+
+    // Next offer exec 1, which is local to all tasks. First task will take it
+    val taskDescriptions0 = taskScheduler
+      .resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+      .flatten
+    assert(taskDescriptions0.length === 1)
+    val task0 = taskDescriptions0.head
     assert(task0.index === 0)
-    taskScheduler.statusUpdate(task0.taskId, TaskState.FINISHED, ByteBuffer.allocate(0))
+    val result0 = new DirectTaskResult[Int](valueSer.serialize(0), Seq(), Array())
+    taskScheduler.statusUpdate(task0.taskId, TaskState.FINISHED, valueSer.serialize(result0))
 
+    // clock advances, increasing data locality level to ANY
     clock.advance(LOCALITY_WAIT_MS * 2)
-    val task1 = taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
-      .flatten.head
+    // Local resource is utilized again
+    val taskDescriptions1 = taskScheduler
+      .resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+      .flatten
+    assert(taskDescriptions1.length === 1)
+    val task1 = taskDescriptions1.head
     assert(task1.index === 1)
-    val task2 = taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec2", "host2", 1)))
-      .flatten.head
+
+    // Even though we launched a local task, we still utilize non-local exec2
+    // This is the behavior change to fix SPARK-18886. Also tested below in this same test
+    // This is because we have not utilized our 2 possible slots (exec1 + exec2) within
+    // the locality wait period.
+    val taskDescriptions2 = taskScheduler
+      .resourceOffers(IndexedSeq(WorkerOffer("exec2", "host2", 1)))
+      .flatten
+    assert(taskDescriptions2.length === 1)
+    val task2 = taskDescriptions2.head
     assert(task2.index === 2)
-    taskScheduler.statusUpdate(task1.taskId, TaskState.FINISHED, ByteBuffer.allocate(0))
-    assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
-      .flatten.head.index === 3)
+
+    // Finish data local task on exec 1 (still 1 task running on exec 2)
+    val result1 = new DirectTaskResult[Int](valueSer.serialize(1), Seq(), Array())
+    taskScheduler.statusUpdate(task1.taskId, TaskState.FINISHED, valueSer.serialize(result1))
+
+    // Local resource will again be utilized. Data locality is reset to PROCESS_LOCAL
+    val taskDescriptions3 = taskScheduler
+      .resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+      .flatten
+    assert(taskDescriptions3.length === 1)
+    val task3 = taskDescriptions3.head
+    assert(task3.index === 3)
+    // Complete task 2 and  task 3, no more tasks running
+    val result2 = new DirectTaskResult[Int](valueSer.serialize(2), Seq(), Array())
+    taskScheduler.statusUpdate(task2.taskId, TaskState.FINISHED, valueSer.serialize(result2))
+    val result3 = new DirectTaskResult[Int](valueSer.serialize(3), Seq(), Array())
+    taskScheduler.statusUpdate(task3.taskId, TaskState.FINISHED, valueSer.serialize(result3))
+
+    // Non-local resource will not be utilized, since data locality was reset,
+    // and locality wait has not expired
     assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec2", "host2", 1)))
       .flatten.isEmpty)
+
+    // Offer a total of 2 slots on exec 3, non are taken since non-local
+    assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec3", "host3", 2)))
+      .flatten.isEmpty)
+
+    taskScheduler.executorLost("exec1", LossReasonPending)
+    taskScheduler.executorLost("exec2", LossReasonPending)
+
+    // clock advances, increasing data locality level to ANY
+    clock.advance(LOCALITY_WAIT_MS * 2)
+
+    // Accept non local resource, total utilized slots is 1
+    val taskDescriptions4 = taskScheduler
+      .resourceOffers(IndexedSeq(WorkerOffer("exec3", "host3", 1)))
+      .flatten
+    assert(taskDescriptions4.length === 1)
+    val task4 = taskDescriptions4.head
+    assert(task4.index === 4)
+
+
+    // Offer data local exec 1. This increases total slots to 3, but only 2 are utilized now
+    val taskDescriptions5 = taskScheduler
+      .resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+      .flatten
+    assert(taskDescriptions5.length === 1)
+    val task5 = taskDescriptions5.head
+    assert(task5.index === 5)
+
+    // Data locality has not been reset, still accept non local resource.
+    // Now 3 tasks are running. 3 slots are being used.
+    // Locality timer is reset, but locality level is not
+    val taskDescriptions6 = taskScheduler
+      .resourceOffers(IndexedSeq(WorkerOffer("exec3", "host3", 1)))
+      .flatten
+    assert(taskDescriptions6.length === 1)
+    val task6 = taskDescriptions6.head
+    assert(task6.index === 6)
+
+    // Finish data local task. Only 2 running tasks now
+    val result5 = new DirectTaskResult[Int](valueSer.serialize(5), Seq(), Array())
+    taskScheduler.statusUpdate(task5.taskId, TaskState.FINISHED, valueSer.serialize(result5))
+
+    // Offer data local exec 1 again. 3 slots used and is data local,
+    // so locality is reset to PROCESS_LOCAL and timer is reset
+    val taskDescriptions7 = taskScheduler
+      .resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+      .flatten
+    assert(taskDescriptions7.length === 1)
+    val task7 = taskDescriptions7.head
+    assert(task7.index === 7)
+
+    // Non-local resource will not be utilized, since data locality was reset,
+    // and locality wait has not expired
+    assert(taskScheduler.resourceOffers(IndexedSeq(WorkerOffer("exec3", "host3", 1)))
+      .flatten.isEmpty)
+
+    // Complete task on data local node (exec1)
+    val result7 = new DirectTaskResult[Int](valueSer.serialize(7), Seq(), Array())
+    taskScheduler.statusUpdate(task7.taskId, TaskState.FINISHED, valueSer.serialize(result7))
+
+    // Complete the last task locally
+    val taskDescriptions8 = taskScheduler
+      .resourceOffers(IndexedSeq(WorkerOffer("exec1", "host1", 1)))
+      .flatten
+    assert(taskDescriptions8.length === 1)
+    val task8 = taskDescriptions8.head
+    assert(task8.index === 8)
   }
 
   test("Scheduler does not crash when tasks are not serializable") {

From f9a85f9475d5851f5bdd79e18141950330421445 Mon Sep 17 00:00:00 2001
From: Nicholas Marcott <481161+bmarcott@users.noreply.github.com>
Date: Wed, 18 Dec 2019 21:50:12 -0800
Subject: [PATCH 4/4] eliminate unnecessary threadsafe queue and shortcut on 0
 slots in Pool.updateAvailableSlots

---
 .../org/apache/spark/scheduler/Pool.scala     | 33 +++++++++++--------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/Pool.scala b/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
index e25fad3f95805..6f17c27b08792 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
@@ -17,6 +17,7 @@
 
 package org.apache.spark.scheduler
 
+import java.util
 import java.util.concurrent.{ConcurrentHashMap, ConcurrentLinkedQueue}
 
 import scala.collection.JavaConverters._
@@ -131,15 +132,17 @@ private[spark] class Pool(
   override def updateAvailableSlots(numSlots: Float): Unit = {
     schedulingMode match {
       case SchedulingMode.FAIR =>
-        val queueCopy = new ConcurrentLinkedQueue[Schedulable](schedulableQueue)
-        var shouldRedistribute = false
+        val queueCopy = new util.LinkedList[Schedulable](schedulableQueue)
+        var shouldRedistribute = true
         var totalWeights = schedulableQueue.asScala.map(_.weight).sum
         var totalSlots = numSlots
-        do {
+        while (totalSlots > 0 && shouldRedistribute) {
           shouldRedistribute = false
           var nextWeights = totalWeights
           var nextSlots = totalSlots
-          queueCopy.forEach(schedulable => {
+          val iterator = queueCopy.iterator()
+          while (iterator.hasNext) {
+            val schedulable = iterator.next()
             val numTasksRemaining = schedulable.getSortedTaskSetQueue
               .map(tsm => tsm.tasks.length - tsm.tasksSuccessful).sum
             val allocatedSlots = Math.max(
@@ -150,14 +153,14 @@ private[spark] class Pool(
               nextWeights -= schedulable.weight
               nextSlots -= numTasksRemaining
               shouldRedistribute = true
-              queueCopy.remove(schedulable)
+              iterator.remove()
             }
-          })
+          }
           totalWeights = nextWeights
           totalSlots = nextSlots
-        } while (shouldRedistribute)
+        }
 
-        // All schedulables remaining have more remaining tasks than their share of slots,
+        // All schedulables remaining have more or equal remaining tasks than their share of slots,
         // so no need to recalculate remaining tasks. Just give them their total share of slots.
         queueCopy.forEach(schedulable => {
           val allocatedSlots = Math.max(
@@ -170,11 +173,15 @@ private[spark] class Pool(
           schedulableQueue.asScala.toSeq.sortWith(taskSetSchedulingAlgorithm.comparator)
         var remainingSlots = numSlots
         for (schedulable <- sortedSchedulableQueue) {
-          val numTasksRemaining = schedulable.getSortedTaskSetQueue
-            .map(tsm => tsm.tasks.length - tsm.tasksSuccessful).sum
-          val toAssign = Math.min(numTasksRemaining, remainingSlots)
-          schedulable.updateAvailableSlots(toAssign)
-          remainingSlots -= toAssign
+          if (remainingSlots == 0) {
+            schedulable.updateAvailableSlots(0)
+          } else {
+            val numTasksRemaining = schedulable.getSortedTaskSetQueue
+              .map(tsm => tsm.tasks.length - tsm.tasksSuccessful).sum
+            val toAssign = Math.min(numTasksRemaining, remainingSlots)
+            schedulable.updateAvailableSlots(toAssign)
+            remainingSlots -= toAssign
+          }
         }
       case _ =>
         val msg = s"Unsupported scheduling mode: $schedulingMode. Use FAIR or FIFO instead."