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 9d3a5b50b [CELEBORN-2413] Fail transport requests retriably when the 
C++ connection is closed or a write fails
9d3a5b50b is described below

commit 9d3a5b50bd8a7dce4f3445b0a0d54d1187f9f5b1
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]>
---
 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();

Reply via email to