(spark) branch branch-4.x updated: [SPARK-59845][CORE] Register an accumulator with the TaskContext only after it is fully deserialized

[we...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.x by this push:
     new ba033ef70ab7 [SPARK-59845][CORE] Register an accumulator with the TaskContext only after it is fully deserialized
ba033ef70ab7 is described below

commit ba033ef70ab71a5607dd508e5d241afc5d72a54b
Author: Ziqi Liu <zi...@databricks.com>
AuthorDate: Tue Sep 29 08:49:18 2026 +0000

    [SPARK-59845][CORE] Register an accumulator with the TaskContext only after it is fully deserialized
    
    ### What changes were proposed in this pull request?
    
    Move the `TaskContext` registration out of `readObject` and into `readResolve`.
    
    `AccumulatorV2` currently registers from `readObject`, which runs at AccumulatorV2's own slot. Java serialization reads hierarchy data **superclass-first**, so no subclass field is assigned yet — while `registerAccumulator` publishes `this` into `TaskMetrics` where other threads can read it. The PR drops `readObject` (beyond registering, all it did was in.defaultReadObject(), which the framework does anyway) and registers from `final protected def readResolve()` instead, which runs on [...]
    
    It's `final` for the same reason `writeReplace` already is: `readObject` must be `private` so it's invoked per class, whereas `readResolve` is inherited — a subclass defining one would silently stop registration.
    
    ### Why are the changes needed?
    
    Subclasses are published before their state exists, so a subclass holding state in a `val` exposes a `null`. The reader that hits it is the executor heartbeater, calling `isZero` on every registered accumulator. The NPE propagates out of a `scheduleAtFixedRate` task, so `ScheduledThreadPoolExecutor` stops rescheduling — the executor never heartbeats again, the driver removes it on timeout.
    
    Plus the JLS 17.5.3 aspect: a val is a final field written reflectively after publication, so a reader has no guarantee of seeing it and the JIT may fold the read.
    
    SPARK-20977 fixed this locally in CollectionAccumulator; every subclass since has re-derived the workaround and not all have — SetAccumulator in sql/execution/debug still dereferences private val _set from isZero, and EventTimeStatsAccum escapes the NPE only incidentally via Scala's null-safe ==, reporting the wrong answer instead.
    
    ### Does this PR introduce _any_ user-facing change?
    NO
    
    ### How was this patch tested?
    New test case
    
    ### Was this patch authored or co-authored using generative AI tooling?
    Yes. Claude Opus 5.0
    
    Closes #59124 from liuzqt/accumulator-register-after-deserialization.
    
    Authored-by: Ziqi Liu <zi...@databricks.com>
    Signed-off-by: Wenchen Fan <we...@databricks.com>
    (cherry picked from commit 5b7d3fe1d74733926ad60c108cd3f2f502e85498)
    Signed-off-by: Wenchen Fan <we...@databricks.com>
---
 .../org/apache/spark/util/AccumulatorV2.scala      | 15 ++++--
 .../org/apache/spark/util/AccumulatorV2Suite.scala | 62 ++++++++++++++++++++++
 2 files changed, 73 insertions(+), 4 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/util/AccumulatorV2.scala b/core/src/main/scala/org/apache/spark/util/AccumulatorV2.scala
index 7e2047526da0..20b2baef0f23 100644
--- a/core/src/main/scala/org/apache/spark/util/AccumulatorV2.scala
+++ b/core/src/main/scala/org/apache/spark/util/AccumulatorV2.scala
@@ -18,7 +18,6 @@
 package org.apache.spark.util
 
 import java.{lang => jl}
-import java.io.ObjectInputStream
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.atomic.AtomicLong
 
@@ -197,9 +196,16 @@ abstract class AccumulatorV2[IN, OUT] extends Serializable {
     }
   }
 
