add setException for all Throwable

zaynt4606 · Dec 11, 2024 · 3f4efdb · 3f4efdb
1 parent dec0f31
commit 3f4efdb
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 9 deletions.
diff --git a/common/src/main/java/org/apache/celeborn/common/util/ExceptionUtils.java b/common/src/main/java/org/apache/celeborn/common/util/ExceptionUtils.java
@@ -44,6 +44,14 @@ public static Throwable wrapIOExceptionToUnRetryable(Throwable throwable) {
     }
   }
 
+  public static IOException wrapThrowableToIOException(Throwable throwable) {
+    if (throwable instanceof IOException) {
+      return (IOException) throwable;
+    } else {
+      return new IOException(throwable.getMessage(), throwable);
+    }
+  }
+
   public static String stringifyException(Throwable exception) {
     if (exception == null) {
       return "(null)";

diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/Flusher.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/Flusher.scala
@@ -31,7 +31,7 @@ import org.apache.celeborn.common.internal.Logging
 import org.apache.celeborn.common.meta.{DiskStatus, TimeWindow}
 import org.apache.celeborn.common.metrics.source.{AbstractSource, ThreadPoolSource}
 import org.apache.celeborn.common.protocol.StorageInfo
-import org.apache.celeborn.common.util.{ThreadUtils, Utils}
+import org.apache.celeborn.common.util.{ExceptionUtils, ThreadUtils, Utils}
 import org.apache.celeborn.service.deploy.worker.WorkerSource
 import org.apache.celeborn.service.deploy.worker.WorkerSource.FLUSH_WORKING_QUEUE_SIZE
 import org.apache.celeborn.service.deploy.worker.congestcontrol.CongestionController
@@ -79,14 +79,9 @@ abstract private[worker] class Flusher(
                   }
                 } catch {
                   case t: Throwable =>
-                    t match {
-                      case exception: IOException =>
-                        task.notifier.setException(exception)
-                        processIOException(
-                          exception,
-                          DiskStatus.READ_OR_WRITE_FAILURE)
-                      case _ =>
-                    }
+                    val e = ExceptionUtils.wrapThrowableToIOException(t)
+                    task.notifier.setException(e)
+                    processIOException(e, DiskStatus.READ_OR_WRITE_FAILURE)
                     logWarning(s"Flusher-$this-thread-$index encounter exception.", t)
                 }
                 lastBeginFlushTime.set(index, -1)