yugan95 commented on code in PR #3801:
URL: https://github.com/apache/celeborn/pull/3801#discussion_r3800420526


##########
cpp/celeborn/network/MessageDispatcher.cpp:
##########
@@ -191,30 +447,86 @@ 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;
+  }

Review Comment:
   Right — added a `catch (...)` next to it. Covered by 
`nonStdThrowingWriteWithoutResponseIsReported`, whose mock handler throws a 
`std::string`.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to