[CELEBORN-1103][BUG] only clean up expire data for good disks

### What changes were proposed in this pull request? **When a bad disk occurs, cleaning up expired shuffle keys can cause NullPointerException appearing in the thread pool obtained from `diskOperators` in `StorageManager`. Therefore, only cleaning up expired shuffle keys from good disks will not cause the above problems.** https://issues.apache.org/jira/browse/CELEBORN-1103 ### Why are the changes needed? bugfix ### Does this PR introduce _any_ user-facing change? ### How was this patch tested? Closes apache#2060 from suizhe007/CELEBORN-1103. Lead-authored-by: qinrui <qr7972@gmail.com> Co-authored-by: qinrui <51885730+suizhe007@users.noreply.github.com> Signed-off-by: zky.zhoukeyong <zky.zhoukeyong@alibaba-inc.com>
shohamy7 · Oct 31, 2023 · 232a44b · 232a44b
1 parent 11fe324
commit 232a44b
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
@@ -560,7 +560,9 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
           }
         }
         val (appId, shuffleId) = Utils.splitShuffleKey(shuffleKey)
-        disksSnapshot().filter(_.status != DiskStatus.IO_HANG).foreach { diskInfo =>
+        disksSnapshot().filter(diskInfo =>
+          diskInfo.status == DiskStatus.HEALTHY
+            || diskInfo.status == DiskStatus.HIGH_DISK_USAGE).foreach { diskInfo =>
           diskInfo.dirs.foreach { dir =>
             val file = new File(dir, s"$appId/$shuffleId")
             deleteDirectory(file, diskOperators.get(diskInfo.mountPoint))
@@ -606,7 +608,9 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
 
   private def cleanupExpiredAppDirs(expireDuration: Long): Unit = {
     val diskInfoAndAppDirs = disksSnapshot()
-      .filter(_.status != DiskStatus.IO_HANG)
+      .filter(diskInfo =>
+        diskInfo.status == DiskStatus.HEALTHY
+          || diskInfo.status == DiskStatus.HIGH_DISK_USAGE)
       .map { case diskInfo =>
         (diskInfo, diskInfo.dirs.filter(_.exists).flatMap(_.listFiles()))
       }