fix: Fix the server crash caused by remote exchange error (facebookincubator#13905)

Summary:
Pull Request resolved: facebookincubator#13905

The remote exchange might run into error so merge operator check if a remote merge source is blocked or not, it might trigger merge source to fetch data from the exchange client which in turn try to dequeue from the exchange source. The dequeue operation might throw if there is already an error. The recent merge spill support moves the merge source is blocked check inside SCOPE_EXIT which converts an throw to a server crash as we can't throw inside a dtor. This changes move the merge source is blocked check out of SCOPE_EXIT and add a unit test to reproduce and verify

Reviewed By: tanjialiang

Differential Revision: D77423282

fbshipit-source-id: b9fd6a62749a8e344b935a94ed46b74016fd4eb2

Author: xiaoxmeng

velox/exec/ExchangeQueue.cpp

@@ -16,6 +16,10 @@
 #include "velox/exec/ExchangeQueue.h"
 #include <algorithm>
 
+#include "velox/common/testutil/TestValue.h"
+
+using facebook::velox::common::testutil::TestValue;
+
 namespace facebook::velox::exec {
 
 SerializedPage::SerializedPage(
@@ -136,6 +140,8 @@ std::vector<std::unique_ptr<SerializedPage>> ExchangeQueue::dequeueLocked(
     ContinueFuture* future,
     ContinuePromise* stalePromise) {
   VELOX_CHECK_NOT_NULL(future);
+  TestValue::adjust(
+      "facebook::velox::exec::ExchangeQueue::dequeueLocked", this);
   if (!error_.empty()) {
     *atEnd = true;
     VELOX_FAIL(error_);

velox/exec/Merge.cpp

@@ -91,6 +91,10 @@ BlockingReason Merge::isBlocked(ContinueFuture* future) {
 
   maybeStartNextMergeSourceGroup();
 
+  if (sourceMerger_ != nullptr) {
+    sourceMerger_->isBlocked(sourceBlockingFutures_);
+  }
+
   if (sourceBlockingFutures_.empty()) {
     return BlockingReason::kNotBlocked;
   }
@@ -176,12 +180,6 @@ void Merge::setupSpillMerger() {
 }
 
 void Merge::maybeStartNextMergeSourceGroup() {
-  SCOPE_EXIT {
-    if (sourceMerger_ != nullptr) {
-      sourceMerger_->isBlocked(sourceBlockingFutures_);
-    }
-  };
-
   if (sourceMerger_ != nullptr || numStartedSources_ >= sources_.size()) {
     return;
   }

velox/exec/tests/MultiFragmentTest.cpp

@@ -1609,6 +1609,56 @@ TEST_P(MultiFragmentTest, mergeExchangeOverEmptySources) {
   }
 }
 
+DEBUG_ONLY_TEST_P(MultiFragmentTest, mergeExchangeFailureOnStart) {
+  std::vector<std::shared_ptr<Task>> tasks;
+  std::vector<std::string> leafTaskIds;
+
+  const auto injectErrorMsg{"injectError"};
+  SCOPED_TESTVALUE_SET(
+      "facebook::velox::exec::ExchangeQueue::dequeueLocked",
+      std::function<void(ExchangeQueue*)>(
+          ([&](ExchangeQueue* /*unused*/) { VELOX_FAIL(injectErrorMsg); })));
+
+  auto data = makeRowVector(rowType_, 0);
+
+  for (int i = 0; i < 2; ++i) {
+    auto taskId = makeTaskId("leaf-", i);
+    leafTaskIds.push_back(taskId);
+    auto plan =
+        PlanBuilder()
+            .values({data})
+            .partitionedOutput({}, 1, /*outputLayout=*/{}, GetParam().serdeKind)
+            .planNode();
+
+    auto task = makeTask(taskId, plan, tasks.size());
+    tasks.push_back(task);
+    task->start(4);
+  }
+
+  auto exchangeTaskId = makeTaskId("exchange-", 0);
+  auto plan = PlanBuilder()
+                  .mergeExchange(rowType_, {"c0"}, GetParam().serdeKind)
+                  .singleAggregation({"c0"}, {"count(1)"})
+                  .planNode();
+
+  std::vector<Split> leafTaskSplits;
+  for (auto leafTaskId : leafTaskIds) {
+    leafTaskSplits.emplace_back(remoteSplit(leafTaskId));
+  }
+  VELOX_ASSERT_THROW(
+      test::AssertQueryBuilder(plan, duckDbQueryRunner_)
+          .splits(std::move(leafTaskSplits))
+          .config(
+              core::QueryConfig::kShuffleCompressionKind,
+              common::compressionKindToString(GetParam().compressionKind))
+          .assertResults(""),
+      injectErrorMsg);
+
+  for (auto& task : tasks) {
+    ASSERT_TRUE(waitForTaskCompletion(task.get())) << task->taskId();
+  }
+}
+
 namespace {
 core::PlanNodePtr makeJoinOverExchangePlan(
     const RowTypePtr& exchangeType,