[CELEBORN-2413] Fail retriably instead of asserting when the C++ transport connection is closed

[yugan95]

What changes were proposed in this pull request?

When the C++ transport connection is already closed, or when the socket write itself fails, MessageDispatcher either raises a non-retriable error or leaves the registered request promise pending until its timeout. Both send paths guard with a hard CELEBORN_CHECK(!closed_) and drop the write future:

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;
}

and TransportClientFactory::createClient raises a non-retriable CELEBORN_FAIL when a connection cannot be established. (CELEBORN_CHECK / CELEBORN_FAIL throw a CelebornRuntimeError with isRetriable=false; they do not abort the process.)

This PR:

• MessageDispatcher::operator() and sendFetchChunkRequest: the leading CELEBORN_CHECK(!closed_) becomes a fast path that returns a ready retriable future; the trailing check now fulfils the just-registered promise with a retriable exception (see the race note below) rather than asserting.
• Observes the write future on both paths. wangle::AsyncSocketHandler::write returns an already-failed future when !socket_->good() ("socket is closed in write()"), and otherwise fails it later from AsyncTransport::WriteCallback::writeErr; neither necessarily flips closed_ first, so discarding the future left the registered promise pending until the request timeout. The error continuation now removes the registry entry and fails it retriably. This is the analogue of Java's TransportClient StdChannelListener#operationComplete.
• MessageDispatcher::cleanup(): the two outstanding-request failures switch from a plain std::runtime_error to the same retriable exception, so every request failed on close carries consistent classification.
• Preserves the retriable classification through TransportClient. Previously sendRpcRequestSync re-threw via CELEBORN_FAIL (hardcoded isRetriable=false) and the push/fetch thenError continuations flattened the cause into a plain std::runtime_error, so the classification was discarded before any caller could observe it. The three paths now forward or re-wrap the cause keeping isRetriable. A non-retriable cause keeps byte-identical std::runtime_error(message) behaviour, so ShuffleClientImpl::getPushDataFailCause string matching is unaffected.
• TransportClientFactory::createClient: the connect failure throws a retriable CelebornRuntimeError instead of CELEBORN_FAIL.

The write-failure message is phrased "Failed to send request {}, errorMsg: {}", matching the Java listener's wording, so ShuffleClientImpl::getPushDataFailCause → connectFail() classifies it as PUSH_DATA_CONNECTION_EXCEPTION_PRIMARY rather than as the non-critical default.

Why are the changes needed?

Sending on a closed connection, or having a socket write fail, is a normal recoverable condition — the peer closed the socket because of a worker restart or an idle timeout — not an invariant violation. Raising a non-retriable error on that condition is wrong, and dropping the write future is worse: the promise stays pending until the request timeout, which is exactly the stall this change is meant to avoid.

This aligns the C++ dispatcher's connection-failure semantics with the Java client. In Java the request path never hard-fails on a dead channel:

• TransportResponseHandler#channelInactive calls failOutstandingRequests(new IOException("Connection from ... closed")), invoking every outstanding fetch/rpc/push callback's onFailure(cause). This is the analogue of the C++ cleanup() and of the trailing-check path below.
• On the send path, TransportClient#sendRpc / fetchChunk / pushData attach StdChannelListener, whose operationComplete closes the channel and calls handleFailure → the callback's onFailure when the write fails. The .thenError continuation on the write future is the wangle equivalent.

In both cases Java fails the request rather than raising a non-retriable error, and CelebornInputStream#createReaderWithRetry (and the push revive path) then retry or fail over to the replica.

One honest scope note: there is no retriable flag on the Java exceptions, and no C++ path branches on isRetriable() yet — both clients currently catch broadly and retry. What this PR fixes is that the classification was previously unobservable at the public API at all; a consumer that branches on it is follow-up work.

Does the close-during-send race matter?

close() sets closed_ = true before cleanup() locks the registry. If cleanup() runs between the leading fast-path check and the promise registration, the promise would be registered into an already-drained registry and never fulfilled — the future would hang until the request times out. The trailing closed_.load() check re-inspects the registry after registration and fails the promise with a retriable exception if it is still there, so the caller sees a ready retriable error instead of a stall. closeFailsInFlightRequestsRetriably covers the in-flight case.

Does this PR resolve a correctness bug?

• Yes

Shuffle output is unaffected. This is a robustness / failure-handling fix: a routine connection close or a failed write no longer raises a non-retriable error, and no longer leaves the request pending until its timeout; it surfaces as a retriable error the existing retry/failover paths can act on.

Does this PR introduce any user-facing change?

• Yes

No config, API, or shuffle-behaviour 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 unit tests in cpp/celeborn/network/tests/MessageDispatcherTest.cpp:

• sendRpcRequestAfterCloseFailsRetriably — a send issued after close() returns a ready, retriable exception instead of tripping the assertion.
• sendFetchChunkRequestAfterCloseFailsRetriably — same for the fetch path.
• closeFailsInFlightRequestsRetriably — close() fails an already-registered in-flight request retriably rather than leaving its future pending.
• sendRpcRequestFailedWriteFailsRequestRetriably / sendFetchChunkRequestFailedWriteFailsRetriably — a MockHandler whose write() returns a failed future fails the request retriably while the dispatcher stays available, covering the rpc/push and fetch paths.

New unit tests in cpp/celeborn/network/tests/TransportClientTest.cpp — a MockDispatcher returning a retriable failure, asserting the classification survives to the public API on the sync, push and fetch paths:

• sendRpcRequestSyncPreservesRetriableFailure
• pushDataAsyncPreservesRetriableFailure
• fetchChunkAsyncPreservesRetriableFailure

