diff --git a/cpp/celeborn/network/MessageDispatcher.cpp b/cpp/celeborn/network/MessageDispatcher.cpp index c9d87351418..6509927835a 100644 --- a/cpp/celeborn/network/MessageDispatcher.cpp +++ b/cpp/celeborn/network/MessageDispatcher.cpp @@ -17,10 +17,41 @@ #include "celeborn/network/MessageDispatcher.h" +#include + #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 -- the peer closed the socket because +// of a worker restart or an idle timeout -- not an invariant violation. +// Reporting it through the promise instead of asserting mirrors the Java +// client, where an inactive channel surfaces as a retriable IOException via +// TransportResponseHandler#channelInactive -> failOutstandingRequests and a +// failed write surfaces via TransportClient's StdChannelListener, which +// CelebornInputStream then retries or fails over to the replica. +// +// The caller passes its own __FILE__/__LINE__/__FUNCTION__ so the exception +// points at the failing site rather than at this helper. +folly::exception_wrapper makeRetriableTransportError( + const char* file, + size_t line, + const char* function, + const std::string& detail) { + return folly::make_exception_wrapper( + file, + line, + function, + /*expression=*/"", + /*message=*/detail, + utils::error_source::kErrorSourceRuntime.c_str(), + utils::error_code::kInvalidState.c_str(), + /*isRetriable=*/true); +} +} // namespace + void MessageDispatcher::read(Context*, std::unique_ptr toRecvMsg) { switch (toRecvMsg->type()) { case Message::RPC_RESPONSE: { @@ -138,7 +169,6 @@ void MessageDispatcher::read(Context*, std::unique_ptr toRecvMsg) { folly::Future> MessageDispatcher::operator()( std::unique_ptr toSendMsg) { - CELEBORN_CHECK(!closed_); auto currTime = std::chrono::system_clock::now(); long requestId; switch (toSendMsg->type()) { @@ -163,6 +193,18 @@ folly::Future> MessageDispatcher::operator()( } } + // Fast path: the connection is already closed. Fail with a retriable error + // rather than asserting, so the caller's retry/failover logic can recover. + if (closed_.load()) { + return folly::makeFuture>( + makeRetriableTransportError( + __FILE__, + __LINE__, + __FUNCTION__, + fmt::format( + "connection closed before sending requestId {}", requestId))); + } + auto f = requestIdRegistry_.withLock( [&](auto& registry) -> folly::Future> { auto& holder = registry[requestId]; @@ -176,12 +218,88 @@ folly::Future> MessageDispatcher::operator()( return p.getFuture(); }); - this->pipeline_->write(std::move(toSendMsg)); + // 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 from AsyncTransport::WriteCallback::writeErr. Neither necessarily + // flips closed_ before the check below, so dropping the future would leave + // the promise pending until the request timeout. + this->pipeline_->write(std::move(toSendMsg)) + .thenError([this, requestId](const folly::exception_wrapper& e) { + // Phrased like the Java listener's message so that + // ShuffleClientImpl::getPushDataFailCause classifies it as a + // connection failure rather than a non-critical one. + failPendingRequest( + requestId, + fmt::format( + "Failed to send request {}, errorMsg: {}", + requestId, + e.what().toStdString())); + }); - CELEBORN_CHECK(!closed_); + // close() flips closed_ before cleanup() locks the registry. If cleanup() ran + // between the fast-path check and the registration above, our promise would + // never be fulfilled. Detect that and fail it here so the caller sees a + // retriable error instead of hanging until the request times out. + if (closed_.load()) { + failPendingRequest( + requestId, + fmt::format("connection closed while sending requestId {}", requestId)); + } return f; } +void MessageDispatcher::failPendingRequest( + long requestId, + const std::string& detail) { + bool found = true; + auto holder = requestIdRegistry_.withLock([&](auto& registry) { + auto search = registry.find(requestId); + if (search == registry.end()) { + // Already fulfilled, cleaned up or interrupted. + found = false; + return MsgPromiseHolder{}; + } + auto result = std::move(search->second); + registry.erase(search); + return std::move(result); + }); + if (!found) { + return; + } + LOG(WARNING) << detail; + // Fulfil outside the registry lock, like read() does: setException may run + // the caller's continuation inline -- a failed write is reported from the + // socket's write callback, by which time the continuation is attached -- and + // that continuation may re-enter the dispatcher. + holder.msgPromise.setException( + makeRetriableTransportError(__FILE__, __LINE__, __FUNCTION__, detail)); +} + +void MessageDispatcher::failPendingFetch( + const protocol::StreamChunkSlice& streamChunkSlice, + const std::string& detail) { + bool found = true; + auto holder = streamChunkSliceRegistry_.withLock([&](auto& registry) { + auto search = registry.find(streamChunkSlice); + if (search == registry.end()) { + // Already fulfilled, cleaned up or interrupted. + found = false; + return MsgPromiseHolder{}; + } + auto result = std::move(search->second); + registry.erase(search); + return std::move(result); + }); + if (!found) { + return; + } + LOG(WARNING) << detail; + // Fulfil outside the registry lock: see failPendingRequest. + holder.msgPromise.setException( + makeRetriableTransportError(__FILE__, __LINE__, __FUNCTION__, detail)); +} + folly::Future> MessageDispatcher::sendPushDataRequest( std::unique_ptr toSendMsg) { return (*this)(std::move(toSendMsg)); @@ -191,8 +309,22 @@ folly::Future> MessageDispatcher::sendFetchChunkRequest( const protocol::StreamChunkSlice& streamChunkSlice, std::unique_ptr toSendMsg) { - CELEBORN_CHECK(!closed_); CELEBORN_CHECK(toSendMsg->type() == Message::RPC_REQUEST); + + // Fast path: the connection is already closed. Fail with a retriable error + // rather than asserting, so CelebornInputStream can retry or fail over to a + // replica. + if (closed_.load()) { + return folly::makeFuture>( + makeRetriableTransportError( + __FILE__, + __LINE__, + __FUNCTION__, + fmt::format( + "connection closed before fetching streamChunkSlice {}", + streamChunkSlice.toString()))); + } + auto f = streamChunkSliceRegistry_.withLock([&](auto& registry) { auto& holder = registry[streamChunkSlice]; holder.requestTime = std::chrono::system_clock::now(); @@ -206,8 +338,26 @@ MessageDispatcher::sendFetchChunkRequest( }); return p.getFuture(); }); - this->pipeline_->write(std::move(toSendMsg)); - CELEBORN_CHECK(!closed_); + // Write-failure handling: see operator(). + this->pipeline_->write(std::move(toSendMsg)) + .thenError([this, streamChunkSlice](const folly::exception_wrapper& e) { + failPendingFetch( + streamChunkSlice, + fmt::format( + "Failed to send request for streamChunkSlice {}, errorMsg: {}", + streamChunkSlice.toString(), + e.what().toStdString())); + }); + + // Race handling: see operator(). If cleanup() ran before we registered above, + // fail our promise here rather than leaking it. + if (closed_.load()) { + failPendingFetch( + streamChunkSlice, + fmt::format( + "connection closed while fetching streamChunkSlice {}", + streamChunkSlice.toString())); + } return f; } @@ -275,7 +425,8 @@ void MessageDispatcher::cleanup() { auto errorMsg = fmt::format("Client closed, cancel ongoing requestId {}", requestId); LOG(WARNING) << errorMsg; - promiseHolder.msgPromise.setException(std::runtime_error(errorMsg)); + promiseHolder.msgPromise.setException(makeRetriableTransportError( + __FILE__, __LINE__, __FUNCTION__, errorMsg)); } registry.clear(); }); @@ -285,7 +436,8 @@ void MessageDispatcher::cleanup() { "Client closed, cancel ongoing streamChunkSlice {}", streamChunkSlice.toString()); LOG(WARNING) << errorMsg; - promiseHolder.msgPromise.setException(std::runtime_error(errorMsg)); + promiseHolder.msgPromise.setException(makeRetriableTransportError( + __FILE__, __LINE__, __FUNCTION__, errorMsg)); } registry.clear(); }); diff --git a/cpp/celeborn/network/MessageDispatcher.h b/cpp/celeborn/network/MessageDispatcher.h index 7ce6a7afe0b..e584d091933 100644 --- a/cpp/celeborn/network/MessageDispatcher.h +++ b/cpp/celeborn/network/MessageDispatcher.h @@ -96,6 +96,16 @@ class MessageDispatcher : public wangle::ClientDispatcherBase< private: void cleanup(); + // Fails the pending rpc/push request with a retriable error carrying + // `detail`, and unregisters it. A no-op when the request is no longer + // registered, i.e. it has already been fulfilled, cleaned up or interrupted. + void failPendingRequest(long requestId, const std::string& detail); + + // The sendFetchChunkRequest counterpart of failPendingRequest. + void failPendingFetch( + const protocol::StreamChunkSlice& streamChunkSlice, + const std::string& detail); + using MsgPromise = folly::Promise>; struct MsgPromiseHolder { MsgPromise msgPromise; diff --git a/cpp/celeborn/network/TransportClient.cpp b/cpp/celeborn/network/TransportClient.cpp index c754eecd037..860d03758d4 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(&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 toCallbackException( + const folly::exception_wrapper& e) { + std::unique_ptr retriableFailure; + e.with_exception([&](const utils::CelebornRuntimeError& error) { + if (error.isRetriable()) { + retriableFailure = std::make_unique(error); + } + }); + if (retriableFailure) { + return retriableFailure; + } + return std::make_unique(e.what().toStdString()); +} + +// The counterpart of toCallbackException for a synchronously thrown cause: +// wraps `errorMsg` while keeping the cause's retriable classification. +std::unique_ptr 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(errorMsg); + } + return std::make_unique( + 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 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(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(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(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(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(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 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 959aca6e32d..dc22141cee5 100644 --- a/cpp/celeborn/network/tests/MessageDispatcherTest.cpp +++ b/cpp/celeborn/network/tests/MessageDispatcherTest.cpp @@ -32,16 +32,26 @@ class MockHandler : public wangle::Handler< public: MockHandler(std::unique_ptr& 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. + MockHandler(std::unique_ptr& writedMsg, std::string writeError) + : writedMsg_(writedMsg), writeError_(std::move(writeError)) {} + void read(Context* ctx, std::unique_ptr msg) override {} folly::Future write(Context* ctx, std::unique_ptr msg) override { writedMsg_ = std::move(msg); + if (!writeError_.empty()) { + return folly::makeFuture(std::runtime_error(writeError_)); + } return {}; } private: std::unique_ptr& writedMsg_; + const std::string writeError_; }; SerializePipeline::Ptr createMockedPipeline(MockHandler&& mockHandler) { @@ -66,6 +76,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>&& 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 +323,123 @@ 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 sentMsg; + MockHandler mockHandler(sentMsg); + auto mockPipeline = createMockedPipeline(std::move(mockHandler)); + auto dispatcher = std::make_unique(); + 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( + 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 sentMsg; + MockHandler mockHandler(sentMsg); + auto mockPipeline = createMockedPipeline(std::move(mockHandler)); + auto dispatcher = std::make_unique(); + 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( + 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 sentMsg; + MockHandler mockHandler(sentMsg); + auto mockPipeline = createMockedPipeline(std::move(mockHandler)); + auto dispatcher = std::make_unique(); + dispatcher->setPipeline(mockPipeline.get()); + + const long requestId = 3001; + const std::string requestBody = "test-request-body"; + auto rpcRequest = std::make_unique( + 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. 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 closed_ is not necessarily set at that point. +TEST(MessageDispatcherTest, sendRpcRequestFailedWriteFailsRequestRetriably) { + std::unique_ptr sentMsg; + MockHandler mockHandler(sentMsg, "socket is closed in write()"); + auto mockPipeline = createMockedPipeline(std::move(mockHandler)); + auto dispatcher = std::make_unique(); + dispatcher->setPipeline(mockPipeline.get()); + + const long requestId = 4001; + const std::string requestBody = "test-request-body"; + auto rpcRequest = std::make_unique( + requestId, toReadOnlyByteBuffer(requestBody)); + auto future = dispatcher->sendRpcRequest(std::move(rpcRequest)); + + // The dispatcher is still open: only the write failed. + EXPECT_TRUE(dispatcher->isAvailable()); + ASSERT_TRUE(future.isReady()); + ASSERT_TRUE(future.hasException()); + EXPECT_TRUE(failedRetriably(std::move(future))); +} + +TEST(MessageDispatcherTest, sendFetchChunkRequestFailedWriteFailsRetriably) { + std::unique_ptr sentMsg; + MockHandler mockHandler(sentMsg, "socket is closed in write()"); + auto mockPipeline = createMockedPipeline(std::move(mockHandler)); + auto dispatcher = std::make_unique(); + dispatcher->setPipeline(mockPipeline.get()); + + const protocol::StreamChunkSlice streamChunkSlice{4001, 4002, 4003, 4004}; + const long requestId = 4001; + const std::string requestBody = "test-request-body"; + auto rpcRequest = std::make_unique( + requestId, toReadOnlyByteBuffer(requestBody)); + auto future = dispatcher->sendFetchChunkRequest( + streamChunkSlice, std::move(rpcRequest)); + + EXPECT_TRUE(dispatcher->isAvailable()); + ASSERT_TRUE(future.isReady()); + ASSERT_TRUE(future.hasException()); + EXPECT_TRUE(failedRetriably(std::move(future))); +} + TEST(MessageDispatcherTest, heartbeatIsSilentlyConsumed) { std::unique_ptr sentMsg; MockHandler mockHandler(sentMsg); diff --git a/cpp/celeborn/network/tests/TransportClientTest.cpp b/cpp/celeborn/network/tests/TransportClientTest.cpp index 0135d61e57f..df432f498cc 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( + __FILE__, + static_cast(__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(exception); + return celebornException != nullptr && celebornException->isRetriable(); +} + class MockDispatcher : public MessageDispatcher { public: folly::Future> sendRpcRequest( std::unique_ptr toSendMsg) override { sentMsg_ = std::move(toSendMsg); + if (connectionClosed_) { + return folly::makeFuture>( + retriableConnectionClosed()); + } msgPromise_ = MsgPromise(); return msgPromise_.getFuture(); } @@ -42,6 +69,10 @@ class MockDispatcher : public MessageDispatcher { folly::Future> sendPushDataRequest( std::unique_ptr toSendMsg) override { sentMsg_ = std::move(toSendMsg); + if (connectionClosed_) { + return folly::makeFuture>( + retriableConnectionClosed()); + } msgPromise_ = MsgPromise(); return msgPromise_.getFuture(); } @@ -50,6 +81,10 @@ class MockDispatcher : public MessageDispatcher { const protocol::StreamChunkSlice& streamChunkSlice, std::unique_ptr toSendMsg) override { sentMsg_ = std::move(toSendMsg); + if (connectionClosed_) { + return folly::makeFuture>( + 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 sentMsg_; MsgPromise msgPromise_; + bool connectionClosed_{false}; }; std::unique_ptr 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->setConnectionClosed(); + TransportClient client(nullptr, std::move(mockDispatcher), MS(10000)); + + const long requestId = 1001; + auto rpcRequest = std::make_unique( + 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->setConnectionClosed(); + TransportClient client(nullptr, std::move(mockDispatcher), MS(10000)); + auto mockRpcResponseCallback = std::make_shared(); + + auto pushData = std::make_unique( + /*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->setConnectionClosed(); + TransportClient client(nullptr, std::move(mockDispatcher), MS(10000)); + + const protocol::StreamChunkSlice streamChunkSlice{1, 2, 3, 4}; + auto rpcRequest = std::make_unique( + /*requestId=*/1001, toReadOnlyByteBuffer("test-request-body")); + std::unique_ptr onSuccessBuffer; + FetchChunkSuccessCallback onSuccess = + [&](protocol::StreamChunkSlice slice, + std::unique_ptr buffer) { + onSuccessBuffer = std::move(buffer); + }; + std::unique_ptr onFailureException; + FetchChunkFailureCallback onFailure = + [&](protocol::StreamChunkSlice slice, + std::unique_ptr 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(); auto rawMockDispatcher = mockDispatcher.get();