This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch branch-0.7
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/branch-0.7 by this push:
new e7cc4058c [CELEBORN-2413] Fail transport requests retriably when the
C++ connection is closed or a write fails
e7cc4058c is described below
commit e7cc4058cc8c09da305e66ade021b4c335bcf243
Author: Yu Gan <[email protected]>
AuthorDate: Tue Aug 18 17:08:05 2026 +0800
[CELEBORN-2413] Fail transport requests retriably when the C++ connection
is closed or a write fails
### What changes were proposed in this pull request?
`MessageDispatcher` guards both send paths with a hard
`CELEBORN_CHECK(!closed_)` and drops the future returned by the write, so
sending on a closed connection raises a **non-retriable** error and a failed
write is never reported at all — the request's promise stays pending until its
timeout:
```cpp
folly::Future<std::unique_ptr<Message>> MessageDispatcher::operator()(...) {
CELEBORN_CHECK(!closed_); // throws a non-retriable CelebornRuntimeError
...
this->pipeline_->write(std::move(toSendMsg)); // future discarded
CELEBORN_CHECK(!closed_); // throws a non-retriable CelebornRuntimeError
return f;
}
```
This PR:
- Merges the two request registries and the closed flag into one
`ConnectionState` under a single mutex, so a request is either registered
before the connection is retired — and the retirement then fails it — or
refused outright with a ready retriable future. Both `CELEBORN_CHECK`s go away.
- Routes every connection-killing event through
`ConnectionState::retire(reason)`, the analogue of Java's
`failOutstandingRequests`: it marks the connection unavailable, so
`TransportClient::active()` reports false and the client pool stops handing it
out, then fails everything outstanding retriably. It is idempotent, and
replaces `cleanup()` on the close path with the per-request messages unchanged.
- Observes the write future on all three send paths, which is where a
socket write failure surfaces (`wangle::AsyncSocketHandler::write` fails it
either immediately or later from `writeErr`, without necessarily flipping the
closed flag first). Like Java's `StdChannelListener`, the connection is retired
before the failure is reported: `retire()` flips the flag and drains the
registries in one critical section, then fails the promises — which matters
because `setException` runs the call [...]
- Handles the other way a write can fail: by throwing. Messages are
serialized on the way down the pipeline and `wangle::Pipeline::write` has no
try/catch, so that failure never reaches the write future and left the request
registered forever. It is now unregistered, and the exception keeps propagating
— a broken message is not a broken connection. The fire-and-forget path logs
it, since `~WorkerPartitionReader` sends `BufferStreamEnd` there.
- Keeps the continuations independent of the dispatcher's lifetime:
`ConnectionState` is refcounted separately and every continuation and interrupt
handler holds a `weak_ptr` to it. A write failure is reported from the socket's
write callback, so it can arrive during teardown, where capturing `this` is a
deterministic segfault. Member reordering can't fix that —
`~ClientDispatcherBase` does `pipeline_->remove(this).finalize()`, so the
dispatcher must be destroyed while its pipeline is [...]
- Preserves `isRetriable` through `TransportClient`, which previously
re-threw via `CELEBORN_FAIL` or flattened the cause into a plain
`std::runtime_error`. Non-retriable causes keep byte-identical messages, so
`getPushDataFailCause` string matching is unaffected.
- Makes `TransportClientFactory::createClient` throw a retriable error on a
failed connect.
- Reports a duplicate registration through `CELEBORN_CHECK` instead of
`registries[key]` handing back the existing holder, where `setInterruptHandler`
threw a bare `std::logic_error` from inside the critical section.
The retirement reason is phrased `"Failed to send request {}, errorMsg:
{}"` to match the Java listener, so `connectFail()` classifies it — and the
requests the dead connection takes down with it — as
`PUSH_DATA_CONNECTION_EXCEPTION_PRIMARY` rather than as the non-critical
default.
### Why are the changes needed?
A closed connection or a failed write is a normal recoverable condition — a
worker restart, an idle timeout — not an invariant violation. Java fails the
affected requests (`TransportResponseHandler#channelInactive` →
`failOutstandingRequests`, `StdChannelListener#operationComplete` →
`handleFailure`) and `CelebornInputStream#createReaderWithRetry` retries or
fails over to the replica; the C++ client hard-failed instead, or stalled until
the request timeout. Leaving the connection avai [...]
No C++ path branches on `isRetriable()` yet, but until now the
classification was not observable at the public API at all.
### Does this PR resolve a correctness bug?
- [ ] Yes
Shuffle output is unaffected: this is a failure-handling fix.
### Does this PR introduce _any_ user-facing change?
- [ ] Yes
No config or API change. A connection-closed or failed-write condition that
previously failed non-retriably, or hung until timeout, now fails the affected
request retriably.
### How was this patch tested?
New tests in `MessageDispatcherTest.cpp`:
- `sendRpcRequestAfterCloseFailsRetriably`,
`sendFetchChunkRequestAfterCloseFailsRetriably`,
`closeFailsInFlightRequestsRetriably` — sends after `close()`, and requests in
flight when it happens.
- `sendRpcRequestFailedWriteRetiresConnection`,
`sendFetchChunkRequestFailedWriteRetiresConnection` — a `MockHandler` failing
from the second write on, so an earlier request is still in flight on the
connection the failing write kills. Each asserts both requests fail retriably,
`isAvailable()` is false, and a further send fails fast.
- `failedWriteWithoutResponseRetiresConnection`,
`sendWithoutResponseAfterCloseIsSkipped` — the fire-and-forget path, previously
uncovered, where the old `CELEBORN_CHECK` threw out of a destructor.
- `throwingWriteUnregistersTheRequest`,
`throwingFetchWriteUnregistersTheRequest`,
`throwingWriteWithoutResponseIsReported`,
`nonStdThrowingWriteWithoutResponseIsReported` — the mock rejects one message
by throwing, with and without deriving from `std::exception`; the connection
stays available and the rejected request is no longer registered.
- `failedWriteMakesTransportClientInactive` — `TransportClient::active()`
is false after a failed write, which is what the factory gates connection reuse
on.
- `failedWriteAfterDispatcherDestroyedIsIgnored`,
`failedFetchWriteAfterDispatcherDestroyedIsIgnored` — the mock hands its write
promise back to the test, which destroys the dispatcher before failing the
write, as `AsyncSocket` does during teardown. Both segfault against a
continuation capturing `this`.
New tests in `TransportClientTest.cpp` asserting the classification
survives to the public API: `sendRpcRequestSyncPreservesRetriableFailure`,
`pushDataAsyncPreservesRetriableFailure`,
`fetchChunkAsyncPreservesRetriableFailure`.
Not covered: the factory dropping a cached client and connecting again end
to end, which needs a listening socket that no C++ unit test brings up today.
The C++ suite passes locally on macOS/arm64, and is covered by the
`Celeborn Cpp Integration Test` workflow.
Closes #3801 from yugan95/CELEBORN-2413.
Authored-by: Yu Gan <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
(cherry picked from commit 9d3a5b50bd8a7dce4f3445b0a0d54d1187f9f5b1)
Signed-off-by: Nicholas Jiang <[email protected]>
---
cpp/celeborn/network/MessageDispatcher.cpp | 568 ++++++++++++++++-----
cpp/celeborn/network/MessageDispatcher.h | 47 +-
cpp/celeborn/network/TransportClient.cpp | 111 +++-
.../network/tests/MessageDispatcherTest.cpp | 456 +++++++++++++++++
cpp/celeborn/network/tests/TransportClientTest.cpp | 114 +++++
5 files changed, 1129 insertions(+), 167 deletions(-)
diff --git a/cpp/celeborn/network/MessageDispatcher.cpp
b/cpp/celeborn/network/MessageDispatcher.cpp
index c9d873514..90dc15b50 100644
--- a/cpp/celeborn/network/MessageDispatcher.cpp
+++ b/cpp/celeborn/network/MessageDispatcher.cpp
@@ -17,58 +17,301 @@
#include "celeborn/network/MessageDispatcher.h"
+#include <atomic>
+#include <chrono>
+#include <mutex>
+#include <optional>
+#include <string_view>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include <fmt/format.h>
+#include <folly/ExceptionWrapper.h>
+#include <folly/Synchronized.h>
+
#include "celeborn/protocol/TransportMessage.h"
namespace celeborn {
namespace network {
+namespace {
+// Builds a retriable transport error. A closed connection or a failed socket
+// write is a normal recoverable condition -- a worker restart or an idle
+// timeout -- not an invariant violation, so it is reported through the promise
+// instead of asserting, as the Java client does. The caller passes its own
+// __FILE__/__LINE__/__FUNCTION__ so the exception points at the failing site.
+folly::exception_wrapper makeRetriableTransportError(
+ const char* file,
+ size_t line,
+ const char* function,
+ const std::string& detail) {
+ return folly::make_exception_wrapper<utils::CelebornRuntimeError>(
+ file,
+ line,
+ function,
+ /*expression=*/"",
+ /*message=*/detail,
+ utils::error_source::kErrorSourceRuntime.c_str(),
+ utils::error_code::kInvalidState.c_str(),
+ /*isRetriable=*/true);
+}
+
+// The reason reported on the requests that are still outstanding when the
+// client goes away, either through close() or by being destroyed.
+constexpr std::string_view kClientClosed = "Client closed";
+} // namespace
+
+// The requests outstanding on one connection, and the flag that retires it.
+//
+// Both registries and the flag live under a single mutex, so that registering
a
+// request cannot interleave with retiring the connection: a request is either
+// registered before the retirement -- which then fails it -- or refused
+// outright, leaving no close-during-send window to patch up afterwards.
+//
+// Promises are always fulfilled after the mutex has been released: fulfilling
+// one runs the caller's continuation inline, and that continuation may
re-enter
+// the dispatcher or drop the last reference to the TransportClient owning it.
+class MessageDispatcher::ConnectionState
+ : public std::enable_shared_from_this<MessageDispatcher::ConnectionState> {
+ public:
+ using MsgPromise = folly::Promise<std::unique_ptr<Message>>;
+
+ struct MsgPromiseHolder {
+ MsgPromise msgPromise;
+ std::chrono::time_point<std::chrono::system_clock> requestTime;
+ };
+
+ bool retired() const {
+ return retired_.load();
+ }
+
+ // Registers an rpc/push request and returns the future its response is
+ // delivered to, or std::nullopt when the connection has already been retired
+ // and must not be written to.
+ std::optional<folly::Future<std::unique_ptr<Message>>> registerRequest(
+ long requestId);
+
+ // The fetch counterpart of registerRequest.
+ std::optional<folly::Future<std::unique_ptr<Message>>> registerFetch(
+ const protocol::StreamChunkSlice& streamChunkSlice);
+
+ // Unregisters the request a received response belongs to, or returns
+ // std::nullopt when it is no longer registered -- it may have been failed,
+ // drained or interrupted in the meantime.
+ std::optional<MsgPromiseHolder> takeRequest(long requestId);
+
+ // The fetch counterpart of takeRequest.
+ std::optional<MsgPromiseHolder> takeFetch(
+ const protocol::StreamChunkSlice& streamChunkSlice);
+
+ // Retires the connection: marks it unavailable, so that
+ // TransportClient::active() reports false and TransportClientFactory stops
+ // handing it out, and fails everything outstanding on it with a retriable
+ // error carrying `reason` -- the analogue of Java's
+ // TransportResponseHandler#failOutstandingRequests. Idempotent.
+ void retire(std::string_view reason);
+
+ private:
+ struct Registries {
+ std::unordered_map<long, MsgPromiseHolder> requests;
+ std::unordered_map<
+ protocol::StreamChunkSlice,
+ MsgPromiseHolder,
+ protocol::StreamChunkSlice::Hasher>
+ fetches;
+ };
+
+ folly::Synchronized<Registries, std::mutex> registries_;
+
+ // Mutated only while `registries_` is held, so that a registration cannot
+ // slip past a retirement. Read without the lock by isAvailable(), which sits
+ // on the connection-reuse path of every request.
+ std::atomic<bool> retired_{false};
+};
+
+std::optional<folly::Future<std::unique_ptr<Message>>>
+MessageDispatcher::ConnectionState::registerRequest(long requestId) {
+ using Result = std::optional<folly::Future<std::unique_ptr<Message>>>;
+ return registries_.withLock([&](Registries& registries) -> Result {
+ if (retired_.load()) {
+ return std::nullopt;
+ }
+ // requestIds come from a monotonic counter, so a key can only collide if
+ // the same request is sent twice. Report that rather than overwriting the
+ // holder: folly throws from setInterruptHandler when a promise already has
+ // one, which would escape from under this lock as a bare std::logic_error.
+ auto [entry, registered] = registries.requests.try_emplace(requestId);
+ CELEBORN_CHECK(
+ registered,
+ fmt::format("requestId {} is already outstanding", requestId));
+ auto& holder = entry->second;
+ holder.requestTime = std::chrono::system_clock::now();
+ holder.msgPromise.setInterruptHandler(
+ [requestId,
+ weakState = weak_from_this()](const folly::exception_wrapper&) {
+ LOG(WARNING) << "rpc request interrupted, requestId: " << requestId;
+ if (auto state = weakState.lock()) {
+ state->registries_.lock()->requests.erase(requestId);
+ }
+ });
+ return holder.msgPromise.getFuture();
+ });
+}
+
+std::optional<folly::Future<std::unique_ptr<Message>>>
+MessageDispatcher::ConnectionState::registerFetch(
+ const protocol::StreamChunkSlice& streamChunkSlice) {
+ using Result = std::optional<folly::Future<std::unique_ptr<Message>>>;
+ return registries_.withLock([&](Registries& registries) -> Result {
+ if (retired_.load()) {
+ return std::nullopt;
+ }
+ // See registerRequest(): a StreamChunkSlice identifies one chunk of one
+ // stream, and a stream is opened per reader, so a key can only collide if
+ // the same chunk is fetched twice on the same connection.
+ auto [entry, registered] =
registries.fetches.try_emplace(streamChunkSlice);
+ CELEBORN_CHECK(
+ registered,
+ fmt::format(
+ "streamChunkSlice {} is already outstanding",
+ streamChunkSlice.toString()));
+ auto& holder = entry->second;
+ holder.requestTime = std::chrono::system_clock::now();
+ holder.msgPromise.setInterruptHandler(
+ [streamChunkSlice,
+ weakState = weak_from_this()](const folly::exception_wrapper&) {
+ LOG(WARNING) << "fetchChunk request interrupted, streamChunkSlice: "
+ << streamChunkSlice.toString();
+ if (auto state = weakState.lock()) {
+ state->registries_.lock()->fetches.erase(streamChunkSlice);
+ }
+ });
+ return holder.msgPromise.getFuture();
+ });
+}
+
+std::optional<MessageDispatcher::ConnectionState::MsgPromiseHolder>
+MessageDispatcher::ConnectionState::takeRequest(long requestId) {
+ using Result = std::optional<MsgPromiseHolder>;
+ return registries_.withLock([&](Registries& registries) -> Result {
+ auto search = registries.requests.find(requestId);
+ if (search == registries.requests.end()) {
+ return std::nullopt;
+ }
+ auto holder = std::move(search->second);
+ registries.requests.erase(search);
+ return std::move(holder);
+ });
+}
+
+std::optional<MessageDispatcher::ConnectionState::MsgPromiseHolder>
+MessageDispatcher::ConnectionState::takeFetch(
+ const protocol::StreamChunkSlice& streamChunkSlice) {
+ using Result = std::optional<MsgPromiseHolder>;
+ return registries_.withLock([&](Registries& registries) -> Result {
+ auto search = registries.fetches.find(streamChunkSlice);
+ if (search == registries.fetches.end()) {
+ return std::nullopt;
+ }
+ auto holder = std::move(search->second);
+ registries.fetches.erase(search);
+ return std::move(holder);
+ });
+}
+
+void MessageDispatcher::ConnectionState::retire(std::string_view reason) {
+ std::vector<std::pair<std::string, MsgPromiseHolder>> failures;
+ const bool firstToRetire = registries_.withLock([&](Registries& registries) {
+ for (auto& [requestId, holder] : registries.requests) {
+ failures.emplace_back(
+ fmt::format("{}, cancel ongoing requestId {}", reason, requestId),
+ std::move(holder));
+ }
+ registries.requests.clear();
+ for (auto& [streamChunkSlice, holder] : registries.fetches) {
+ failures.emplace_back(
+ fmt::format(
+ "{}, cancel ongoing streamChunkSlice {}",
+ reason,
+ streamChunkSlice.toString()),
+ std::move(holder));
+ }
+ registries.fetches.clear();
+ return !retired_.exchange(true);
+ });
+ if (firstToRetire) {
+ LOG(WARNING) << reason;
+ }
+ for (auto& [detail, holder] : failures) {
+ LOG(WARNING) << detail;
+ holder.msgPromise.setException(
+ makeRetriableTransportError(__FILE__, __LINE__, __FUNCTION__, detail));
+ }
+}
+
+MessageDispatcher::MessageDispatcher()
+ : state_(std::make_shared<ConnectionState>()) {}
+
+MessageDispatcher::~MessageDispatcher() {
+ // The dispatcher can be destroyed with requests still outstanding: the
client
+ // pool replaces a client whose connection went bad while it is in use. Fail
+ // those with the same retriable error close() reports, rather than letting
+ // their futures surface folly's BrokenPromise, which carries no cause the
+ // caller could classify.
+ state_->retire(kClientClosed);
+}
+
+bool MessageDispatcher::isAvailable() {
+ return !state_->retired();
+}
+
+folly::Future<folly::Unit> MessageDispatcher::writeToPipeline(
+ std::unique_ptr<Message> toSendMsg,
+ const std::function<void()>& onThrow) {
+ try {
+ return this->pipeline_->write(std::move(toSendMsg));
+ } catch (...) {
+ onThrow();
+ throw;
+ }
+}
+
void MessageDispatcher::read(Context*, std::unique_ptr<Message> toRecvMsg) {
+ // Hold the state for the whole call: fulfilling a promise below runs the
+ // caller's continuation inline, and that continuation may drop the last
+ // reference to the TransportClient owning this dispatcher, so nothing may
+ // touch `this` once a promise may have been fulfilled.
+ const auto state = state_;
switch (toRecvMsg->type()) {
case Message::RPC_RESPONSE: {
RpcResponse* response = reinterpret_cast<RpcResponse*>(toRecvMsg.get());
- bool found = true;
- auto holder = requestIdRegistry_.withLock([&](auto& registry) {
- auto search = registry.find(response->requestId());
- if (search == registry.end()) {
- LOG(WARNING)
- << "requestId " << response->requestId()
- << " not found when handling RPC_RESPONSE. Might be outdated
already, ignored.";
- found = false;
- return MsgPromiseHolder{};
- }
- auto result = std::move(search->second);
- registry.erase(response->requestId());
- return std::move(result);
- });
- if (found) {
- holder.msgPromise.setValue(std::move(toRecvMsg));
+ auto holder = state->takeRequest(response->requestId());
+ if (!holder) {
+ LOG(WARNING)
+ << "requestId " << response->requestId()
+ << " not found when handling RPC_RESPONSE. Might be outdated
already, ignored.";
+ return;
}
+ holder->msgPromise.setValue(std::move(toRecvMsg));
return;
}
case Message::RPC_FAILURE: {
RpcFailure* failure = reinterpret_cast<RpcFailure*>(toRecvMsg.get());
- bool found = true;
- auto holder = requestIdRegistry_.withLock([&](auto& registry) {
- auto search = registry.find(failure->requestId());
- if (search == registry.end()) {
- LOG(WARNING)
- << "requestId " << failure->requestId()
- << " not found when handling RPC_FAILURE. Might be outdated
already, ignored.";
- found = false;
- return MsgPromiseHolder{};
- }
- auto result = std::move(search->second);
- registry.erase(failure->requestId());
- return std::move(result);
- });
+ auto holder = state->takeRequest(failure->requestId());
+ if (!holder) {
+ LOG(WARNING)
+ << "requestId " << failure->requestId()
+ << " not found when handling RPC_FAILURE. Might be outdated
already, ignored.";
+ }
const std::string errorMsg = failure->errorMsg();
LOG(ERROR) << "Rpc failed, requestId: " << failure->requestId()
<< " errorMsg: " << errorMsg << std::endl;
- if (found) {
+ if (holder) {
// Carry the worker's error message on the exception so the push/fetch
// callbacks can recover the precise cause via
// ShuffleClientImpl::getPushDataFailCause. A blank std::exception
// would collapse every failure into the non-critical default.
- holder.msgPromise.setException(
+ holder->msgPromise.setException(
folly::make_exception_wrapper<std::runtime_error>(errorMsg));
}
return;
@@ -77,52 +320,35 @@ void MessageDispatcher::read(Context*,
std::unique_ptr<Message> toRecvMsg) {
ChunkFetchSuccess* success =
reinterpret_cast<ChunkFetchSuccess*>(toRecvMsg.get());
auto streamChunkSlice = success->streamChunkSlice();
- bool found = true;
- auto holder = streamChunkSliceRegistry_.withLock([&](auto& registry) {
- auto search = registry.find(streamChunkSlice);
- if (search == registry.end()) {
- LOG(WARNING)
- << "streamChunkSlice " << streamChunkSlice.toString()
- << " not found when handling CHUNK_FETCH_SUCCESS. Might be
outdated already, ignored.";
- found = false;
- return MsgPromiseHolder{};
- }
- auto result = std::move(search->second);
- registry.erase(streamChunkSlice);
- return std::move(result);
- });
- if (found) {
- holder.msgPromise.setValue(std::move(toRecvMsg));
+ auto holder = state->takeFetch(streamChunkSlice);
+ if (!holder) {
+ LOG(WARNING)
+ << "streamChunkSlice " << streamChunkSlice.toString()
+ << " not found when handling CHUNK_FETCH_SUCCESS. Might be
outdated already, ignored.";
+ return;
}
+ holder->msgPromise.setValue(std::move(toRecvMsg));
return;
}
case Message::CHUNK_FETCH_FAILURE: {
ChunkFetchFailure* failure =
reinterpret_cast<ChunkFetchFailure*>(toRecvMsg.get());
auto streamChunkSlice = failure->streamChunkSlice();
- bool found = true;
- auto holder = streamChunkSliceRegistry_.withLock([&](auto& registry) {
- auto search = registry.find(streamChunkSlice);
- if (search == registry.end()) {
- LOG(WARNING)
- << "streamChunkSlice " << streamChunkSlice.toString()
- << " not found when handling CHUNK_FETCH_FAILURE. Might be
outdated already, ignored.";
- found = false;
- return MsgPromiseHolder{};
- }
- auto result = std::move(search->second);
- registry.erase(streamChunkSlice);
- return std::move(result);
- });
- std::string errorMsg = fmt::format(
+ auto holder = state->takeFetch(streamChunkSlice);
+ if (!holder) {
+ LOG(WARNING)
+ << "streamChunkSlice " << streamChunkSlice.toString()
+ << " not found when handling CHUNK_FETCH_FAILURE. Might be
outdated already, ignored.";
+ }
+ const std::string errorMsg = fmt::format(
"fetchChunk failed, streamChunkSlice: {}, errorMsg: {}",
streamChunkSlice.toString(),
failure->errorMsg());
LOG(ERROR) << errorMsg;
- if (found) {
+ if (holder) {
// Carry the streamChunkSlice context and the worker's error message so
// the reader's fetch-failure path sees the real cause.
- holder.msgPromise.setException(
+ holder->msgPromise.setException(
folly::make_exception_wrapper<std::runtime_error>(errorMsg));
}
return;
@@ -138,8 +364,6 @@ void MessageDispatcher::read(Context*,
std::unique_ptr<Message> toRecvMsg) {
folly::Future<std::unique_ptr<Message>> MessageDispatcher::operator()(
std::unique_ptr<Message> toSendMsg) {
- CELEBORN_CHECK(!closed_);
- auto currTime = std::chrono::system_clock::now();
long requestId;
switch (toSendMsg->type()) {
case Message::RPC_REQUEST: {
@@ -163,23 +387,55 @@ folly::Future<std::unique_ptr<Message>>
MessageDispatcher::operator()(
}
}
- auto f = requestIdRegistry_.withLock(
- [&](auto& registry) -> folly::Future<std::unique_ptr<Message>> {
- auto& holder = registry[requestId];
- holder.requestTime = currTime;
- auto& p = holder.msgPromise;
- p.setInterruptHandler([requestId,
- this](const folly::exception_wrapper&) {
- this->requestIdRegistry_.lock()->erase(requestId);
- LOG(WARNING) << "rpc request interrupted, requestId: " << requestId;
- });
- return p.getFuture();
- });
-
- this->pipeline_->write(std::move(toSendMsg));
+ // Hold the state for the whole call: see read().
+ const auto state = state_;
+ auto future = state->registerRequest(requestId);
+ if (!future) {
+ // The connection has been retired. Fail with a retriable error rather than
+ // asserting, so the caller's retry/failover logic can recover, and do not
+ // write to a socket that is known to be dead.
+ return folly::makeFuture<std::unique_ptr<Message>>(
+ makeRetriableTransportError(
+ __FILE__,
+ __LINE__,
+ __FUNCTION__,
+ fmt::format(
+ "connection closed before sending requestId {}", requestId)));
+ }
- CELEBORN_CHECK(!closed_);
- return f;
+ // Observe the write future, like Java's TransportClient does with
+ // StdChannelListener. wangle's AsyncSocketHandler::write returns an
+ // already-failed future when the socket is no longer good, and otherwise
+ // fails it later from AsyncTransport::WriteCallback::writeErr; dropping the
+ // future would leave the registered request pending until its timeout.
+ const std::weak_ptr<ConnectionState> weakState = state;
+ auto written = writeToPipeline(std::move(toSendMsg), [&]() {
+ // A handler threw instead of returning a failed future -- the message is
+ // serialized on the way down, by MessageSerializeHandler, and
+ // wangle::Pipeline::write has no try/catch. That is a violation of our own
+ // encoding invariants rather than a connection failure, so it keeps
+ // propagating to the caller; but nothing was sent, so the request just
+ // registered must not be left behind waiting for a response.
+ state->takeRequest(requestId);
+ });
+ std::move(written).thenError(
+ [weakState, requestId](const folly::exception_wrapper& e) {
+ // The socket is dead, so the whole connection is retired -- and
+ // everything outstanding on it failed -- before anything is reported,
+ // the way StdChannelListener closes the channel before calling
+ // handleFailure. That order matters: setException runs the caller's
+ // continuation inline, and it may immediately ask
+ // TransportClientFactory for a client, which must not be the
connection
+ // that just died. Once the dispatcher is gone there is nothing left to
+ // retire; its requests were failed when it was destroyed.
+ if (auto state = weakState.lock()) {
+ state->retire(fmt::format(
+ "Failed to send request {}, errorMsg: {}",
+ requestId,
+ e.what().toStdString()));
+ }
+ });
+ return std::move(*future);
}
folly::Future<std::unique_ptr<Message>> MessageDispatcher::sendPushDataRequest(
@@ -191,30 +447,92 @@ folly::Future<std::unique_ptr<Message>>
MessageDispatcher::sendFetchChunkRequest(
const protocol::StreamChunkSlice& streamChunkSlice,
std::unique_ptr<Message> toSendMsg) {
- CELEBORN_CHECK(!closed_);
CELEBORN_CHECK(toSendMsg->type() == Message::RPC_REQUEST);
- auto f = streamChunkSliceRegistry_.withLock([&](auto& registry) {
- auto& holder = registry[streamChunkSlice];
- holder.requestTime = std::chrono::system_clock::now();
- auto& p = holder.msgPromise;
- p.setInterruptHandler(
- [streamChunkSlice, this](const folly::exception_wrapper&) {
- LOG(WARNING) << "fetchChunk request interrupted, "
- "streamChunkSlice: "
- << streamChunkSlice.toString();
- this->streamChunkSliceRegistry_.lock()->erase(streamChunkSlice);
- });
- return p.getFuture();
- });
- this->pipeline_->write(std::move(toSendMsg));
- CELEBORN_CHECK(!closed_);
- return f;
+
+ // Hold the state for the whole call: see read().
+ const auto state = state_;
+ auto future = state->registerFetch(streamChunkSlice);
+ if (!future) {
+ // The connection has been retired: fail retriably rather than asserting,
so
+ // CelebornInputStream can retry or fail over to a replica.
+ return folly::makeFuture<std::unique_ptr<Message>>(
+ makeRetriableTransportError(
+ __FILE__,
+ __LINE__,
+ __FUNCTION__,
+ fmt::format(
+ "connection closed before fetching streamChunkSlice {}",
+ streamChunkSlice.toString())));
+ }
+
+ // Write-failure handling: see operator().
+ const std::weak_ptr<ConnectionState> weakState = state;
+ auto written = writeToPipeline(
+ std::move(toSendMsg), [&]() { state->takeFetch(streamChunkSlice); });
+ std::move(written).thenError(
+ [weakState, streamChunkSlice](const folly::exception_wrapper& e) {
+ if (auto state = weakState.lock()) {
+ state->retire(fmt::format(
+ "Failed to send request for streamChunkSlice {}, errorMsg: {}",
+ streamChunkSlice.toString(),
+ e.what().toStdString()));
+ }
+ });
+ return std::move(*future);
}
void MessageDispatcher::sendRpcRequestWithoutResponse(
std::unique_ptr<Message> toSendMsg) {
CELEBORN_CHECK(toSendMsg->type() == Message::RPC_REQUEST);
- this->pipeline_->write(std::move(toSendMsg));
+ const long requestId =
+ reinterpret_cast<RpcRequest*>(toSendMsg.get())->requestId();
+
+ // Hold the state for the whole call: see read().
+ const auto state = state_;
+ if (state->retired()) {
+ // There is no promise to fail for this send -- the caller does not wait
for
+ // a response -- and no live socket to write to either.
+ LOG(WARNING) << "connection closed before sending requestId " << requestId
+ << " without response";
+ return;
+ }
+
+ // Unlike the paths above, this one cannot make the check and the send
atomic:
+ // there is no promise to register, so nothing to register it against. Losing
+ // the race is harmless -- the write then fails and retires the connection
+ // below -- it just means the send reached a socket already known to be dead.
+ //
+ // A failed write leaves nothing to fail here either, but it still means the
+ // connection is dead, so it has to be retired all the same: otherwise the
+ // client pool keeps handing it to the next caller.
+ const std::weak_ptr<ConnectionState> weakState = state;
+ folly::Future<folly::Unit> written = folly::makeFuture();
+ try {
+ written = this->pipeline_->write(std::move(toSendMsg));
+ } catch (const std::exception& e) {
+ // A handler rejected the message by throwing -- see writeToPipeline().
With
+ // no promise and no caller waiting, logging is the only way to report it,
+ // and it must not propagate: ~WorkerPartitionReader sends BufferStreamEnd
+ // this way, and an exception escaping a destructor aborts the process.
+ LOG(ERROR) << "failed to send requestId " << requestId
+ << " without response, errorMsg: " << e.what();
+ return;
+ } catch (...) {
+ // A handler is free to throw something that does not derive from
+ // std::exception, and that must not escape a destructor either.
+ LOG(ERROR) << "failed to send requestId " << requestId
+ << " without response, errorMsg: unknown exception";
+ return;
+ }
+ std::move(written).thenError(
+ [weakState, requestId](const folly::exception_wrapper& e) {
+ if (auto state = weakState.lock()) {
+ state->retire(fmt::format(
+ "Failed to send request {} without response, errorMsg: {}",
+ requestId,
+ e.what().toStdString()));
+ }
+ });
}
void MessageDispatcher::readEOF(Context* ctx) {
@@ -252,43 +570,23 @@ folly::Future<folly::Unit>
MessageDispatcher::writeException(
}
folly::Future<folly::Unit> MessageDispatcher::close() {
- if (!closed_) {
- closed_ = true;
- cleanup();
- }
- return ClientDispatcherBase::close();
+ // Hold the state for the whole call: see read().
+ const auto state = state_;
+ // Close the pipeline before failing what was outstanding on it: the order
+ // Java uses, and the only safe one here, since retiring fulfils promises
+ // inline and a continuation may drop the last reference to the
+ // TransportClient owning this dispatcher.
+ auto result = ClientDispatcherBase::close();
+ state->retire(kClientClosed);
+ return result;
}
folly::Future<folly::Unit> MessageDispatcher::close(Context* ctx) {
- if (!closed_) {
- closed_ = true;
- cleanup();
- }
-
- return ClientDispatcherBase::close(ctx);
-}
-
-void MessageDispatcher::cleanup() {
- LOG(WARNING) << "Cleaning up client!";
- requestIdRegistry_.withLock([&](auto& registry) {
- for (auto& [requestId, promiseHolder] : registry) {
- auto errorMsg =
- fmt::format("Client closed, cancel ongoing requestId {}", requestId);
- LOG(WARNING) << errorMsg;
- promiseHolder.msgPromise.setException(std::runtime_error(errorMsg));
- }
- registry.clear();
- });
- streamChunkSliceRegistry_.withLock([&](auto& registry) {
- for (auto& [streamChunkSlice, promiseHolder] : registry) {
- auto errorMsg = fmt::format(
- "Client closed, cancel ongoing streamChunkSlice {}",
- streamChunkSlice.toString());
- LOG(WARNING) << errorMsg;
- promiseHolder.msgPromise.setException(std::runtime_error(errorMsg));
- }
- registry.clear();
- });
+ // Ordering: see close().
+ const auto state = state_;
+ auto result = ClientDispatcherBase::close(ctx);
+ state->retire(kClientClosed);
+ return result;
}
} // namespace network
} // namespace celeborn
diff --git a/cpp/celeborn/network/MessageDispatcher.h
b/cpp/celeborn/network/MessageDispatcher.h
index 7ce6a7afe..73464d849 100644
--- a/cpp/celeborn/network/MessageDispatcher.h
+++ b/cpp/celeborn/network/MessageDispatcher.h
@@ -17,6 +17,9 @@
#pragma once
+#include <functional>
+#include <memory>
+
#include <wangle/bootstrap/ClientBootstrap.h>
#include <wangle/channel/AsyncSocketHandler.h>
#include <wangle/channel/EventBaseHandler.h>
@@ -53,6 +56,10 @@ class MessageDispatcher : public
wangle::ClientDispatcherBase<
std::unique_ptr<Message>,
std::unique_ptr<Message>> {
public:
+ MessageDispatcher();
+
+ ~MessageDispatcher() override;
+
void read(Context*, std::unique_ptr<Message> toRecvMsg) override;
virtual folly::Future<std::unique_ptr<Message>> sendRpcRequest(
@@ -89,28 +96,28 @@ class MessageDispatcher : public
wangle::ClientDispatcherBase<
folly::Future<folly::Unit> close(Context* ctx) override;
- bool isAvailable() override {
- return !closed_;
- }
+ bool isAvailable() override;
private:
- void cleanup();
-
- using MsgPromise = folly::Promise<std::unique_ptr<Message>>;
- struct MsgPromiseHolder {
- MsgPromise msgPromise;
- std::chrono::time_point<std::chrono::system_clock> requestTime;
- };
- folly::Synchronized<std::unordered_map<long, MsgPromiseHolder>, std::mutex>
- requestIdRegistry_;
- folly::Synchronized<
- std::unordered_map<
- protocol::StreamChunkSlice,
- MsgPromiseHolder,
- protocol::StreamChunkSlice::Hasher>,
- std::mutex>
- streamChunkSliceRegistry_;
- std::atomic<bool> closed_{false};
+ // The outstanding requests of the connection, and the flag that retires it.
+ // Defined in the .cpp, and refcounted separately from the dispatcher,
because
+ // the continuations the dispatcher attaches asynchronously can run after it
+ // is gone: a failed write is reported from the socket's write callback, and
+ // AsyncSocket fails whatever is still pending while it is torn down, by
which
+ // time TransportClient has destroyed its dispatcher. Keeping the dispatcher
+ // alive instead is not an option -- ~ClientDispatcherBase unregisters it
from
+ // the pipeline, so it must be destroyed while that pipeline still lives.
+ class ConnectionState;
+
+ // Writes a message down the pipeline and returns the future reporting the
+ // write's outcome. A handler may fail by throwing instead of by failing that
+ // future, in which case `onThrow` runs and the exception keeps propagating.
+ // See the call sites for when that happens.
+ folly::Future<folly::Unit> writeToPipeline(
+ std::unique_ptr<Message> toSendMsg,
+ const std::function<void()>& onThrow);
+
+ const std::shared_ptr<ConnectionState> state_;
};
} // namespace network
} // namespace celeborn
diff --git a/cpp/celeborn/network/TransportClient.cpp
b/cpp/celeborn/network/TransportClient.cpp
index c754eecd0..860d03758 100644
--- a/cpp/celeborn/network/TransportClient.cpp
+++ b/cpp/celeborn/network/TransportClient.cpp
@@ -22,6 +22,80 @@
namespace celeborn {
namespace network {
+namespace {
+// True when the cause was already classified as retriable by the transport.
+bool isRetriableCause(const std::exception& e) {
+ const auto* celebornException =
+ dynamic_cast<const utils::CelebornException*>(&e);
+ return celebornException != nullptr && celebornException->isRetriable();
+}
+
+// Rebuilds the failure handed to a push/fetch callback. The transport marks a
+// recoverable failure -- a closed connection, a failed write, a connect
+// failure -- with a retriable CelebornException; flattening it into a plain
+// std::runtime_error here would drop that classification before any caller can
+// observe it, so the retriable exception is forwarded as-is.
+//
+// Matching CelebornRuntimeError alone is enough: the CELEBORN_CHECK /
+// CELEBORN_FAIL macros all hardcode isRetriable=false, so a retriable cause is
+// always one of the CelebornRuntimeErrors the transport builds explicitly.
+std::unique_ptr<std::exception> toCallbackException(
+ const folly::exception_wrapper& e) {
+ std::unique_ptr<std::exception> retriableFailure;
+ e.with_exception([&](const utils::CelebornRuntimeError& error) {
+ if (error.isRetriable()) {
+ retriableFailure = std::make_unique<utils::CelebornRuntimeError>(error);
+ }
+ });
+ if (retriableFailure) {
+ return retriableFailure;
+ }
+ return std::make_unique<std::runtime_error>(e.what().toStdString());
+}
+
+// The counterpart of toCallbackException for a synchronously thrown cause:
+// wraps `errorMsg` while keeping the cause's retriable classification.
+std::unique_ptr<std::exception> wrapCallbackException(
+ const char* file,
+ size_t line,
+ const char* function,
+ const std::string& errorMsg,
+ const std::exception& cause) {
+ if (!isRetriableCause(cause)) {
+ return std::make_unique<std::runtime_error>(errorMsg);
+ }
+ return std::make_unique<utils::CelebornRuntimeError>(
+ file,
+ line,
+ function,
+ /*expression=*/"",
+ /*message=*/errorMsg,
+ utils::error_source::kErrorSourceRuntime.c_str(),
+ utils::error_code::kInvalidState.c_str(),
+ /*isRetriable=*/true);
+}
+
+// Rethrows `errorMsg` preserving the cause's retriable classification.
+// CELEBORN_FAIL would hardcode isRetriable=false and hide a recoverable
+// transport failure from the retry/failover paths.
+[[noreturn]] void failPreservingRetriable(
+ const char* file,
+ size_t line,
+ const char* function,
+ const std::string& errorMsg,
+ const std::exception& cause) {
+ throw utils::CelebornRuntimeError(
+ file,
+ line,
+ function,
+ /*expression=*/"",
+ /*message=*/errorMsg,
+ utils::error_source::kErrorSourceRuntime.c_str(),
+ utils::error_code::kInvalidState.c_str(),
+ /*isRetriable=*/isRetriableCause(cause));
+}
+} // namespace
+
void MessageSerializeHandler::read(
Context* ctx,
std::unique_ptr<folly::IOBuf> msg) {
@@ -61,7 +135,7 @@ RpcResponse TransportClient::sendRpcRequestSync(
timeout,
folly::exceptionStr(e).toStdString());
LOG(ERROR) << errorMsg;
- CELEBORN_FAIL(errorMsg);
+ failPreservingRetriable(__FILE__, __LINE__, __FUNCTION__, errorMsg, e);
}
}
@@ -100,8 +174,7 @@ void TransportClient::pushDataAsync(
}
})
.thenError([_callback = callback](const folly::exception_wrapper& e) {
- _callback->onFailure(
- std::make_unique<std::runtime_error>(e.what().toStdString()));
+ _callback->onFailure(toCallbackException(e));
});
} catch (std::exception& e) {
@@ -112,7 +185,8 @@ void TransportClient::pushDataAsync(
pushData.mode(),
e.what());
LOG(ERROR) << errorMsg;
- callback->onFailure(std::make_unique<std::runtime_error>(errorMsg));
+ callback->onFailure(
+ wrapCallbackException(__FILE__, __LINE__, __FUNCTION__, errorMsg, e));
}
}
@@ -137,8 +211,7 @@ void TransportClient::pushMergedDataAsync(
}
})
.thenError([_callback = callback](const folly::exception_wrapper& e) {
- _callback->onFailure(
- std::make_unique<std::runtime_error>(e.what().toStdString()));
+ _callback->onFailure(toCallbackException(e));
});
} catch (std::exception& e) {
@@ -148,7 +221,8 @@ void TransportClient::pushMergedDataAsync(
pushMergedData.mode(),
e.what());
LOG(ERROR) << errorMsg;
- callback->onFailure(std::make_unique<std::runtime_error>(errorMsg));
+ callback->onFailure(
+ wrapCallbackException(__FILE__, __LINE__, __FUNCTION__, errorMsg, e));
}
}
@@ -178,12 +252,12 @@ void TransportClient::fetchChunkAsync(
})
.thenError(
[=, _onFailure = onFailure](const folly::exception_wrapper& e) {
- _onFailure(
- streamChunkSlice,
-
std::make_unique<std::runtime_error>(e.what().toStdString()));
+ _onFailure(streamChunkSlice, toCallbackException(e));
});
} catch (std::exception& e) {
- CELEBORN_FAIL(e.what());
+ LOG(ERROR) << "fetchChunk failed. streamChunkSlice: "
+ << streamChunkSlice.toString() << ", errorMsg: " << e.what();
+ failPreservingRetriable(__FILE__, __LINE__, __FUNCTION__, e.what(), e);
}
}
@@ -264,7 +338,20 @@ std::shared_ptr<TransportClient>
TransportClientFactory::createClient(
connectTimeout_,
folly::exceptionStr(e).toStdString());
LOG(ERROR) << errorMsg;
- CELEBORN_FAIL(errorMsg);
+ // Failing to establish a connection is transient: the peer may be
+ // restarting, or the network may be briefly unavailable. Classify it as
+ // retriable rather than as an invariant violation, so that the
+ // createReaderWithRetry and push failover paths can tell it apart from a
+ // terminal failure.
+ throw utils::CelebornRuntimeError(
+ __FILE__,
+ __LINE__,
+ __FUNCTION__,
+ /*expression=*/"",
+ /*message=*/errorMsg,
+ utils::error_source::kErrorSourceRuntime.c_str(),
+ utils::error_code::kInvalidState.c_str(),
+ /*isRetriable=*/true);
}
}
}
diff --git a/cpp/celeborn/network/tests/MessageDispatcherTest.cpp
b/cpp/celeborn/network/tests/MessageDispatcherTest.cpp
index 959aca6e3..8cdf80cff 100644
--- a/cpp/celeborn/network/tests/MessageDispatcherTest.cpp
+++ b/cpp/celeborn/network/tests/MessageDispatcherTest.cpp
@@ -19,11 +19,20 @@
#include "celeborn/network/FrameDecoder.h"
#include "celeborn/network/MessageDispatcher.h"
+#include "celeborn/network/TransportClient.h"
using namespace celeborn;
using namespace celeborn::network;
namespace {
+// How write() reports a failure: through the returned future, the way wangle's
+// AsyncSocketHandler does, or by throwing, the way a handler that serializes
+// the message does -- MessageSerializeHandler encodes it, and Message::encode
+// checks its own invariants.
+// kThrowNonStd throws something that does not derive from std::exception,
which
+// a handler is free to do.
+enum class WriteFailure { kFailFuture, kThrow, kThrowNonStd };
+
class MockHandler : public wangle::Handler<
std::unique_ptr<folly::IOBuf>,
std::unique_ptr<Message>,
@@ -32,16 +41,64 @@ class MockHandler : public wangle::Handler<
public:
MockHandler(std::unique_ptr<Message>& writedMsg) : writedMsg_(writedMsg) {}
+ // When writeError is set, write() reports the failure through the returned
+ // future, the way wangle's AsyncSocketHandler does for a socket that is no
+ // longer good or whose write callback fails. The writes before
+ // failFromWrite succeed, so that a test can leave an earlier request in
+ // flight on the connection that the failing write kills.
+ MockHandler(
+ std::unique_ptr<Message>& writedMsg,
+ std::string writeError,
+ int failFromWrite = 1,
+ WriteFailure writeFailure = WriteFailure::kFailFuture)
+ : writedMsg_(writedMsg),
+ writeError_(std::move(writeError)),
+ failFromWrite_(failFromWrite),
+ writeFailure_(writeFailure) {}
+
+ // Hands the write future back to the test, which completes it whenever it
+ // wants -- AsyncSocketHandler reports a failed write from the socket's write
+ // callback, so it can arrive long after write() returned, including while
the
+ // connection is being torn down.
+ MockHandler(
+ std::unique_ptr<Message>& writedMsg,
+ folly::Promise<folly::Unit>& writePromise)
+ : writedMsg_(writedMsg), writePromise_(&writePromise) {}
+
void read(Context* ctx, std::unique_ptr<folly::IOBuf> msg) override {}
folly::Future<folly::Unit> write(Context* ctx, std::unique_ptr<Message> msg)
override {
writedMsg_ = std::move(msg);
+ ++numWrites_;
+ if (writePromise_ != nullptr) {
+ return writePromise_->getFuture();
+ }
+ if (writeFailure_ == WriteFailure::kThrow ||
+ writeFailure_ == WriteFailure::kThrowNonStd) {
+ // Only the one malformed message throws: what a handler rejects is the
+ // message, not the connection, so the writes around it go through.
+ if (numWrites_ == failFromWrite_) {
+ if (writeFailure_ == WriteFailure::kThrowNonStd) {
+ throw writeError_;
+ }
+ throw std::runtime_error(writeError_);
+ }
+ return {};
+ }
+ if (!writeError_.empty() && numWrites_ >= failFromWrite_) {
+ return folly::makeFuture<folly::Unit>(std::runtime_error(writeError_));
+ }
return {};
}
private:
std::unique_ptr<Message>& writedMsg_;
+ const std::string writeError_;
+ const int failFromWrite_{1};
+ const WriteFailure writeFailure_{WriteFailure::kFailFuture};
+ folly::Promise<folly::Unit>* writePromise_{nullptr};
+ int numWrites_{0};
};
SerializePipeline::Ptr createMockedPipeline(MockHandler&& mockHandler) {
@@ -66,6 +123,15 @@ std::string takeExceptionMessage(
return std::move(future).result().exception().what().toStdString();
}
+// Returns true when the future failed with a CelebornException marked
+// retriable.
+bool failedRetriably(folly::Future<std::unique_ptr<Message>>&& future) {
+ bool retriable = false;
+ const bool matched = std::move(future).result().exception().with_exception(
+ [&](const utils::CelebornException& e) { retriable = e.isRetriable(); });
+ return matched && retriable;
+}
+
} // namespace
TEST(MessageDispatcherTest, sendRpcRequestAndReceiveResponse) {
@@ -304,6 +370,396 @@ TEST(MessageDispatcherTest,
sendFetchChunkRequestAndReceiveFailure) {
EXPECT_NE(exceptionMsg.find(streamChunkSlice.toString()), std::string::npos);
}
+// A send issued after the connection is closed must fail gracefully with a
+// ready, retriable exception instead of tripping an assertion. This mirrors
the
+// Java client, where a send on an inactive channel surfaces as a retriable
+// IOException that CelebornInputStream retries.
+TEST(MessageDispatcherTest, sendRpcRequestAfterCloseFailsRetriably) {
+ std::unique_ptr<Message> sentMsg;
+ MockHandler mockHandler(sentMsg);
+ auto mockPipeline = createMockedPipeline(std::move(mockHandler));
+ auto dispatcher = std::make_unique<MessageDispatcher>();
+ dispatcher->setPipeline(mockPipeline.get());
+
+ dispatcher->close();
+ EXPECT_FALSE(dispatcher->isAvailable());
+
+ const long requestId = 2001;
+ const std::string requestBody = "test-request-body";
+ auto rpcRequest = std::make_unique<RpcRequest>(
+ requestId, toReadOnlyByteBuffer(requestBody));
+ auto future = dispatcher->sendRpcRequest(std::move(rpcRequest));
+
+ ASSERT_TRUE(future.isReady());
+ ASSERT_TRUE(future.hasException());
+ EXPECT_TRUE(failedRetriably(std::move(future)));
+}
+
+TEST(MessageDispatcherTest, sendFetchChunkRequestAfterCloseFailsRetriably) {
+ std::unique_ptr<Message> sentMsg;
+ MockHandler mockHandler(sentMsg);
+ auto mockPipeline = createMockedPipeline(std::move(mockHandler));
+ auto dispatcher = std::make_unique<MessageDispatcher>();
+ dispatcher->setPipeline(mockPipeline.get());
+
+ dispatcher->close();
+ EXPECT_FALSE(dispatcher->isAvailable());
+
+ const protocol::StreamChunkSlice streamChunkSlice{2001, 2002, 2003, 2004};
+ const long requestId = 2001;
+ const std::string requestBody = "test-request-body";
+ auto rpcRequest = std::make_unique<RpcRequest>(
+ requestId, toReadOnlyByteBuffer(requestBody));
+ auto future = dispatcher->sendFetchChunkRequest(
+ streamChunkSlice, std::move(rpcRequest));
+
+ ASSERT_TRUE(future.isReady());
+ ASSERT_TRUE(future.hasException());
+ EXPECT_TRUE(failedRetriably(std::move(future)));
+}
+
+// close() must fail any in-flight request rather than leaving its future
+// pending forever, matching Java's failOutstandingRequests on channelInactive.
+TEST(MessageDispatcherTest, closeFailsInFlightRequestsRetriably) {
+ 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 = 3001;
+ 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());
+
+ dispatcher->close();
+
+ ASSERT_TRUE(future.isReady());
+ ASSERT_TRUE(future.hasException());
+ EXPECT_TRUE(failedRetriably(std::move(future)));
+}
+
+// A failed write must fail the registered request instead of leaving its
future
+// pending until the request timeout, and it must retire the connection.
+// wangle's AsyncSocketHandler reports such a failure through the write future
+// -- immediately when the socket is no longer good, or later from its write
+// callback -- without going through transportInactive first, so the dispatcher
+// is not closed at that point: TransportClient::active() would keep reporting
+// true and TransportClientFactory would hand the same dead connection to every
+// retry. Java's StdChannelListener closes the channel before reporting the
+// failure, and closing it fails whatever else was outstanding.
+TEST(MessageDispatcherTest, sendRpcRequestFailedWriteRetiresConnection) {
+ std::unique_ptr<Message> sentMsg;
+ // The first write succeeds and leaves its request in flight; the second one
+ // fails.
+ MockHandler mockHandler(
+ sentMsg, "socket is closed in write()", /*failFromWrite=*/2);
+ auto mockPipeline = createMockedPipeline(std::move(mockHandler));
+ auto dispatcher = std::make_unique<MessageDispatcher>();
+ dispatcher->setPipeline(mockPipeline.get());
+
+ const std::string requestBody = "test-request-body";
+ auto inFlight = dispatcher->sendRpcRequest(std::make_unique<RpcRequest>(
+ /*requestId=*/4001, toReadOnlyByteBuffer(requestBody)));
+ EXPECT_FALSE(inFlight.isReady());
+
+ auto future = dispatcher->sendRpcRequest(std::make_unique<RpcRequest>(
+ /*requestId=*/4002, toReadOnlyByteBuffer(requestBody)));
+
+ ASSERT_TRUE(future.isReady());
+ ASSERT_TRUE(future.hasException());
+ EXPECT_TRUE(failedRetriably(std::move(future)));
+ // The connection is retired, so the client pool stops handing it out.
+ EXPECT_FALSE(dispatcher->isAvailable());
+ // And the request that was still outstanding on it is failed too, rather
than
+ // waiting for its timeout on a dead connection.
+ ASSERT_TRUE(inFlight.isReady());
+ ASSERT_TRUE(inFlight.hasException());
+ EXPECT_TRUE(failedRetriably(std::move(inFlight)));
+ // A caller still holding the retired connection fails fast on it instead of
+ // writing to a dead socket.
+ auto rejected = dispatcher->sendRpcRequest(std::make_unique<RpcRequest>(
+ /*requestId=*/4003, toReadOnlyByteBuffer(requestBody)));
+ ASSERT_TRUE(rejected.isReady());
+ EXPECT_TRUE(failedRetriably(std::move(rejected)));
+}
+
+// The hop from a retired connection to the client pool:
TransportClient::active
+// reports the dispatcher's availability, and TransportClientFactory only
reuses
+// a cached client while that is true, so a failed write must make the client
+// report itself inactive.
+TEST(MessageDispatcherTest, failedWriteMakesTransportClientInactive) {
+ std::unique_ptr<Message> sentMsg;
+ MockHandler mockHandler(sentMsg, "socket is closed in write()");
+ auto mockPipeline = createMockedPipeline(std::move(mockHandler));
+ auto dispatcher = std::make_unique<MessageDispatcher>();
+ dispatcher->setPipeline(mockPipeline.get());
+ auto* rawDispatcher = dispatcher.get();
+ TransportClient client(
+ /*client=*/nullptr, std::move(dispatcher), Timeout(10000));
+ EXPECT_TRUE(client.active());
+
+ auto future = rawDispatcher->sendRpcRequest(std::make_unique<RpcRequest>(
+ /*requestId=*/4101, toReadOnlyByteBuffer("test-request-body")));
+
+ ASSERT_TRUE(future.isReady());
+ EXPECT_TRUE(failedRetriably(std::move(future)));
+ EXPECT_FALSE(client.active());
+}
+
+TEST(MessageDispatcherTest, sendFetchChunkRequestFailedWriteRetiresConnection)
{
+ std::unique_ptr<Message> sentMsg;
+ MockHandler mockHandler(
+ sentMsg, "socket is closed in write()", /*failFromWrite=*/2);
+ auto mockPipeline = createMockedPipeline(std::move(mockHandler));
+ auto dispatcher = std::make_unique<MessageDispatcher>();
+ dispatcher->setPipeline(mockPipeline.get());
+
+ const std::string requestBody = "test-request-body";
+ const protocol::StreamChunkSlice inFlightSlice{4001, 4002, 4003, 4004};
+ auto inFlight = dispatcher->sendFetchChunkRequest(
+ inFlightSlice,
+ std::make_unique<RpcRequest>(
+ /*requestId=*/4001, toReadOnlyByteBuffer(requestBody)));
+ EXPECT_FALSE(inFlight.isReady());
+
+ const protocol::StreamChunkSlice streamChunkSlice{4002, 4002, 4003, 4004};
+ auto future = dispatcher->sendFetchChunkRequest(
+ streamChunkSlice,
+ std::make_unique<RpcRequest>(
+ /*requestId=*/4002, toReadOnlyByteBuffer(requestBody)));
+
+ ASSERT_TRUE(future.isReady());
+ ASSERT_TRUE(future.hasException());
+ EXPECT_TRUE(failedRetriably(std::move(future)));
+ EXPECT_FALSE(dispatcher->isAvailable());
+ ASSERT_TRUE(inFlight.isReady());
+ ASSERT_TRUE(inFlight.hasException());
+ EXPECT_TRUE(failedRetriably(std::move(inFlight)));
+ const protocol::StreamChunkSlice rejectedSlice{4003, 4002, 4003, 4004};
+ auto rejected = dispatcher->sendFetchChunkRequest(
+ rejectedSlice,
+ std::make_unique<RpcRequest>(
+ /*requestId=*/4003, toReadOnlyByteBuffer(requestBody)));
+ ASSERT_TRUE(rejected.isReady());
+ EXPECT_TRUE(failedRetriably(std::move(rejected)));
+}
+
+// The write failure may be reported after the dispatcher is destroyed: it
comes
+// from the socket's write callback, and AsyncSocket fails whatever is still
+// pending when it is torn down. TransportClient destroys its dispatcher before
+// the bootstrap that owns the pipeline -- and it has to, because
+// ~ClientDispatcherBase unregisters itself from that pipeline -- so the write
+// continuation must not depend on the dispatcher being alive.
+TEST(MessageDispatcherTest, failedWriteAfterDispatcherDestroyedIsIgnored) {
+ std::unique_ptr<Message> sentMsg;
+ folly::Promise<folly::Unit> writePromise;
+ MockHandler mockHandler(sentMsg, writePromise);
+ auto mockPipeline = createMockedPipeline(std::move(mockHandler));
+ auto dispatcher = std::make_unique<MessageDispatcher>();
+ dispatcher->setPipeline(mockPipeline.get());
+
+ auto future = dispatcher->sendRpcRequest(std::make_unique<RpcRequest>(
+ /*requestId=*/5001, toReadOnlyByteBuffer("test-request-body")));
+ EXPECT_FALSE(future.isReady());
+
+ dispatcher.reset();
+ // The destroyed dispatcher failed the request it still had outstanding, with
+ // the same retriable error it reports on close() rather than with folly's
+ // BrokenPromise, which carries no cause for the caller to classify.
+ ASSERT_TRUE(future.isReady());
+ ASSERT_TRUE(future.hasException());
+ EXPECT_TRUE(failedRetriably(std::move(future)));
+
+ // The write fails only now, with no dispatcher left to report it to. The
+ // continuation must be a no-op instead of reaching into freed memory.
+ writePromise.setException(
+ std::runtime_error("socket is closed during teardown"));
+}
+
+TEST(MessageDispatcherTest, failedFetchWriteAfterDispatcherDestroyedIsIgnored)
{
+ std::unique_ptr<Message> sentMsg;
+ folly::Promise<folly::Unit> writePromise;
+ MockHandler mockHandler(sentMsg, writePromise);
+ auto mockPipeline = createMockedPipeline(std::move(mockHandler));
+ auto dispatcher = std::make_unique<MessageDispatcher>();
+ dispatcher->setPipeline(mockPipeline.get());
+
+ const protocol::StreamChunkSlice streamChunkSlice{5001, 5002, 5003, 5004};
+ auto future = dispatcher->sendFetchChunkRequest(
+ streamChunkSlice,
+ std::make_unique<RpcRequest>(
+ /*requestId=*/5001, toReadOnlyByteBuffer("test-request-body")));
+ EXPECT_FALSE(future.isReady());
+
+ dispatcher.reset();
+ ASSERT_TRUE(future.isReady());
+ ASSERT_TRUE(future.hasException());
+ EXPECT_TRUE(failedRetriably(std::move(future)));
+
+ writePromise.setException(
+ std::runtime_error("socket is closed during teardown"));
+}
+
+// A handler may also fail by throwing rather than by failing the write future:
+// the message is serialized on the way down the pipeline, by
+// MessageSerializeHandler, and wangle::Pipeline::write has no try/catch of its
+// own. The connection itself is fine in that case, so the exception keeps
+// propagating to the caller -- a violation of our own encoding invariants is
+// not retriable -- but the request must not be left registered on the
+// connection, since nothing was sent and no response will ever arrive for it.
+TEST(MessageDispatcherTest, throwingWriteUnregistersTheRequest) {
+ std::unique_ptr<Message> sentMsg;
+ MockHandler mockHandler(
+ sentMsg,
+ "encoded length mismatch",
+ /*failFromWrite=*/2,
+ WriteFailure::kThrow);
+ auto mockPipeline = createMockedPipeline(std::move(mockHandler));
+ auto dispatcher = std::make_unique<MessageDispatcher>();
+ dispatcher->setPipeline(mockPipeline.get());
+
+ const std::string requestBody = "test-request-body";
+ auto inFlight = dispatcher->sendRpcRequest(std::make_unique<RpcRequest>(
+ /*requestId=*/7001, toReadOnlyByteBuffer(requestBody)));
+ EXPECT_FALSE(inFlight.isReady());
+
+ const long requestId = 7002;
+ EXPECT_THROW(
+ dispatcher->sendRpcRequest(std::make_unique<RpcRequest>(
+ requestId, toReadOnlyByteBuffer(requestBody))),
+ std::runtime_error);
+
+ // The connection is untouched: it is still usable, and what was outstanding
+ // on it is still outstanding.
+ EXPECT_TRUE(dispatcher->isAvailable());
+ EXPECT_FALSE(inFlight.isReady());
+ // And the request whose write threw is gone from the registry. Registering
+ // the same id again would otherwise find the leftover entry, whose future
has
+ // already been handed out.
+ auto retried = dispatcher->sendRpcRequest(std::make_unique<RpcRequest>(
+ requestId, toReadOnlyByteBuffer(requestBody)));
+ EXPECT_FALSE(retried.isReady());
+}
+
+TEST(MessageDispatcherTest, throwingFetchWriteUnregistersTheRequest) {
+ std::unique_ptr<Message> sentMsg;
+ MockHandler mockHandler(
+ sentMsg,
+ "encoded length mismatch",
+ /*failFromWrite=*/1,
+ WriteFailure::kThrow);
+ auto mockPipeline = createMockedPipeline(std::move(mockHandler));
+ auto dispatcher = std::make_unique<MessageDispatcher>();
+ dispatcher->setPipeline(mockPipeline.get());
+
+ const std::string requestBody = "test-request-body";
+ const protocol::StreamChunkSlice streamChunkSlice{7003, 7004, 7005, 7006};
+ EXPECT_THROW(
+ dispatcher->sendFetchChunkRequest(
+ streamChunkSlice,
+ std::make_unique<RpcRequest>(
+ /*requestId=*/7003, toReadOnlyByteBuffer(requestBody))),
+ std::runtime_error);
+
+ EXPECT_TRUE(dispatcher->isAvailable());
+ auto retried = dispatcher->sendFetchChunkRequest(
+ streamChunkSlice,
+ std::make_unique<RpcRequest>(
+ /*requestId=*/7003, toReadOnlyByteBuffer(requestBody)));
+ EXPECT_FALSE(retried.isReady());
+}
+
+// A send that expects no response has no promise to fail, but a failed write
+// still means the connection is dead: it must be retired, or the client pool
+// keeps handing it to the next caller. ~WorkerPartitionReader takes this path
+// to send BufferStreamEnd.
+TEST(MessageDispatcherTest, failedWriteWithoutResponseRetiresConnection) {
+ std::unique_ptr<Message> sentMsg;
+ // The first write succeeds and leaves its request in flight; the second one
+ // fails.
+ MockHandler mockHandler(
+ sentMsg, "socket is closed in write()", /*failFromWrite=*/2);
+ auto mockPipeline = createMockedPipeline(std::move(mockHandler));
+ auto dispatcher = std::make_unique<MessageDispatcher>();
+ dispatcher->setPipeline(mockPipeline.get());
+
+ const std::string requestBody = "test-request-body";
+ auto inFlight = dispatcher->sendRpcRequest(std::make_unique<RpcRequest>(
+ /*requestId=*/6001, toReadOnlyByteBuffer(requestBody)));
+ EXPECT_FALSE(inFlight.isReady());
+
+ dispatcher->sendRpcRequestWithoutResponse(std::make_unique<RpcRequest>(
+ /*requestId=*/6002, toReadOnlyByteBuffer(requestBody)));
+
+ EXPECT_FALSE(dispatcher->isAvailable());
+ ASSERT_TRUE(inFlight.isReady());
+ ASSERT_TRUE(inFlight.hasException());
+ EXPECT_TRUE(failedRetriably(std::move(inFlight)));
+}
+
+// A handler that rejects the message by throwing must not propagate out of a
+// send that expects no response either: ~WorkerPartitionReader sends
+// BufferStreamEnd this way, and an exception escaping a destructor aborts the
+// process.
+TEST(MessageDispatcherTest, throwingWriteWithoutResponseIsReported) {
+ std::unique_ptr<Message> sentMsg;
+ MockHandler mockHandler(
+ sentMsg,
+ "encoded length mismatch",
+ /*failFromWrite=*/1,
+ WriteFailure::kThrow);
+ auto mockPipeline = createMockedPipeline(std::move(mockHandler));
+ auto dispatcher = std::make_unique<MessageDispatcher>();
+ dispatcher->setPipeline(mockPipeline.get());
+
+ EXPECT_NO_THROW(
+ dispatcher->sendRpcRequestWithoutResponse(std::make_unique<RpcRequest>(
+ /*requestId=*/7007, toReadOnlyByteBuffer("test-request-body"))));
+
+ // The message was rejected, not the connection.
+ EXPECT_TRUE(dispatcher->isAvailable());
+}
+
+// Including when what it throws does not derive from std::exception.
+TEST(MessageDispatcherTest, nonStdThrowingWriteWithoutResponseIsReported) {
+ std::unique_ptr<Message> sentMsg;
+ MockHandler mockHandler(
+ sentMsg,
+ "encoded length mismatch",
+ /*failFromWrite=*/1,
+ WriteFailure::kThrowNonStd);
+ auto mockPipeline = createMockedPipeline(std::move(mockHandler));
+ auto dispatcher = std::make_unique<MessageDispatcher>();
+ dispatcher->setPipeline(mockPipeline.get());
+
+ EXPECT_NO_THROW(
+ dispatcher->sendRpcRequestWithoutResponse(std::make_unique<RpcRequest>(
+ /*requestId=*/7008, toReadOnlyByteBuffer("test-request-body"))));
+
+ EXPECT_TRUE(dispatcher->isAvailable());
+}
+
+// And once the connection is retired, such a send is skipped rather than
+// written to a socket that is known to be dead.
+TEST(MessageDispatcherTest, sendWithoutResponseAfterCloseIsSkipped) {
+ std::unique_ptr<Message> sentMsg;
+ MockHandler mockHandler(sentMsg);
+ auto mockPipeline = createMockedPipeline(std::move(mockHandler));
+ auto dispatcher = std::make_unique<MessageDispatcher>();
+ dispatcher->setPipeline(mockPipeline.get());
+
+ dispatcher->close();
+ ASSERT_FALSE(dispatcher->isAvailable());
+
+ dispatcher->sendRpcRequestWithoutResponse(std::make_unique<RpcRequest>(
+ /*requestId=*/6003, toReadOnlyByteBuffer("test-request-body")));
+
+ EXPECT_EQ(sentMsg, nullptr);
+}
+
TEST(MessageDispatcherTest, heartbeatIsSilentlyConsumed) {
std::unique_ptr<Message> sentMsg;
MockHandler mockHandler(sentMsg);
diff --git a/cpp/celeborn/network/tests/TransportClientTest.cpp
b/cpp/celeborn/network/tests/TransportClientTest.cpp
index 0135d61e5..df432f498 100644
--- a/cpp/celeborn/network/tests/TransportClientTest.cpp
+++ b/cpp/celeborn/network/tests/TransportClientTest.cpp
@@ -25,11 +25,38 @@ using namespace celeborn::network;
namespace {
using MS = std::chrono::milliseconds;
+
+// The retriable error MessageDispatcher reports when the connection is closed
+// or the write fails.
+folly::exception_wrapper retriableConnectionClosed() {
+ return folly::make_exception_wrapper<utils::CelebornRuntimeError>(
+ __FILE__,
+ static_cast<size_t>(__LINE__),
+ __FUNCTION__,
+ /*expression=*/"",
+ /*message=*/"connection closed",
+ utils::error_source::kErrorSourceRuntime.c_str(),
+ utils::error_code::kInvalidState.c_str(),
+ /*isRetriable=*/true);
+}
+
+// True when the exception handed to the caller kept the retriable
+// classification the dispatcher attached to the failure.
+bool failedRetriably(const std::exception* exception) {
+ const auto* celebornException =
+ dynamic_cast<const utils::CelebornException*>(exception);
+ return celebornException != nullptr && celebornException->isRetriable();
+}
+
class MockDispatcher : public MessageDispatcher {
public:
folly::Future<std::unique_ptr<Message>> sendRpcRequest(
std::unique_ptr<Message> toSendMsg) override {
sentMsg_ = std::move(toSendMsg);
+ if (connectionClosed_) {
+ return folly::makeFuture<std::unique_ptr<Message>>(
+ retriableConnectionClosed());
+ }
msgPromise_ = MsgPromise();
return msgPromise_.getFuture();
}
@@ -42,6 +69,10 @@ class MockDispatcher : public MessageDispatcher {
folly::Future<std::unique_ptr<Message>> sendPushDataRequest(
std::unique_ptr<Message> toSendMsg) override {
sentMsg_ = std::move(toSendMsg);
+ if (connectionClosed_) {
+ return folly::makeFuture<std::unique_ptr<Message>>(
+ retriableConnectionClosed());
+ }
msgPromise_ = MsgPromise();
return msgPromise_.getFuture();
}
@@ -50,6 +81,10 @@ class MockDispatcher : public MessageDispatcher {
const protocol::StreamChunkSlice& streamChunkSlice,
std::unique_ptr<Message> toSendMsg) override {
sentMsg_ = std::move(toSendMsg);
+ if (connectionClosed_) {
+ return folly::makeFuture<std::unique_ptr<Message>>(
+ retriableConnectionClosed());
+ }
msgPromise_ = MsgPromise();
return msgPromise_.getFuture();
}
@@ -62,10 +97,17 @@ class MockDispatcher : public MessageDispatcher {
msgPromise_.setValue(std::move(msg));
}
+ // Makes every send fail retriably, as MessageDispatcher does once the
+ // connection is closed.
+ void setConnectionClosed() {
+ connectionClosed_ = true;
+ }
+
private:
using MsgPromise = folly::Promise<std::unique_ptr<Message>>;
std::unique_ptr<Message> sentMsg_;
MsgPromise msgPromise_;
+ bool connectionClosed_{false};
};
std::unique_ptr<memory::ReadOnlyByteBuffer> toReadOnlyByteBuffer(
@@ -209,6 +251,78 @@ TEST_F(TransportClientTest, sendRpcRequestSyncTimeout) {
EXPECT_TRUE(timeoutHappened);
}
+// The retriable classification the dispatcher attaches to a recoverable
+// transport failure must survive TransportClient's public API, otherwise no
+// caller can tell such a failure apart from a terminal one.
+TEST_F(TransportClientTest, sendRpcRequestSyncPreservesRetriableFailure) {
+ auto mockDispatcher = std::make_unique<MockDispatcher>();
+ mockDispatcher->setConnectionClosed();
+ TransportClient client(nullptr, std::move(mockDispatcher), MS(10000));
+
+ const long requestId = 1001;
+ auto rpcRequest = std::make_unique<RpcRequest>(
+ requestId, toReadOnlyByteBuffer("test-request-body"));
+
+ bool failed = false;
+ bool retriable = false;
+ try {
+ client.sendRpcRequestSync(*rpcRequest, MS(10000));
+ } catch (const std::exception& e) {
+ failed = true;
+ retriable = failedRetriably(&e);
+ }
+ EXPECT_TRUE(failed);
+ EXPECT_TRUE(retriable);
+}
+
+TEST_F(TransportClientTest, pushDataAsyncPreservesRetriableFailure) {
+ auto mockDispatcher = std::make_unique<MockDispatcher>();
+ mockDispatcher->setConnectionClosed();
+ TransportClient client(nullptr, std::move(mockDispatcher), MS(10000));
+ auto mockRpcResponseCallback = std::make_shared<MockRpcResponseCallback>();
+
+ auto pushData = std::make_unique<PushData>(
+ /*requestId=*/1001,
+ /*mode=*/2,
+ "test-shuffle-key",
+ "test-partition-id",
+ toReadOnlyByteBuffer("test-request-body"));
+ client.pushDataAsync(*pushData, MS(10000), mockRpcResponseCallback);
+
+ auto onFailureException = mockRpcResponseCallback->getOnFailureException();
+ EXPECT_FALSE(mockRpcResponseCallback->getOnSuccessBuffer());
+ ASSERT_TRUE(onFailureException);
+ EXPECT_TRUE(failedRetriably(onFailureException.get()));
+}
+
+TEST_F(TransportClientTest, fetchChunkAsyncPreservesRetriableFailure) {
+ auto mockDispatcher = std::make_unique<MockDispatcher>();
+ mockDispatcher->setConnectionClosed();
+ TransportClient client(nullptr, std::move(mockDispatcher), MS(10000));
+
+ const protocol::StreamChunkSlice streamChunkSlice{1, 2, 3, 4};
+ auto rpcRequest = std::make_unique<RpcRequest>(
+ /*requestId=*/1001, toReadOnlyByteBuffer("test-request-body"));
+ std::unique_ptr<memory::ReadOnlyByteBuffer> onSuccessBuffer;
+ FetchChunkSuccessCallback onSuccess =
+ [&](protocol::StreamChunkSlice slice,
+ std::unique_ptr<memory::ReadOnlyByteBuffer> buffer) {
+ onSuccessBuffer = std::move(buffer);
+ };
+ std::unique_ptr<std::exception> onFailureException;
+ FetchChunkFailureCallback onFailure =
+ [&](protocol::StreamChunkSlice slice,
+ std::unique_ptr<std::exception> exception) {
+ onFailureException = std::move(exception);
+ };
+
+ client.fetchChunkAsync(streamChunkSlice, *rpcRequest, onSuccess, onFailure);
+
+ EXPECT_FALSE(onSuccessBuffer);
+ ASSERT_TRUE(onFailureException);
+ EXPECT_TRUE(failedRetriably(onFailureException.get()));
+}
+
TEST_F(TransportClientTest, sendRpcRequestWithoutResponse) {
auto mockDispatcher = std::make_unique<MockDispatcher>();
auto rawMockDispatcher = mockDispatcher.get();
