[CELEBORN-2460] Reply SHUFFLE_EXPIRED for a released shuffle instead of an empty file group

client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java

@@ -1936,6 +1936,7 @@ protected Tuple3<ReduceFileGroups, String, Exception> loadFileGroupInternal(
               null);
         case STAGE_END_TIMEOUT:
         case SHUFFLE_DATA_LOST:
+        case SHUFFLE_EXPIRED:
           exceptionMsg =
               String.format(
                   "Request %s return %s for %s.",

client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala

@@ -90,6 +90,7 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
   private val shufflePartitionType = JavaUtils.newConcurrentHashMap[Int, PartitionType]()
   private val rangeReadFilter = conf.shuffleRangeReadFilterEnabled
   private val unregisterShuffleTime = JavaUtils.newConcurrentHashMap[Int, Long]()
+  private val expiredShuffleIds = ConcurrentHashMap.newKeySet[Int]()
 
   val registeredShuffle = ConcurrentHashMap.newKeySet[Int]()
   val shuffleCount = new LongAdder()
@@ -971,6 +972,15 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
     // If isSegmentGranularityVisible is set to true, the downstream reduce task may start early than upstream map task, e.g. flink hybrid shuffle.
     // Under these circumstances, there's a possibility that the shuffle might not yet be registered when the downstream reduce task send GetReduceFileGroup request,
     // so we shouldn't send a SHUFFLE_NOT_REGISTERED response directly, should enqueue this request to pending list, and response to the downstream reduce task the ReduceFileGroup when the upstream map task register shuffle done
+    if (expiredShuffleIds.contains(shuffleId)) {
+      logWarning(s"[handleGetReducerFileGroup] shuffle $shuffleId has been released, its data must be recomputed.")
+      context.reply(GetReducerFileGroupResponse(
+        StatusCode.SHUFFLE_EXPIRED,
+        JavaUtils.newConcurrentHashMap(),
+        Array.empty,
+        serdeVersion = serdeVersion))
+      return
+    }
     if (!registeredShuffle.contains(shuffleId) && !isSegmentGranularityVisible) {
       logWarning(s"[handleGetReducerFileGroup] shuffle $shuffleId not registered, maybe no shuffle data within this stage.")
       context.reply(GetReducerFileGroupResponse(
@@ -1279,6 +1289,12 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
         }
       }
     }
+    if (celebornShuffleIdToAppShuffleIdMap.containsKey(shuffleId)) {
+      val appShuffleId = celebornShuffleIdToAppShuffleIdMap.get(shuffleId)
+      if (shuffleIdMapping.containsKey(appShuffleId) && registeredShuffle.contains(shuffleId)) {
+        expiredShuffleIds.add(shuffleId)
+      }
+    }
     celebornShuffleIdToAppShuffleIdMap.remove(shuffleId)
     // add shuffleKey to delay shuffle removal set
     unregisterShuffleTime.put(shuffleId, System.currentTimeMillis())

common/src/main/java/org/apache/celeborn/common/protocol/message/StatusCode.java

@@ -93,7 +93,8 @@ public enum StatusCode {
   SEGMENT_START_FAIL_PRIMARY(53),
   NO_SPLIT(54),
   WORKER_UNRESPONSIVE(55),
-  READ_REDUCER_PARTITION_END_FAILED(56);
+  READ_REDUCER_PARTITION_END_FAILED(56),
+  SHUFFLE_EXPIRED(57);
 
   private final byte value;
 

tests/spark-it/src/test/scala/org/apache/celeborn/tests/client/LifecycleManagerUnregisterShuffleSuite.scala

@@ -27,6 +27,9 @@ import org.scalatest.time.SpanSugar.convertIntToGrainOfTime
 
 import org.apache.celeborn.client.{LifecycleManager, WithShuffleClientSuite}
 import org.apache.celeborn.common.CelebornConf
+import org.apache.celeborn.common.network.protocol.SerdeVersion
+import org.apache.celeborn.common.protocol.{PbGetShuffleId, PbGetShuffleIdResponse}
+import org.apache.celeborn.common.protocol.message.ControlMessages.{GetReducerFileGroup, GetReducerFileGroupResponse}
 import org.apache.celeborn.common.protocol.message.StatusCode
 import org.apache.celeborn.service.deploy.MiniClusterFeature
 
@@ -122,6 +125,50 @@ class LifecycleManagerUnregisterShuffleSuite extends WithShuffleClientSuite
     lifecycleManager.stop()
   }
 
+  test(
+    "released shuffle is reported as expired and an empty map stage is still reported as empty") {
+    val conf = celebornConf.clone
+    val lifecycleManager: LifecycleManager = new LifecycleManager(APP, conf)
+    val ids =
+      new util.ArrayList[Integer]((0 until 10).toList.map(x => Integer.valueOf(x)).asJava)
+    val appShuffleId = 101
+    val neverMappedShuffleId = 202
+
+    def getReducerFileGroup(shuffleId: Int): GetReducerFileGroupResponse =
+      lifecycleManager.self.askSync[GetReducerFileGroupResponse](
+        GetReducerFileGroup(shuffleId, false, SerdeVersion.V1))
+
+    // allocate the celeborn shuffle id the way a map task writer does, so the
+    // app shuffle stays mapped while the copy is released mid-job
+    val releasedShuffleId = lifecycleManager.self.askSync[PbGetShuffleIdResponse](
+      PbGetShuffleId.newBuilder()
+        .setAppShuffleId(appShuffleId)
+        .setAppShuffleIdentifier(s"$appShuffleId-0-0")
+        .setIsShuffleWriter(true)
+        .setIsBarrierStage(false)
+        .build()).getShuffleId
+    assert(lifecycleManager.requestMasterRequestSlotsWithRetry(
+      releasedShuffleId,
+      ids).status == StatusCode.SUCCESS)
+    lifecycleManager.registeredShuffle.add(releasedShuffleId)
+    lifecycleManager.commitManager.setStageEnd(releasedShuffleId)
+    assert(getReducerFileGroup(releasedShuffleId).status == StatusCode.SUCCESS)
+
+    lifecycleManager.unregisterShuffle(releasedShuffleId)
+    eventually(timeout(120.seconds), interval(1.seconds)) {
+      assert(!lifecycleManager.registeredShuffle.contains(releasedShuffleId))
+    }
+    assert(getReducerFileGroup(releasedShuffleId).status == StatusCode.SHUFFLE_EXPIRED)
+
+    // a final release of an id that no app shuffle maps any more is not a tombstone
+    lifecycleManager.unregisterShuffle(neverMappedShuffleId)
+    val emptyStage = getReducerFileGroup(neverMappedShuffleId)
+    assert(emptyStage.status == StatusCode.SHUFFLE_UNREGISTERED)
+    assert(emptyStage.fileGroup.isEmpty)
+
+    lifecycleManager.stop()
+  }
+
   override def afterAll(): Unit = {
     logInfo("all test complete , stop celeborn mini cluster")
     shutdownMiniCluster()

tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornExpiredShuffleReadSuite.scala

@@ -0,0 +1,190 @@
+/*
+ * 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.tests.spark
+
+import java.io.File
+import java.util.concurrent.{CountDownLatch, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicReference}
+
+import org.apache.spark.{FetchFailed, SparkConf, TaskContext, TaskEndReason}
+import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}
+import org.apache.spark.shuffle.ShuffleHandle
+import org.apache.spark.shuffle.celeborn.{CelebornShuffleHandle, ShuffleManagerHook, SparkUtils, TestCelebornShuffleManager}
+import org.apache.spark.sql.SparkSession
+import org.scalatest.BeforeAndAfterEach
+import org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.celeborn.client.ShuffleClient
+import org.apache.celeborn.common.CelebornConf
+import org.apache.celeborn.common.protocol.ShuffleMode
+import org.apache.celeborn.common.protocol.message.StatusCode
+import org.apache.celeborn.service.deploy.worker.Worker
+
+/**
+ * Replays a reader that resolved a celeborn shuffle id before the copy was
+ * released by a stage rerun, then asked for its file group only after the
+ * failed-shuffle cleaner had removed the copy from the workers.
+ */
+class CelebornExpiredShuffleReadSuite extends AnyFunSuite
+  with SparkTestBase
+  with BeforeAndAfterEach {
+
+  override def beforeAll(): Unit = {
+    logInfo("test initialized , setup Celeborn mini cluster")
+    setupMiniClusterWithRandomPorts(workerNum = 1)
+  }
+
+  override def beforeEach(): Unit = {
+    ShuffleClient.reset()
+  }
+
+  override def afterEach(): Unit = {
+    System.gc()
+  }
+
+  override def createWorker(map: Map[String, String]): Worker = {
+    val storageDir = createTmpDir()
+    workerDirs = workerDirs :+ storageDir
+    super.createWorker(map ++ Map("celeborn.master.heartbeat.worker.timeout" -> "10s"), storageDir)
+  }
+
+  test("a reader holding a released shuffle id fails instead of reading it as empty") {
+    if (Spark3OrNewer) {
+      val sparkConf = new SparkConf().setAppName("rss-demo").setMaster("local[2,3]")
+      val sparkSession = SparkSession.builder()
+        .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
+        .config("spark.celeborn.shuffle.forceFallback.partition.enabled", false)
+        .config("spark.celeborn.client.spark.stageRerun.enabled", "true")
+        .config("spark.celeborn.client.spark.fetch.cleanFailedShuffle", "true")
+        .config("spark.celeborn.client.shuffle.expired.checkInterval", "5s")
+        .config(
+          "spark.shuffle.manager",
+          "org.apache.spark.shuffle.celeborn.TestCelebornShuffleManager")
+        .getOrCreate()
+      val celebornConf = SparkUtils.fromSparkConf(sparkSession.sparkContext.getConf)
+
+      val staleTaskFinished = new CountDownLatch(1)
+      val staleTaskEnd = new AtomicReference[TaskEndReason]()
+      val staleTaskRecordsRead = new AtomicReference[java.lang.Long](-1L)
+      val hook = new StaleReaderHook(celebornConf, workerDirs, staleTaskFinished)
+      TestCelebornShuffleManager.registerReaderGetHook(hook)
+      sparkSession.sparkContext.addSparkListener(new SparkListener {
+        override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = {
+          if (taskEnd.stageAttemptId == 0 && taskEnd.stageId == hook.reduceStageId.get() &&
+            taskEnd.taskInfo.index == StaleReaderHook.StalePartition) {
+            staleTaskEnd.set(taskEnd.reason)
+            if (taskEnd.taskMetrics != null) {
+              staleTaskRecordsRead.set(taskEnd.taskMetrics.shuffleReadMetrics.recordsRead)
+            }
+            staleTaskFinished.countDown()
+          }
+        }
+      })
+
+      val tuples = sparkSession.sparkContext.parallelize(1 to 10000, 2)
+        .map { i => (i, i) }.groupByKey(4).collect()
+
+      assert(hook.observedRelease.get(), "stale reader never saw the old copy leave the workers")
+      assert(staleTaskFinished.await(60, TimeUnit.SECONDS))
+      staleTaskEnd.get() match {
+        case FetchFailed(_, _, _, _, _, message) =>
+          assert(message.contains(StatusCode.SHUFFLE_EXPIRED.toString), message)
+        case other =>
+          fail(s"stale reader ended with $other after reading " +
+            s"${staleTaskRecordsRead.get()} records from a released shuffle")
+      }
+      assert(tuples.length == 10000)
+      for (elem <- tuples) {
+        elem._2.foreach(i => assert(i.equals(elem._1)))
+      }
+      sparkSession.stop()
+    }
+  }
+}
+
+object StaleReaderHook {
+  val StalePartition = 0
+}
+
+/**
+ * For the first attempt of one reduce partition: resolve the celeborn shuffle id
+ * the way the real reader does, delete the shuffle files so every other reader
+ * of this copy hits a fetch failure and triggers the rerun, then block until the
+ * failed-shuffle cleaner has removed the copy from the worker. The reader then
+ * proceeds with the id it resolved before the release and no cached file group,
+ * exactly like a task that was stuck on the file-group broadcast in production.
+ * The rerun's reader of the
+ * same partition is held back until the stale one has finished, so a wrong empty
+ * success would be the result Spark records.
+ */
+class StaleReaderHook(
+    conf: CelebornConf,
+    workerDirs: Seq[String],
+    staleTaskFinished: CountDownLatch)
+  extends ShuffleManagerHook {
+
+  val observedRelease = new AtomicBoolean(false)
+  val reduceStageId = new AtomicInteger(-1)
+  private val armed = new AtomicBoolean(false)
+
+  private def shuffleDirs(appUniqueId: String, celebornShuffleId: Int): Seq[File] =
+    workerDirs.map { dir =>
+      new File(s"$dir/celeborn-worker/shuffle_data/$appUniqueId/$celebornShuffleId")
+    }
+
+  override def exec(
+      handle: ShuffleHandle,
+      startPartition: Int,
+      endPartition: Int,
+      context: TaskContext): Unit = {
+    if (startPartition != StaleReaderHook.StalePartition) {
+      return
+    }
+    if (context.stageAttemptNumber() != 0) {
+      staleTaskFinished.await(60, TimeUnit.SECONDS)
+      return
+    }
+    if (!armed.compareAndSet(false, true)) {
+      return
+    }
+    val h = handle.asInstanceOf[CelebornShuffleHandle[_, _, _]]
+    reduceStageId.set(context.stageId())
+    val shuffleClient = ShuffleClient.get(
+      h.appUniqueId,
+      h.lifecycleManagerHost,
+      h.lifecycleManagerPort,
+      conf,
+      h.userIdentifier,
+      h.extension)
+    val celebornShuffleId = SparkUtils.celebornShuffleId(shuffleClient, h, context, false)
+    val dirs = shuffleDirs(h.appUniqueId, celebornShuffleId)
+    val dataFiles = dirs.filter(_.exists()).flatMap(_.listFiles())
+    if (dataFiles.isEmpty) {
+      throw new RuntimeException("unexpected, there must be some data file")
+    }
+    dataFiles.foreach(_.delete())
+    val deadline = System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(120)
+    while (dirs.exists(_.exists()) && System.currentTimeMillis() < deadline) {
+      Thread.sleep(500)
+    }
+    observedRelease.set(!dirs.exists(_.exists()))
+    // the executors in production held no usable file group for the released copy
+    // (their broadcast fetch had failed), so the reader reloaded it over RPC
+    shuffleClient.cleanupShuffle(celebornShuffleId)
+  }
+}