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

SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git


The following commit(s) were added to refs/heads/main by this push:
     new e3895e6815 [CELEBORN-2508] Fail safely after AQE integrity validation 
errors
e3895e6815 is described below

commit e3895e681515ad702528d2d02e2a232b08deea4e
Author: zhengtao <[email protected]>
AuthorDate: Fri Oct 30 23:00:55 2026 +0800

    [CELEBORN-2508] Fail safely after AQE integrity validation errors
    
    JIRA: https://issues.apache.org/jira/browse/CELEBORN-2508
    
    ### What changes were proposed in this pull request?
    
    Retain a REDUCE partial-read integrity validation failure for the lifetime 
of the logical shuffle. Reject subsequent completion reports and shuffle 
generation requests, and decline conversion of that failure into a Spark fetch 
failure. Clear the failure when the logical shuffle is unregistered.
    
    Limit this terminal state to `PartitionType.REDUCE`. Flink MAP partitions 
reuse the same completion RPC, where finite bounds identify an independently 
validated subpartition range. A failed MAP report must still fail validation, 
while other valid ranges and subsequent corrected reports remain eligible for 
validation.
    
    ### Why are the changes needed?
    
    AQE can split one reducer partition across several tasks. Earlier tasks may 
return corrupted data before the last task reports the aggregate CRC mismatch. 
Spark can then regenerate the shuffle but rerun only the final missing result 
task, retaining the earlier corrupted results. In a controlled Spark 
reproduction, the query returned all 16,000 rows, including one corrupted row, 
after detecting the mismatch and rerunning the stage.
    
    A new physical shuffle does not invalidate those already accepted results. 
Fail the query after a detected partial-read integrity failure until recovery 
can guarantee that every affected result is recomputed. Ordinary fetch failures 
and whole-partition integrity failures retain stage rerun.
    
    ### Does this PR resolve a correctness bug?
    
    - [x] Yes
    
    ### Does this PR introduce _any_ user-facing change?
    
    - [x] Yes
    
    A detected AQE aggregate integrity failure fails the query instead of 
attempting a stage rerun that may preserve corrupted results. No new 
configurations, RPC messages, or worker changes are required. Healthy reads do 
not add RPCs or per-record processing.
    
    ### How was this patch tested?
    
    - Current head: ran the complete `dev/reformat` workflow; all 16 formatter 
invocations succeeded.
    - Current head: 62 client lifecycle/commit tests and 11 Flink reader tests 
passed with the Flink 1.20 profile. The two added regressions exercise real 
`MapPartitionCommitHandler` checksum validation through real `LifecycleManager` 
RPC dispatch for single- and multi-subpartition ranges. Both failed on the 
previous PR head because an unrelated valid range received the retained MAP 
error. With this change, corrupt reports still fail, while another subpartition 
range, another map partitio [...]
    - Previous head `61de621`: compared actual community-client artifacts on 
unmodified Apache Spark 3.5.3 and on Spark 3.5.3 with the Celeborn community 
Spark patch applied and compiled. A real `REBALANCE` workload with production 
`SparkShuffleManager`, 28 mappers, two remote workers, and controlled payload 
corruption reproduced a successful 16,000-row result containing one corrupt row 
after CRC detection and stage rerun in the legacy map-range path on both 
runtimes. The client with this [...]
    - Previous head `61de621`: a combined community build with #3914 also 
passed distributed YARN validation on EMR Spark 3.5.3 with seven executors on 
seven compute nodes and two remote workers. Healthy AQE reads returned all 
expected rows, both partial-read modes rejected injected corruption, ordinary 
whole-partition corruption recovered through stage rerun, and healthy post-read 
task retries completed with correct results. These runtime checks predate the 
MAP scope follow-up; current-h [...]
    
    Closes #3905 from zaynt4606/codex/aqe-safe-failure-community-20261028.
    
    Authored-by: zhengtao <[email protected]>
    Signed-off-by: Nicholas Jiang <[email protected]>
---
 .../apache/celeborn/client/LifecycleManager.scala  |  74 +++++++-
 .../LifecycleManagerIntegrityFailureSuite.scala    | 207 +++++++++++++++++++++
 .../commit/MapPartitionCommitHandlerTest.scala     |  85 ++++++++-
 3 files changed, 358 insertions(+), 8 deletions(-)