The C++ unit test suite passes locally on macOS/arm64 with no regressions, and is covered by the Celeborn Cpp Integration Test workflow (Run Unittests of Celeborn Cpp).

[Copilot]

Pull request overview

This PR adjusts the C++ client’s network failure semantics so that attempts to send/fetch on a closed transport connection (and connection-establishment failures) surface as retriable CelebornRuntimeErrors instead of non-retriable assertion-style failures, aligning behavior with the Java client’s “fail outstanding requests on channel inactive” contract.

Changes:

• Convert “connection closed” conditions in MessageDispatcher send/fetch paths from CELEBORN_CHECK to ready futures / promise failures carrying retriable CelebornRuntimeErrors (including close-during-send race handling).
• Make TransportClientFactory::createClient connect failures throw a retriable CelebornRuntimeError instead of CELEBORN_FAIL.
• Add unit tests verifying post-close send/fetch and in-flight request failure behavior is retriable and immediately observable via ready futures.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
cpp/celeborn/network/TransportClient.cpp	Marks connect-establishment failures as retriable by throwing CelebornRuntimeError(isRetriable=true) instead of CELEBORN_FAIL.
cpp/celeborn/network/MessageDispatcher.cpp	Replaces hard closed-connection assertions with retriable promise/future failures and adds race handling to prevent hung futures.
cpp/celeborn/network/tests/MessageDispatcherTest.cpp	Adds tests asserting closed-connection send/fetch and close()-failed in-flight requests produce ready futures with retriable exceptions.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[yugan95]

Valid — every failure reported the helper's own location, so all of them collapsed onto one line.

Fixed: the helper now takes file/line/function from the caller, and each failing site (operator(), sendFetchChunkRequest, failPendingRequest, failPendingFetch, cleanup()) passes its own __FILE__/__LINE__/__FUNCTION__.

[SteNicholas]

@yugan95, thanks for contribution. I have left some comments for this pull request. PTAL.

[SteNicholas]

Please observe and handle the write future here. Wangle's AsyncSocketHandler::write returns an exceptional future immediately when the socket is no longer good, and its asynchronous writeErr also completes that future exceptionally. Those failures do not necessarily set this dispatcher's closed_ before the trailing check, so discarding the future can leave the registered request promise pending until its timeout—the hang this change is intended to avoid. Please attach an error continuation that removes and fails the corresponding registry entry retriably, and cover both RPC/push and fetch with a mock handler that returns a failed write future. Reference: https://github.com/facebook/wangle/blob/v2024.07.01.00/wangle/channel/AsyncSocketHandler.h#L95-L113

[yugan95]

You're right, and this is the exact mechanism the description cites as the Java analogue. wangle::AsyncSocketHandler::write returns an already-failed future when !socket_->good() ("socket is closed in write()", AsyncSocketHandler.h), and otherwise fails it later from AsyncTransport::WriteCallback::writeErr. Neither path necessarily flips closed_ before the trailing check, so discarding the future left the registered promise pending until the request timeout — the hang this PR is about.

Both send paths now attach .thenError, which removes the registry entry and fails it retriably: failPendingRequest(requestId, ...) for rpc/push and failPendingFetch(streamChunkSlice, ...) for fetch. Both are no-ops when the entry is already gone (fulfilled, cleaned up or interrupted).

The message is phrased "Failed to send request {}, errorMsg: {}" to match Java's StdChannelListener, so ShuffleClientImpl::getPushDataFailCause -> connectFail() classifies it as PUSH_DATA_CONNECTION_EXCEPTION_PRIMARY instead of the non-critical default.

Covered by two new tests using a MockHandler whose write() returns a failed future: sendRpcRequestFailedWriteFailsRequestRetriably and sendFetchChunkRequestFailedWriteFailsRetriably — each asserts the dispatcher is still available (only the write failed) and the future is ready with a retriable CelebornException.

[SteNicholas]

This retriable classification is discarded at the next public API layer. TransportClient::sendRpcRequestSync catches the exception and rethrows via CELEBORN_FAIL (isRetriable=false), while the push and fetch thenError handlers convert it to std::runtime_error. Consequently callers cannot observe the isRetriable=true value asserted by these dispatcher-only tests. Please preserve the CelebornException classification through TransportClient and add API-level tests for the sync, push, and fetch paths.

[yugan95]

Confirmed — sendRpcRequestSync re-threw through CELEBORN_FAIL (hardcoded isRetriable=false) and the push/fetch thenError continuations flattened the cause into a plain std::runtime_error, so no caller could observe the flag.

TransportClient now preserves it:

• sendRpcRequestSync and fetchChunkAsync's catch use failPreservingRetriable(...), which carries the cause's isRetriable.
• the push/fetch thenError continuations use toCallbackException(...), forwarding a retriable CelebornRuntimeError as-is.
• the synchronous push catch uses wrapCallbackException(...), which keeps the flag while wrapping the contextual message.

A non-retriable cause still produces a byte-identical std::runtime_error(message), so getPushDataFailCause string matching is unchanged (there's a regression assertion for PUSH_DATA_FAIL_PARTITION_NOT_FOUND). The callback consumers only read what(), so forwarding the exception rather than a copy of its message is behaviour-neutral for them.

Added API-level tests in TransportClientTest.cpp with a MockDispatcher that returns a retriable failure: sendRpcRequestSyncPreservesRetriableFailure, pushDataAsyncPreservesRetriableFailure, fetchChunkAsyncPreservesRetriableFailure.

[yugan95]

Thanks for the review @SteNicholas. All three comments are addressed and pushed — the exception helper now carries the caller's location, both send paths observe the write future and fail the pending request retriably, and TransportClient no longer flattens the retriable classification. Replied inline with the details. PTAL.