cloud-fan commented on code in PR #58960:
URL: https://github.com/apache/spark/pull/58960#discussion_r4068211600
##########
sql/core/src/main/scala/org/apache/spark/sql/util/PartitionKeyedAccumulator.scala:
##########
@@ -34,41 +34,56 @@ import org.apache.spark.util.AccumulatorV2
* failed/interrupted tasks are dropped by the accumulator framework (it is not
* `countFailedValues`), so only complete per-partition values are ever merged.
*
- * Backed by a `ConcurrentHashMap`. Mutations and folds are synchronized so a
caller can atomically
- * verify that every partition completed and read its statistics from the same
stable snapshot.
+ * Backed by a `ConcurrentHashMap`. All access is synchronized so a caller can
atomically verify
+ * that every partition completed and read its statistics from the same stable
snapshot.
* Framework-facing reads remain safe, weakly consistent views.
*
* @tparam T the per-partition value type. Must be non-null
(`ConcurrentHashMap` forbids nulls).
*/
class PartitionKeyedAccumulator[T] extends AccumulatorV2[(Int, T),
java.util.Map[Int, T]] {
- // partition id -> value.
- private val byPartition = new ConcurrentHashMap[Int, T]()
+ // partition id -> value. Deliberately a lazily created `var` rather than a
`val`: Java
+ // deserialization assigns a subclass's fields only AFTER
`AccumulatorV2.readObject` has already
+ // registered `this` with the `TaskContext`, so another thread can reach
this accumulator while
+ // the map is still unset -- in production the executor heartbeater, which
calls `isZero` on every
+ // registered accumulator. A final field would additionally leave that
thread no guarantee of ever
+ // seeing the post-publication reflective write (JLS 17.5.3). All access
goes through
+ // `getOrCreate`; see SPARK-20977, which fixed the same hazard in
`CollectionAccumulator`.
+ private var byPartition: ConcurrentHashMap[Int, T] = _
+
+ private def getOrCreate: ConcurrentHashMap[Int, T] = {
+ if (byPartition == null) {
+ byPartition = new ConcurrentHashMap[Int, T]()
+ }
+ byPartition
+ }
- override def isZero: Boolean = byPartition.isEmpty
+ override def isZero: Boolean = synchronized {
+ getOrCreate.isEmpty
+ }
override def copyAndReset(): PartitionKeyedAccumulator[T] = new
PartitionKeyedAccumulator[T]
override def copy(): PartitionKeyedAccumulator[T] = synchronized {
val newAcc = new PartitionKeyedAccumulator[T]
- newAcc.byPartition.putAll(byPartition)
+ newAcc.getOrCreate.putAll(getOrCreate)
newAcc
}
override def reset(): Unit = synchronized {
- byPartition.clear()
+ getOrCreate.clear()
}
override def add(v: (Int, T)): Unit = synchronized {
- byPartition.put(v._1, v._2)
+ getOrCreate.put(v._1, v._2)
}
override def merge(other: AccumulatorV2[(Int, T), java.util.Map[Int, T]]):
Unit = synchronized {
other match {
case o: PartitionKeyedAccumulator[T] =>
// Last-write-wins per partition id: a partition recorded by more than
one task replaces
// rather than accumulates, keeping any caller-derived aggregate exact.
- byPartition.putAll(o.byPartition)
+ getOrCreate.putAll(o.value)
Review Comment:
**Non-blocking (P2):** `merge` now holds this accumulator's monitor and
calls `o.value`, which acquires the other accumulator's monitor. If two threads
run `a.merge(b)` and `b.merge(a)`, each can hold one lock and wait on the other
indefinitely. Please obtain the source map before entering the receiver's
synchronized section, or otherwise impose a consistent lock order, so
opposite-direction merges cannot deadlock.
##########
sql/core/src/test/scala/org/apache/spark/sql/util/PartitionKeyedAccumulatorSuite.scala:
##########
@@ -126,4 +130,46 @@ class PartitionKeyedAccumulatorSuite extends SparkFunSuite
{
(rows + partitionRows, bytes + partitionBytes)
}.contains((17L, 170L)))
}
+
+ test("SPARK-57547: accessors are null-safe while readObject publishes the
accumulator") {
+ // `AccumulatorV2.readObject` registers `this` with the `TaskContext`
before Java
+ // deserialization has read this subclass's fields, so the backing map is
still unset at that
+ // point. The executor heartbeater reads every registered accumulator and
calls `isZero` on it,
+ // which used to throw a NullPointerException there and kill the heartbeat
thread. Stand in for
+ // that reader by probing the accumulator from `registerAccumulator`,
which `readObject` calls
+ // at exactly that moment.
+ val acc = new PartitionKeyedAccumulator[Stats]
+ acc.metadata = AccumulatorMetadata(AccumulatorContext.newId(), None,
countFailedValues = false)
+ AccumulatorContext.register(acc)
+
+ var probed = false
+ val taskContext = new TaskContextImpl(
+ stageId = 0,
+ stageAttemptNumber = 0,
+ partitionId = 0,
+ taskAttemptId = 0L,
+ attemptNumber = 0,
+ numPartitions = 1,
+ taskMemoryManager = null,
+ localProperties = new Properties,
+ metricsSystem = null,
+ taskMetrics = TaskMetrics.empty,
+ cpuAmount = BigDecimal(1)) {
+ private[spark] override def registerAccumulator(a: AccumulatorV2[_, _]):
Unit = {
+ // The accessors the heartbeat path reaches. Neither may throw on a
half-read instance.
+ assert(a.isZero)
+ assert(a.value.asInstanceOf[java.util.Map[_, _]].isEmpty)
Review Comment:
**Non-blocking (P2):** `isZero` runs first and initializes `byPartition`, so
this assertion still passes if `value` regresses to directly dereferencing the
null field. Also, the executor heartbeat path calls `isZero`, not `value`,
during this publication window. Please probe each accessor first on a separate
fresh deserialization and describe `value` as an additional null-safety check
rather than part of the heartbeat call path.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
