From bf9d308ada8a1259e5b9b8206172e0f67982d18c Mon Sep 17 00:00:00 2001 From: Fei Wang Date: Wed, 21 Feb 2024 23:03:11 -0800 Subject: [PATCH] Fix bug that callback function may hang when unchecked exception missed --- .../common/network/client/TransportClient.java | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java index b5796fe9812..7054c2f93d7 100644 --- a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java +++ b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java @@ -261,11 +261,16 @@ public ByteBuffer sendRpcSync(ByteBuffer message, long timeoutMs) throws IOExcep new RpcResponseCallback() { @Override public void onSuccess(ByteBuffer response) { - ByteBuffer copy = ByteBuffer.allocate(response.remaining()); - copy.put(response); - // flip "copy" to make it readable - copy.flip(); - result.set(copy); + try { + ByteBuffer copy = ByteBuffer.allocate(response.remaining()); + copy.put(response); + // flip "copy" to make it readable + copy.flip(); + result.set(copy); + } catch (Throwable t) { + logger.warn("Error in responding RPC callback", t); + result.setException(t); + } } @Override