Use tryWithSafeFinallyAndFailureCallbacks

apache · Jan 9, 2018 · c1278b8 · c1278b8
1 parent 8e5ffa4
commit c1278b8
Showing 1 changed file with 5 additions and 9 deletions.
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
@@ -332,18 +332,14 @@ private[netty] class NettyRpcEnv(
 
     val pipe = Pipe.open()
     val source = new FileDownloadChannel(pipe.source())
-    var exceptionThrown = true
-    try {
+    Utils.tryWithSafeFinallyAndFailureCallbacks(block = {
       val client = downloadClient(parsedUri.getHost(), parsedUri.getPort())
       val callback = new FileDownloadCallback(pipe.sink(), source, client)
       client.stream(parsedUri.getPath(), callback)
-      exceptionThrown = false
-    } finally {
-      if (exceptionThrown) {
-        pipe.sink().close()
-        source.close()
-      }
-    }
+    })(catchBlock = {
+      pipe.sink().close()
+      source.close()
+    })
 
     source
   }