[SPARK-28160][CORE] Fix a bug that callback function may hang when unchecked exception missed

[LantaoJin]

What changes were proposed in this pull request?

This is very like #23590 .

ByteBuffer.allocate may throw OutOfMemoryError when the response is large but no enough memory is available. However, when this happens, TransportClient.sendRpcSync will just hang forever if the timeout set to unlimited.

This PR catches Throwable and uses the error to complete SettableFuture.

How was this patch tested?

I tested in my IDE by setting the value of size to -1 to verify the result. Without this patch, it won't be finished until timeout (May hang forever if timeout set to MAX_INT), or the expected IllegalArgumentException will be caught.

@Override
      public void onSuccess(ByteBuffer response) {
        try {
          int size = response.remaining();
          ByteBuffer copy = ByteBuffer.allocate(size); // set size to -1 in runtime when debug
          copy.put(response);
          // flip "copy" to make it readable
          copy.flip();
          result.set(copy);
        } catch (Throwable t) {
          result.setException(t);
        }
      }

[LantaoJin]

cc @zsxwing

[srowen]

Seems OK to me; are there other callbacks that might have the same problem?

[LantaoJin]

@srowen , I only find one place in ExternalShuffleClient.removeBlocks which maybe has a similar problem (not OOM, just uncaught runtime exception):

public Future<Integer> removeBlocks(
      String host,
      int port,
      String execId,
      String[] blockIds) throws IOException, InterruptedException {
    checkInit();
    CompletableFuture<Integer> numRemovedBlocksFuture = new CompletableFuture<>();
    ByteBuffer removeBlocksMessage = new RemoveBlocks(appId, execId, blockIds).toByteBuffer();
    final TransportClient client = clientFactory.createClient(host, port);
    client.sendRpc(removeBlocksMessage, new RpcResponseCallback() {
      @Override
      public void onSuccess(ByteBuffer response) {
        BlockTransferMessage msgObj = BlockTransferMessage.Decoder.fromByteBuffer(response);
        numRemovedBlocksFuture.complete(((BlocksRemoved)msgObj).numRemovedBlocks);
        client.close();
      }

I prefer to change to below code since fromByteBuffer could throw IllegalArgumentException

      @Override
      public void onSuccess(ByteBuffer response) {
        try {
          BlockTransferMessage msgObj = BlockTransferMessage.Decoder.fromByteBuffer(response);
          numRemovedBlocksFuture.complete(((BlocksRemoved) msgObj).numRemovedBlocks);
        } catch (Exception e) {
          logger.warn("Error trying to remove RDD blocks " + Arrays.toString(blockIds) +
            " via external shuffle service from executor: " + execId, e);
          numRemovedBlocksFuture.complete(0);
        } finally {
          client.close();
        }
      }

[LantaoJin]

Should I fix above code if needed in this PR or file a new one?

[srowen]

I think you can fix the similar issue here, and update the title/description.

[srowen]

Do you want to catch Throwable here?

[LantaoJin]

I haven’t seen any error could occur here, so I use Exception. Throwable of course is fine.

[srowen]

Do you want to log the throwable here just for completeness?

[LantaoJin]

I will add a warning log.

[LantaoJin]

Gentle ping @srowen

[SparkQA]

Test build #4813 has finished for PR 24964 at commit d2330cc.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[srowen]

Merged to master/2.4/2.3