should reduce active threads when fatal exception

apache · Aug 30, 2020 · eb8b0b3 · eb8b0b3
1 parent cfe012a
commit eb8b0b3
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 0 deletions.
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/Inbox.scala b/core/src/main/scala/org/apache/spark/rpc/netty/Inbox.scala
@@ -200,6 +200,16 @@ private[netty] class Inbox(val endpointName: String, val endpoint: RpcEndpoint)
    * Calls action closure, and calls the endpoint's onError function in the case of exceptions.
    */
   private def safelyCall(endpoint: RpcEndpoint)(action: => Unit): Unit = {
+    def dealWithFatalException(fatal: Throwable): Unit = {
+      inbox.synchronized {
+        // Should reduce the number of active threads before throw exception
+        if (numActiveThreads != 0) {
+          numActiveThreads -= 1
+        }
+      }
+      throw fatal
+    }
+
     try action catch {
       case NonFatal(e) =>
         try endpoint.onError(e) catch {
@@ -209,8 +219,18 @@ private[netty] class Inbox(val endpointName: String, val endpoint: RpcEndpoint)
             } else {
               logError("Ignoring error", ee)
             }
+          case fatal =>
+            dealWithFatalException(fatal)
         }
+      case fatal =>
+        dealWithFatalException(fatal)
     }
   }
 
+  // exposed only for testing
+  def getNumActiveThreads: Int = {
+    inbox.synchronized {
+      inbox.numActiveThreads
+    }
+  }
 }
diff --git a/core/src/test/scala/org/apache/spark/rpc/netty/InboxSuite.scala b/core/src/test/scala/org/apache/spark/rpc/netty/InboxSuite.scala
@@ -136,4 +136,17 @@ class InboxSuite extends SparkFunSuite {
 
     endpoint.verifySingleOnNetworkErrorMessage(cause, remoteAddress)
   }
+
+  test("SPARK-32738: should reduce active thread when fatal exception") {
+    val endpoint = mock(classOf[TestRpcEndpoint])
+    when(endpoint.receive).thenThrow(new OutOfMemoryError())
+
+    val dispatcher = mock(classOf[Dispatcher])
+    val inbox = new Inbox("name", endpoint)
+    inbox.post(OneWayMessage(null, "hi"))
+    intercept[OutOfMemoryError] {
+      inbox.process(dispatcher)
+    }
+    assert(inbox.getNumActiveThreads == 0)
+  }
 }