Trying to prove davies suggestion in https://github.com/apache/spark/pull/10240

spark-pr-10240-suggestion.md

This page tries to prove that the following two are equivalent, as suggested by @davies.

// === (1): The original code in Spark 1.6 before PR 10240

val maxToGrant = math.min(numBytes, math.max(0, maxMemoryPerTask - curMem))
val toGrant = math.min(maxToGrant, memoryFree)

if (curMem < minMemoryPerTask) {
  if (memoryFree >= math.min(maxToGrant, minMemoryPerTask - curMem)) {
    memoryForTask(taskAttemptId) += toGrant
    return toGrant
  } else {
    lock.wait()
  }
} else {
  memoryForTask(taskAttemptId) += toGrant
  return toGrant
}

// === (2) New code suggested by @davies, which he claims is equivalent

val maxToGrant = ...
val toGrant = ...

if (toGrant < numBytes && curMem + toGrant < minMemoryPerTask) {
  lock.wait()
} else {
  memoryForTask(taskAttemptId) += toGrant
  return toGrant
}

Now, I will slowly transform (1) to (2) through a series of steps, assuming all variables are >= 0.

Step 1

Rewrite the nested if-else's into a single if-else:

if (curMem < minMemoryPerTask && memoryFree < math.min(maxToGrant, minMemoryPerTask - curMem)) {
  lock.wait()
} else {
  memoryForTask(taskAttemptId) += toGrant
  return toGrant
}

Now we're left with the same if-else structure as in (2), so all we need to do is to prove that the following two clauses are equal.

// === (1)
curMem < minMemoryPerTask && memoryFree < math.min(maxToGrant, minMemoryPerTask - curMem)

// === (2)
toGrant < numBytes && curMem + toGrant < minMemoryPerTask

Step 2

Expand the math.min to end up with 3 AND expressions.

// === (1)
curMem < minMemoryPerTask &&
  memoryFree < maxToGrant &&
  memoryFree < minMemoryPerTask - curMem

Step 3

Notice that memoryFree < maxToGrant implies toGrant == memoryFree, since we have

val toGrant = math.min(maxToGrant, memoryFree)

so we can replace the memoryFree in the third clause with toGrant and rearrange things in the inequality.

// === (1)
curMem < minMemoryPerTask &&
  memoryFree < maxToGrant &&
  curMem + toGrant < minMemoryPerTask

Step 4

The third clause implies the first clause, so remove the first clause.

// === (1)
memoryFree < maxToGrant && curMem + toGrant < minMemoryPerTask

// === (2)
toGrant < numBytes && curMem + toGrant < minMemoryPerTask

Now (1) and (2) are starting to look really similar. The only difference is in the first clause of each predicate.

Step 5

Notice that memoryFree < maxToGrant is the same as toGrant < maxToGrant, since we have

val maxToGrant = math.min(numBytes, math.max(0, maxMemoryPerTask - curMem))
val toGrant = math.min(maxToGrant, memoryFree)

Further notice that maxToGrant <= numBytes. In other words, toGrant < maxToGrant <= numBytes. Finally, we end up with

// === (1)
toGrant < numBytes && curMem + toGrant < minMemoryPerTask

// === (2)
toGrant < numBytes && curMem + toGrant < minMemoryPerTask

They are the same provided I did not make any mistakes!