-  // Called by Java when deserializing an object
-  private def readObject(in: ObjectInputStream): Unit = Utils.tryOrIOException {
-    in.defaultReadObject()
+  // Called by Java once this object is fully deserialized -- after every class in the hierarchy has
+  // had its fields read, including a subclass's. This must not move to `readObject`, which runs at
+  // this class's slot while a subclass's fields are still unset: `registerAccumulator` publishes
+  // `this` where other threads can reach it, and the executor heartbeater calls `isZero` on every
+  // registered accumulator. Each subclass would then have to tolerate its own state being null,
+  // which is the hazard SPARK-20977 worked around in `CollectionAccumulator`.
+  // `final` for the same reason as `writeReplace` above: unlike `readObject`, which is private and
+  // therefore invoked once per class in the hierarchy, `readResolve` is inherited, so a subclass
+  // defining one would silently replace this and stop registering the accumulator entirely.
+  final protected def readResolve(): Any = Utils.tryOrIOException {
     if (atDriverSide) {
       atDriverSide = false
 
@@ -213,6 +219,7 @@ abstract class AccumulatorV2[IN, OUT] extends Serializable {
     } else {
       atDriverSide = true
     }
+    this
   }
 
   override def toString: String = {
diff --git a/core/src/test/scala/org/apache/spark/util/AccumulatorV2Suite.scala b/core/src/test/scala/org/apache/spark/util/AccumulatorV2Suite.scala
index 172bebbfec61..22e518da5b2c 100644
--- a/core/src/test/scala/org/apache/spark/util/AccumulatorV2Suite.scala
+++ b/core/src/test/scala/org/apache/spark/util/AccumulatorV2Suite.scala
@@ -17,7 +17,10 @@
 
 package org.apache.spark.util
 
+import java.util.Properties
+
 import org.apache.spark._
+import org.apache.spark.executor.TaskMetrics
 
 class AccumulatorV2Suite extends SparkFunSuite {
 
@@ -127,6 +130,65 @@ class AccumulatorV2Suite extends SparkFunSuite {
     assert(acc3.value.isEmpty)
   }
 
+  test("SPARK-59845: a subclass is registered only once it is fully deserialized") {
+    // Deserialization registers the accumulator with the `TaskContext`, which publishes it to other
+    // threads -- in production the executor heartbeater, which calls `isZero` on every registered
+    // accumulator. Registering before the subclass's fields have been read exposes state that is
+    // still unset, so a subclass holding its state in a plain `val` would see a null there.
+    // Probe from `registerAccumulator` to observe that instant without racing a real reader.
+    val acc = new NaiveStateAccumulator
+    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 = {
+        assert(a.isZero, "accumulator state should be readable as soon as it is registered")
+        probed = true
+      }
+    }
+
+    TaskContext.setTaskContext(taskContext)
+    try {
+      Utils.deserialize[NaiveStateAccumulator](Utils.serialize(acc))
+    } finally {
+      TaskContext.unset()
+      AccumulatorContext.remove(acc.id)
+    }
+    assert(probed, "the accumulator was never registered, so nothing was verified")
+  }
 }
 
 class MyData(val i: Int) extends Serializable
+
+/**
+ * Holds its state the way a subclass naturally would, in a `val` dereferenced by `isZero`, with no
+ * null handling. Registration must therefore happen only after this field has been deserialized.
+ */
+private class NaiveStateAccumulator extends AccumulatorV2[Int, java.util.List[Int]] {
+  private val state = new java.util.ArrayList[Int]()
+
+  override def isZero: Boolean = state.isEmpty
+  override def copyAndReset(): NaiveStateAccumulator = new NaiveStateAccumulator
+  override def copy(): NaiveStateAccumulator = {
+    val newAcc = new NaiveStateAccumulator
+    newAcc.state.addAll(state)
+    newAcc
+  }
+  override def reset(): Unit = state.clear()
+  override def add(v: Int): Unit = state.add(v)
+  override def merge(other: AccumulatorV2[Int, java.util.List[Int]]): Unit =
+    state.addAll(other.value)
+  override def value: java.util.List[Int] = state
+}


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org