diff --git 
a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala 
b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
index 7074899fc1..7122df06a3 100644
--- a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
@@ -111,6 +111,9 @@ class LifecycleManager(val appUniqueId: String, val conf: 
CelebornConf) extends
     Int,
     scala.collection.mutable.LinkedHashMap[String, (Int, Boolean)]]()
   private val celebornShuffleIdToAppShuffleIdMap = 
JavaUtils.newConcurrentHashMap[Int, Int]()
+  // Partial-read integrity failures must survive physical shuffle 
generations. Spark can retain
+  // successful sibling task results when rerunning a stage, including results 
from corrupt data.
+  private val failedIntegrityShuffles = JavaUtils.newConcurrentHashMap[Int, 
String]()
   private val shuffleIdGenerator = new AtomicInteger(0)
   // app shuffle id -> whether shuffle is determinate, rerun of a 
indeterminate shuffle gets different result
   private val appShuffleDeterminateMap = JavaUtils.newConcurrentHashMap[Int, 
Boolean]();
@@ -580,12 +583,22 @@ class LifecycleManager(val appUniqueId: String, val conf: 
CelebornConf) extends
       endMapIndex: Int,
       crc32: Int,
       bytesWritten: Long): Unit = {
-    val (isValid, errorMessage) = commitManager.finishPartition(
-      shuffleId,
-      partitionId,
-      startMapIndex,
-      endMapIndex,
-      new CommitMetadata(crc32, bytesWritten))
+    val (isValid, errorMessage) = integrityFailureForShuffle(shuffleId) match {
+      case Some(error) => (false, error)
+      case None =>
+        val result = commitManager.finishPartition(
+          shuffleId,
+          partitionId,
+          startMapIndex,
+          endMapIndex,
+          new CommitMetadata(crc32, bytesWritten))
+        // MAP readers also use finite ranges, but validate each subpartition 
range independently.
+        if (!result._1 && endMapIndex != Integer.MAX_VALUE &&
+          getPartitionType(shuffleId) == PartitionType.REDUCE) {
+          failedIntegrityShuffles.putIfAbsent(shuffleId, result._2)
+        }
+        result
+    }
     var response: PbReadReducerPartitionEndResponse = null
     if (isValid) {
       response =
@@ -598,6 +611,38 @@ class LifecycleManager(val appUniqueId: String, val conf: 
CelebornConf) extends
     context.reply(response)
   }
 
+  private def integrityFailure(
+      shuffleIds: mutable.LinkedHashMap[String, (Int, Boolean)]): 
Option[String] = {
+    // Healthy applications do not scan shuffle mappings or allocate failure 
state per task.
+    if (failedIntegrityShuffles.isEmpty) {
+      None
+    } else {
+      shuffleIds.synchronized {
+        shuffleIds.valuesIterator
+          .map { case (shuffleId, _) => failedIntegrityShuffles.get(shuffleId) 
}
+          .find(_ != null)
+      }
+    }
+  }
+
+  private def integrityFailureForShuffle(shuffleId: Int): Option[String] = {
+    if (failedIntegrityShuffles.isEmpty) {
+      None
+    } else {
+      Option(failedIntegrityShuffles.get(shuffleId)).orElse {
+        shuffleIdMapping.values().asScala.flatMap { shuffleIds =>
+          shuffleIds.synchronized {
+            if (shuffleIds.valuesIterator.exists(_._1 == shuffleId)) {
+              integrityFailure(shuffleIds)
+            } else {
+              None
+            }
+          }
+        }.take(1).toList.headOption
+      }
+    }
+  }
+
   def setupEndpoints(
       workers: util.Set[WorkerInfo],
       shuffleId: Int,
@@ -1084,6 +1129,13 @@ class LifecycleManager(val appUniqueId: String, val 
conf: CelebornConf) extends
     }
 
     shuffleIds.synchronized {
+      if (integrityFailure(shuffleIds).isDefined) {
+        context.reply(PbGetShuffleIdResponse.newBuilder()
+          .setShuffleId(UNKNOWN_APP_SHUFFLE_ID)
+          .setSuccess(false)
+          .build())
+        return
+      }
       if (isWriter) {
         shuffleIds.get(appShuffleIdentifier) match {
           case Some((shuffleId, _)) =>
@@ -1168,6 +1220,14 @@ class LifecycleManager(val appUniqueId: String, val 
conf: CelebornConf) extends
     }
     var ret = true
     shuffleIds.synchronized {
+      if (integrityFailure(shuffleIds).isDefined) {
+        // A FetchFailed would allow Spark to keep results from corrupt 
sibling tasks.
+        // Return false so the reader surfaces the validation error as a task 
failure.
+        context.reply(PbReportShuffleFetchFailureResponse.newBuilder()
+          .setSuccess(false)
+          .build())
+        return
+      }
       shuffleIds.find(e => e._2._1 == shuffleId) match {
         case Some((appShuffleIdentifier, (shuffleId, true))) =>
           if (invokeReportTaskShuffleFetchFailurePreCheck(taskId)) {
@@ -1354,11 +1414,13 @@ class LifecycleManager(val appUniqueId: String, val 
conf: CelebornConf) extends
           shuffleIds.values.map {
             case (shuffleId, _) =>
               unregisterShuffle(shuffleId)
+              failedIntegrityShuffles.remove(shuffleId)
               unregisterShuffleCallback.foreach(c => c.accept(shuffleId))
           })
       }
     } else {
       unregisterShuffle(appShuffleId)
+      failedIntegrityShuffles.remove(appShuffleId)
     }
   }
 
diff --git 
a/client/src/test/scala/org/apache/celeborn/client/LifecycleManagerIntegrityFailureSuite.scala
 
b/client/src/test/scala/org/apache/celeborn/client/LifecycleManagerIntegrityFailureSuite.scala
new file mode 100644
index 0000000000..4d92c35f64
--- /dev/null
+++ 
b/client/src/test/scala/org/apache/celeborn/client/LifecycleManagerIntegrityFailureSuite.scala
@@ -0,0 +1,207 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.client
+
+import java.util.concurrent.atomic.AtomicInteger
+import java.util.function.Consumer
+
+import scala.collection.mutable
+
+import org.mockito.ArgumentCaptor
+import org.mockito.ArgumentMatchers.{any, anyInt}
+import org.mockito.Mockito._
+
+import org.apache.celeborn.CelebornFunSuite
+import org.apache.celeborn.common.{CelebornConf, CommitMetadata}
+import org.apache.celeborn.common.protocol._
+import 
org.apache.celeborn.common.protocol.message.ControlMessages.ReadReducerPartitionEnd
+import org.apache.celeborn.common.protocol.message.StatusCode
+import org.apache.celeborn.common.rpc.RpcCallContext
+import org.apache.celeborn.common.util.JavaUtils
+
+class LifecycleManagerIntegrityFailureSuite extends CelebornFunSuite {
+  private val validationError = "AQE Partition 0 failed validation check"
+
+  // Exercise real RPC dispatch and lifecycle decisions without starting a 
master or RPC threads.
+  private def fixture(): (LifecycleManager, CommitManager, AtomicInteger) = {
+    val manager = mock(classOf[LifecycleManager], CALLS_REAL_METHODS)
+    val commit = mock(classOf[CommitManager])
+    setField(manager, "conf", new CelebornConf())
+    setField(
+      manager,
+      "shuffleIdMapping",
+      JavaUtils.newConcurrentHashMap[
+        Int,
+        mutable.LinkedHashMap[String, (Int, Boolean)]]())
+    setField(manager, "failedIntegrityShuffles", 
JavaUtils.newConcurrentHashMap[Int, String]())
+    setField(manager, "shuffleIdGenerator", new AtomicInteger())
+    setField(manager, "appShuffleDeterminateMap", 
JavaUtils.newConcurrentHashMap[Int, Boolean]())
+    setField(
+      manager,
+      "celebornShuffleIdToAppShuffleIdMap",
+      JavaUtils.newConcurrentHashMap[Int, Int]())
+    setField(manager, "unregisterShuffleCallback", None)
+    setField(manager, "validateCelebornShuffleIdForClean", None)
+    manager.registerReportTaskShuffleFetchFailurePreCheck(new 
java.util.function.Function[
+      java.lang.Long,
+      Boolean] {
+      override def apply(id: java.lang.Long): Boolean = true
+    })
+    setField(manager, "appShuffleTrackerCallback", None)
+    setField(manager, "commitManager", commit)
+    setField(manager, "shufflePartitionType", 
JavaUtils.newConcurrentHashMap[Int, PartitionType]())
+    when(manager.getPartitionType(anyInt())).thenReturn(PartitionType.REDUCE)
+    doNothing().when(manager).unregisterShuffle(anyInt())
+    when(commit.finishPartition(anyInt(), anyInt(), anyInt(), anyInt(), 
any[CommitMetadata]))
+      .thenReturn((true, ""))
+    when(commit.areAllMapperAttemptsFinished(anyInt())).thenReturn(true)
+    val callbacks = new AtomicInteger()
+    manager.registerShuffleTrackerCallback(new Consumer[Integer] {
+      override def accept(id: Integer): Unit = { callbacks.incrementAndGet() }
+    })
+    manager.registerAppShuffleDeterminate(10, determinate = false)
+    (manager, commit, callbacks)
+  }
+
+  private def setField(manager: LifecycleManager, name: String, value: 
AnyRef): Unit = {
+    val field = classOf[LifecycleManager].getDeclaredFields
+      .find(f => f.getName == name || f.getName.endsWith("$$" + name)).get
+    field.setAccessible(true)
+    field.set(manager, value)
+  }
+
+  private def reply(manager: LifecycleManager, request: Any): AnyRef = {
+    val context = mock(classOf[RpcCallContext])
+    manager.receiveAndReply(context)(request)
+    val response = ArgumentCaptor.forClass(classOf[Object])
+    verify(context).reply(response.capture())
+    response.getValue
+  }
+
+  private def shuffleId(
+      manager: LifecycleManager,
+      attempt: Int,
+      appId: Int = 10): PbGetShuffleIdResponse = {
+    reply(
+      manager,
+      PbGetShuffleId.newBuilder()
+        .setAppShuffleId(appId)
+        .setAppShuffleIdentifier(s"$appId-0-$attempt")
+        .setIsShuffleWriter(true)
+        .build()).asInstanceOf[PbGetShuffleIdResponse]
+  }
+
+  private def fetchFailure(manager: LifecycleManager, id: Int, appId: Int = 
10): Boolean = {
+    reply(
+      manager,
+      PbReportShuffleFetchFailure.newBuilder()
+        .setAppShuffleId(appId).setShuffleId(id).build())
+      .asInstanceOf[PbReportShuffleFetchFailureResponse].getSuccess
+  }
+
+  private def finish(
+      manager: LifecycleManager,
+      id: Int,
+      end: Int): PbReadReducerPartitionEndResponse = {
+    reply(manager, ReadReducerPartitionEnd(id, 0, if (end == 0) 2 else 0, end, 
1, 100L))
+      .asInstanceOf[PbReadReducerPartitionEndResponse]
+  }
+
+  Seq(1, 0).foreach { end =>
+    test(
+      s"partial integrity failure cannot trigger stage rerun or accept 
corrected reports, end=$end") {
+      val (manager, commit, callbacks) = fixture()
+      val id = shuffleId(manager, 0).getShuffleId
+      when(commit.finishPartition(anyInt(), anyInt(), anyInt(), anyInt(), 
any[CommitMetadata]))
+        .thenReturn((false, validationError))
+      assert(finish(
+        manager,
+        id,
+        end).getStatus == 
StatusCode.READ_REDUCER_PARTITION_END_FAILED.getValue)
+      assert(!fetchFailure(manager, id))
+      assert(callbacks.get() == 0)
+      assert(!shuffleId(manager, 1).getSuccess)
+      when(commit.finishPartition(anyInt(), anyInt(), anyInt(), anyInt(), 
any[CommitMetadata]))
+        .thenReturn((true, ""))
+      assert(finish(manager, id, end).getErrorMsg == validationError)
+      verify(commit, times(1)).finishPartition(
+        anyInt(),
+        anyInt(),
+        anyInt(),
+        anyInt(),
+        any[CommitMetadata])
+    }
+  }
+
+  test("a physical generation allocated before the integrity failure cannot 
bypass it") {
+    val (manager, commit, _) = fixture()
+    val first = shuffleId(manager, 0).getShuffleId
+    assert(fetchFailure(manager, first))
+    val second = shuffleId(manager, 1).getShuffleId
+    assert(second != first)
+    when(commit.finishPartition(anyInt(), anyInt(), anyInt(), anyInt(), 
any[CommitMetadata]))
+      .thenReturn((false, validationError))
+    finish(manager, first, 1)
+    when(commit.finishPartition(anyInt(), anyInt(), anyInt(), anyInt(), 
any[CommitMetadata]))
+      .thenReturn((true, ""))
+    assert(finish(manager, second, 1).getErrorMsg == validationError)
+    assert(!fetchFailure(manager, second))
+    assert(!shuffleId(manager, 2).getSuccess)
+  }
+
+  test("ordinary fetch failure and full-partition integrity failure retain 
stage rerun") {
+    val (manager, commit, callbacks) = fixture()
+    val first = shuffleId(manager, 0).getShuffleId
+    assert(fetchFailure(manager, first))
+    val second = shuffleId(manager, 1).getShuffleId
+    assert(second != first)
+    when(commit.finishPartition(anyInt(), anyInt(), anyInt(), anyInt(), 
any[CommitMetadata]))
+      .thenReturn((false, "full-partition mismatch"))
+    finish(manager, second, Integer.MAX_VALUE)
+    assert(fetchFailure(manager, second))
+    assert(shuffleId(manager, 2).getSuccess)
+    assert(callbacks.get() == 2)
+  }
+
+  test("integrity failure does not block unrelated logical shuffles") {
+    val (manager, commit, _) = fixture()
+    val bad = shuffleId(manager, 0).getShuffleId
+    when(commit.finishPartition(anyInt(), anyInt(), anyInt(), anyInt(), 
any[CommitMetadata]))
+      .thenReturn((false, validationError))
+    finish(manager, bad, 1)
+    val unrelated = shuffleId(manager, 0, appId = 20)
+    assert(unrelated.getSuccess)
+    when(commit.finishPartition(anyInt(), anyInt(), anyInt(), anyInt(), 
any[CommitMetadata]))
+      .thenReturn((true, ""))
+    assert(finish(manager, unrelated.getShuffleId, 1).getStatus == 
StatusCode.SUCCESS.getValue)
+    assert(fetchFailure(manager, unrelated.getShuffleId, appId = 20))
+  }
+
+  test("logical shuffle unregister removes its integrity failure") {
+    val (manager, commit, _) = fixture()
+    val first = shuffleId(manager, 0).getShuffleId
+    when(commit.finishPartition(anyInt(), anyInt(), anyInt(), anyInt(), 
any[CommitMetadata]))
+      .thenReturn((false, validationError))
+    finish(manager, first, 1)
+    manager.unregisterAppShuffle(10, hasMapping = true)
+    assert(shuffleId(manager, 1).getSuccess)
+    when(commit.finishPartition(anyInt(), anyInt(), anyInt(), anyInt(), 
any[CommitMetadata]))
+      .thenReturn((true, ""))
+    assert(finish(manager, first, 1).getStatus == StatusCode.SUCCESS.getValue)
+  }
+}
diff --git 
a/client/src/test/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandlerTest.scala
 
b/client/src/test/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandlerTest.scala
index 4e1f52d1c6..7ea6158295 100644
--- 
a/client/src/test/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandlerTest.scala
+++ 
b/client/src/test/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandlerTest.scala
@@ -19,12 +19,23 @@ package org.apache.celeborn.client.commit
 
 import java.util.concurrent.{ConcurrentHashMap, ConcurrentLinkedQueue, 
CountDownLatch, ScheduledExecutorService, ThreadPoolExecutor, TimeUnit}
 
+import scala.collection.mutable
+
+import org.mockito.ArgumentCaptor
+import org.mockito.ArgumentMatchers.{any, anyInt}
+import org.mockito.Mockito._
+import org.mockito.invocation.InvocationOnMock
+
 import org.apache.celeborn.CelebornFunSuite
-import org.apache.celeborn.client.{ShuffleCommittedInfo, WorkerStatusTracker}
+import org.apache.celeborn.client.{CommitManager, LifecycleManager, 
ShuffleCommittedInfo, WorkerStatusTracker}
 import org.apache.celeborn.client.CommitManager.CommittedPartitionInfo
 import org.apache.celeborn.client.LifecycleManager.ShuffleAllocatedWorkers
 import org.apache.celeborn.common.{CelebornConf, CommitMetadata}
-import org.apache.celeborn.common.util.ThreadUtils
+import org.apache.celeborn.common.protocol.{PartitionType, 
PbReadReducerPartitionEndResponse}
+import 
org.apache.celeborn.common.protocol.message.ControlMessages.ReadReducerPartitionEnd
+import org.apache.celeborn.common.protocol.message.StatusCode
+import org.apache.celeborn.common.rpc.RpcCallContext
+import org.apache.celeborn.common.util.{JavaUtils, ThreadUtils}
 
 class MapPartitionCommitHandlerTest extends CelebornFunSuite {
 
@@ -297,4 +308,74 @@ class MapPartitionCommitHandlerTest extends 
CelebornFunSuite {
       assert(isValid, s"map partition $mapPartition failed validation: 
$message")
     }
   }
+
+  Seq((0, 0), (0, 1)).foreach { case (startSubIndex, endSubIndex) =>
+    test(s"MAP integrity failure preserves independent reads and retry for 
range [$startSubIndex, $endSubIndex]") {
+      val handler = newHandler(integrityCheckEnabled = true)
+      val (checksums, bytesWritten) = writeSideMetadata(payloads)
+      Seq(mapPartitionId, mapPartitionId + 1).foreach { id =>
+        handler.recordMapPartitionCommitMetadata(
+          shuffleId,
+          id,
+          numPartitions,
+          checksums,
+          bytesWritten)
+      }
+
+      val manager = mock(classOf[LifecycleManager], CALLS_REAL_METHODS)
+      val commit = mock(classOf[CommitManager])
+      when(commit.finishPartition(anyInt(), anyInt(), anyInt(), anyInt(), 
any[CommitMetadata]))
+        .thenAnswer((invocation: InvocationOnMock) =>
+          handler.finishPartition(
+            invocation.getArgument[Int](0),
+            invocation.getArgument[Int](1),
+            invocation.getArgument[Int](2),
+            invocation.getArgument[Int](3),
+            invocation.getArgument[CommitMetadata](4)))
+      val partitionTypes = JavaUtils.newConcurrentHashMap[Int, PartitionType]()
+      partitionTypes.put(shuffleId, PartitionType.MAP)
+      val fields = Map[String, AnyRef](
+        "commitManager" -> commit,
+        "shufflePartitionType" -> partitionTypes,
+        "failedIntegrityShuffles" -> JavaUtils.newConcurrentHashMap[Int, 
String](),
+        "shuffleIdMapping" -> JavaUtils.newConcurrentHashMap[
+          Int,
+          mutable.LinkedHashMap[String, (Int, Boolean)]]())
+      fields.foreach { case (name, value) =>
+        val field = classOf[LifecycleManager].getDeclaredFields
+          .find(f => f.getName == name || f.getName.endsWith("$$" + name)).get
+        field.setAccessible(true)
+        field.set(manager, value)
+      }
+
+      def report(
+          partitionId: Int,
+          start: Int,
+          end: Int,
+          metadata: CommitMetadata): PbReadReducerPartitionEndResponse = {
+        val context = mock(classOf[RpcCallContext])
+        manager.receiveAndReply(context)(ReadReducerPartitionEnd(
+          shuffleId,
+          partitionId,
+          start,
+          end,
+          metadata.getChecksum,
+          metadata.getBytes))
+        val response = ArgumentCaptor.forClass(classOf[Object])
+        verify(context).reply(response.capture())
+        response.getValue.asInstanceOf[PbReadReducerPartitionEndResponse]
+      }
+
+      val valid = readSideMetadata(payloads, startSubIndex, endSubIndex)
+      val corrupt = new CommitMetadata(valid.getChecksum ^ 1, valid.getBytes)
+      assert(report(mapPartitionId, startSubIndex, endSubIndex, corrupt)
+        .getStatus == StatusCode.READ_REDUCER_PARTITION_END_FAILED.getValue)
+      val otherRange = readSideMetadata(payloads, 2, 2)
+      assert(report(mapPartitionId, 2, 2, otherRange).getStatus == 
StatusCode.SUCCESS.getValue)
+      assert(report(mapPartitionId + 1, startSubIndex, endSubIndex, valid)
+        .getStatus == StatusCode.SUCCESS.getValue)
+      assert(report(mapPartitionId, startSubIndex, endSubIndex, valid)
+        .getStatus == StatusCode.SUCCESS.getValue)
+    }
+  }
 }

Reply via email to