This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new b2f1da858f [CELEBORN-2377] Support Heartbeat message decoding in C++
client
b2f1da858f is described below
commit b2f1da858fa374b6755b7e50b9fd875b517962a5
Author: yz <[email protected]>
AuthorDate: Thu Jul 9 19:59:16 2026 +0800
[CELEBORN-2377] Support Heartbeat message decoding in C++ client
### What changes were proposed in this pull request?
This PR adds Heartbeat message decoding support to the C++ client network.
### Why are the changes needed?
The HEARTBEAT type (id 22) is already declared in the Message::Type enum
and recognized by decodeType, but Message::decodeFrom had no case for it, so
any heartbeat frame received by the C++ client hit the unsupported Message
decode type failure path. As the protocol/server side begins to send heartbeat
messages, the C++ client must be able to decode them instead of throwing.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce any user-facing change?
- [ ] Yes
### How was this patch tested?
Added MessageTest.decodeHeartbeat, which constructs a heartbeat wire frame
and asserts that Message::decodeFrom returns a non-null message with type
HEARTBEAT. Built and ran the test with --
gtest_filter="MessageTest.decodeHeartbeat" inside the
jraaaay/celeborn-cpp-dev:0.4 container; it passed.
Closes #3755 from Kalvin2077/main.
Authored-by: yz <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
---
cpp/celeborn/network/Message.cpp | 17 ++++++++++++
cpp/celeborn/network/Message.h | 17 ++++++++++++
cpp/celeborn/network/MessageDispatcher.cpp | 3 +++
.../network/tests/MessageDispatcherTest.cpp | 31 ++++++++++++++++++++++
cpp/celeborn/network/tests/MessageTest.cpp | 27 +++++++++++++++++++
5 files changed, 95 insertions(+)
diff --git a/cpp/celeborn/network/Message.cpp b/cpp/celeborn/network/Message.cpp
index 39911849a5..912a9663ca 100644
--- a/cpp/celeborn/network/Message.cpp
+++ b/cpp/celeborn/network/Message.cpp
@@ -102,6 +102,8 @@ std::unique_ptr<Message> Message::decodeFrom(
return ChunkFetchSuccess::decodeFrom(std::move(data));
case CHUNK_FETCH_FAILURE:
return ChunkFetchFailure::decodeFrom(std::move(data));
+ case HEARTBEAT:
+ return Heartbeat::decodeFrom(std::move(data));
default:
CELEBORN_FAIL("unsupported Message decode type " + std::to_string(type));
}
@@ -218,5 +220,20 @@ void PushMergedData::internalEncodeTo(
protocol::encode(buffer, partitionUniqueIds_);
protocol::encode(buffer, batchOffsets_);
}
+
+std::unique_ptr<Heartbeat> Heartbeat::decodeFrom(
+ std::unique_ptr<memory::ReadOnlyByteBuffer>&& data) {
+ CELEBORN_CHECK_EQ(data->remainingSize(), 1);
+ data->skip(1);
+ return std::make_unique<Heartbeat>();
+}
+
+int Heartbeat::internalEncodedLength() const {
+ return 1;
+}
+
+void Heartbeat::internalEncodeTo(memory::WriteOnlyByteBuffer& buffer) const {
+ buffer.write<uint8_t>(0);
+}
} // namespace network
} // namespace celeborn
diff --git a/cpp/celeborn/network/Message.h b/cpp/celeborn/network/Message.h
index 3b91241385..3b5ef2fa8d 100644
--- a/cpp/celeborn/network/Message.h
+++ b/cpp/celeborn/network/Message.h
@@ -322,5 +322,22 @@ class PushMergedData : public Message {
std::vector<std::string> partitionUniqueIds_;
std::vector<int32_t> batchOffsets_;
};
+
+class Heartbeat : public Message {
+ public:
+ Heartbeat()
+ : Message(HEARTBEAT, memory::ReadOnlyByteBuffer::createEmptyBuffer()) {}
+
+ Heartbeat(const Heartbeat& other)
+ : Message(HEARTBEAT, other.body_->clone()) {}
+
+ static std::unique_ptr<Heartbeat> decodeFrom(
+ std::unique_ptr<memory::ReadOnlyByteBuffer>&& data);
+
+ private:
+ int internalEncodedLength() const override;
+
+ void internalEncodeTo(memory::WriteOnlyByteBuffer& buffer) const override;
+};
} // namespace network
} // namespace celeborn
diff --git a/cpp/celeborn/network/MessageDispatcher.cpp
b/cpp/celeborn/network/MessageDispatcher.cpp
index 86ec141110..f70a396f86 100644
--- a/cpp/celeborn/network/MessageDispatcher.cpp
+++ b/cpp/celeborn/network/MessageDispatcher.cpp
@@ -120,6 +120,9 @@ void MessageDispatcher::read(Context*,
std::unique_ptr<Message> toRecvMsg) {
}
return;
}
+ case Message::HEARTBEAT: {
+ return;
+ }
default: {
LOG(ERROR) << "unsupported msg for dispatcher";
}
diff --git a/cpp/celeborn/network/tests/MessageDispatcherTest.cpp
b/cpp/celeborn/network/tests/MessageDispatcherTest.cpp
index 6127548d67..45d2d4e82b 100644
--- a/cpp/celeborn/network/tests/MessageDispatcherTest.cpp
+++ b/cpp/celeborn/network/tests/MessageDispatcherTest.cpp
@@ -284,3 +284,34 @@ TEST(MessageDispatcherTest,
sendFetchChunkRequestAndReceiveFailure) {
EXPECT_TRUE(future.hasException());
}
+
+TEST(MessageDispatcherTest, heartbeatIsSilentlyConsumed) {
+ std::unique_ptr<Message> sentMsg;
+ MockHandler mockHandler(sentMsg);
+ auto mockPipeline = createMockedPipeline(std::move(mockHandler));
+ auto dispatcher = std::make_unique<MessageDispatcher>();
+ dispatcher->setPipeline(mockPipeline.get());
+
+ const long requestId = 1001;
+ const std::string requestBody = "test-request-body";
+ auto rpcRequest = std::make_unique<RpcRequest>(
+ requestId, toReadOnlyByteBuffer(requestBody));
+ auto future = dispatcher->sendRpcRequest(std::move(rpcRequest));
+
+ EXPECT_FALSE(future.isReady());
+
+ // A heartbeat arriving mid-request must be a no-op: it should not fulfill
+ // the pending future, not throw, and not close the dispatcher.
+ auto heartbeat = std::make_unique<Heartbeat>();
+ dispatcher->read(nullptr, std::move(heartbeat));
+
+ EXPECT_FALSE(future.isReady());
+ EXPECT_TRUE(dispatcher->isAvailable());
+
+ const std::string responseBody = "test-response-body";
+ auto rpcResponse = std::make_unique<RpcResponse>(
+ requestId, toReadOnlyByteBuffer(responseBody));
+ dispatcher->read(nullptr, std::move(rpcResponse));
+
+ EXPECT_TRUE(future.isReady());
+}
diff --git a/cpp/celeborn/network/tests/MessageTest.cpp
b/cpp/celeborn/network/tests/MessageTest.cpp
index 3e18f77cda..09d05826ff 100644
--- a/cpp/celeborn/network/tests/MessageTest.cpp
+++ b/cpp/celeborn/network/tests/MessageTest.cpp
@@ -255,3 +255,30 @@ TEST(MessageTest, pushMergedDataCopyConstructor) {
EXPECT_EQ(copy.batchOffsets().size(), 2);
EXPECT_EQ(copy.batchOffsets()[1], 4);
}
+
+TEST(MessageTest, decodeHeartbeat) {
+ const int headerLength = sizeof(int32_t) + sizeof(uint8_t) + sizeof(int32_t);
+ const int encodedLength = 1;
+ const int bodyLength = 0;
+ size_t size = headerLength + encodedLength + bodyLength;
+ auto writeBuffer = memory::ByteBuffer::createWriteOnly(size);
+ writeBuffer->write<int32_t>(encodedLength);
+ writeBuffer->write<uint8_t>(Message::Type::HEARTBEAT);
+ writeBuffer->write<int32_t>(bodyLength);
+ writeBuffer->write<uint8_t>(0);
+ auto message = Message::decodeFrom(
+ memory::ByteBuffer::toReadOnly(std::move(writeBuffer)));
+ EXPECT_NE(message, nullptr);
+ EXPECT_EQ(message->type(), Message::Type::HEARTBEAT);
+ auto heartbeatBody = message->body();
+ EXPECT_EQ(heartbeatBody->remainingSize(), 0);
+}
+
+TEST(MessageTest, encodeHeartbeat) {
+ Heartbeat heartbeat;
+ auto encodedBuffer = heartbeat.encode();
+ EXPECT_EQ(encodedBuffer->read<int32_t>(), 1);
+ EXPECT_EQ(encodedBuffer->read<uint8_t>(), Message::Type::HEARTBEAT);
+ EXPECT_EQ(encodedBuffer->read<int32_t>(), 0);
+ EXPECT_EQ(encodedBuffer->read<uint8_t>(), 0);
+}