Fix the race between task termination and join table build

velox/exec/HashAggregation.cpp

@@ -457,4 +457,8 @@ void HashAggregation::close() {
   output_ = nullptr;
   groupingSet_.reset();
 }
+
+void HashAggregation::abort() {
+  close();
+}
 } // namespace facebook::velox::exec

velox/exec/HashAggregation.h

@@ -47,6 +47,8 @@ class HashAggregation : public Operator {
 
   void close() override;
 
+  void abort() override;
+
  private:
   void updateRuntimeStats();
 

velox/exec/HashBuild.cpp

@@ -716,6 +716,8 @@ bool HashBuild::finishHashBuild() {
     return false;
   }
 
+  TestValue::adjust("facebook::velox::exec::HashBuild::finishHashBuild", this);
+
   auto promisesGuard = folly::makeGuard([&]() {
     // Realize the promises so that the other Drivers (which were not
     // the last to finish) can continue from the barrier and finish.
@@ -736,7 +738,7 @@ bool HashBuild::finishHashBuild() {
     for (auto& peer : peers) {
       auto op = peer->findOperator(planNodeId());
       HashBuild* build = dynamic_cast<HashBuild*>(op);
-      VELOX_CHECK(build);
+      VELOX_CHECK_NOT_NULL(build);
       if (build->joinHasNullKeys_) {
         joinHasNullKeys_ = true;
         if (isAntiJoin(joinType_) && nullAware_ && !joinNode_->filter()) {
@@ -1055,8 +1057,8 @@ void HashBuild::reclaim(uint64_t /*unused*/) {
   }
 }
 
-void HashBuild::close() {
-  Operator::close();
+void HashBuild::abort() {
+  Operator::abort();
 
   // Free up major memory usage.
   joinBridge_.reset();

velox/exec/HashBuild.h

@@ -79,7 +79,7 @@ class HashBuild final : public Operator {
 
   void reclaim(uint64_t targetBytes) override;
 
-  void close() override;
+  void abort() override;
 
  private:
   void setState(State state);

velox/exec/HashProbe.cpp

@@ -1408,8 +1408,8 @@ void HashProbe::setRunning() {
   setState(ProbeOperatorState::kRunning);
 }
 
-void HashProbe::close() {
-  Operator::close();
+void HashProbe::abort() {
+  Operator::abort();
 
   // Free up major memory usage.
   joinBridge_.reset();

velox/exec/HashProbe.h

@@ -64,7 +64,7 @@ class HashProbe : public Operator {
     return false;
   }
 
-  void close() override;
+  void abort() override;
 
   void clearDynamicFilters() override;
 

velox/exec/Operator.cpp

@@ -499,7 +499,7 @@ void Operator::MemoryReclaimer::abort(memory::MemoryPool* pool) {
       driver->state().isTerminated);
   VELOX_CHECK(driver->task()->isCancelled());
 
-  // Calls operator close to free up major memory usage.
-  op_->close();
+  // Calls operator abort to free up major memory usage.
+  op_->abort();
 }
 } // namespace facebook::velox::exec

velox/exec/Operator.h

@@ -378,6 +378,12 @@ class Operator : public BaseRuntimeStatWriter {
     operatorCtx_->pool()->release();
   }
 
+  // Invoked by memory arbitrator to free up operator's resource on memory
+  // abort. The query will stop running after this.
+  virtual void abort() {
+    close();
+  }
+
   // Returns true if 'this' never has more output rows than input rows.
   virtual bool isFilter() const {
     return false;

velox/exec/OrderBy.cpp

@@ -410,8 +410,8 @@ void OrderBy::prepareOutput() {
   }
 }
 
-void OrderBy::close() {
-  Operator::close();
+void OrderBy::abort() {
+  Operator::abort();
 
   output_ = nullptr;
   spiller_.reset();

velox/exec/OrderBy.h

@@ -58,7 +58,7 @@ class OrderBy : public Operator {
 
   void reclaim(uint64_t targetBytes) override;
 
-  void close() override;
+  void abort() override;
 
  private:
   // Checks if input will fit in the existing memory and increases

velox/exec/tests/HashJoinTest.cpp

@@ -862,6 +862,32 @@ DEBUG_ONLY_TEST_P(MultiThreadedHashJoinTest, parallelJoinBuildCheck) {
   ASSERT_EQ(numDrivers_ == 1, !isParallelBuild);
 }
 
+DEBUG_ONLY_TEST_P(
+    MultiThreadedHashJoinTest,
+    raceBetweenTaskTerminateAndTableBuild) {
+  if (numDrivers_ == 1) {
+    return;
+  }
+  std::atomic<bool> isParallelBuild{false};
+  SCOPED_TESTVALUE_SET(
+      "facebook::velox::exec::HashBuild::finishHashBuild",
+      std::function<void(Operator*)>([&](Operator* op) {
+        auto task = op->testingOperatorCtx()->task();
+        task->requestAbort();
+      }));
+  VELOX_ASSERT_THROW(
+      HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get())
+          .numDrivers(numDrivers_)
+          .keyTypes({BIGINT(), VARCHAR()})
+          .probeVectors(1600, 5)
+          .buildVectors(1500, 5)
+          .referenceQuery(
+              "SELECT t_k0, t_k1, t_data, u_k0, u_k1, u_data FROM t, u WHERE t_k0 = u_k0 AND t_k1 = u_k1")
+          .injectSpill(false)
+          .run(),
+      "Aborted for external error");
+}
+
 TEST_P(MultiThreadedHashJoinTest, allTypes) {
   HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get())
       .keyTypes(