Fix flaky HashJoinTest.failedToReclaimFromHashJoinBuildersInNonReclai (…

…#9962)

Summary:
The test is flaky because of the following race condition. When reclaim happens in join node, there is a chance probe side pipeline/driver has not been initialized, making reclaimer not set for the operator. This results in probe side not reporting non reclaimable attempts because the reclaim codepath is not executed. Added additional coordination of code execution to avoid such inconsistent behavior of the test.

Stress tests have been run to verify the consistency of the fixed test.

Pull Request resolved: #9962

Reviewed By: xiaoxmeng

Differential Revision: D57880512

Pulled By: tanjialiang

fbshipit-source-id: bcc829301911f151d51aa68e8b3c75a93ebf5293

velox/exec/tests/HashJoinTest.cpp

@@ -6771,10 +6771,25 @@ DEBUG_ONLY_TEST_F(
       runHashJoinTask(vectors, queryCtx, numDrivers, pool(), false).data;
 
   std::atomic_bool nonReclaimableSectionWaitFlag{true};
+  std::atomic_bool reclaimerInitializationWaitFlag{true};
   folly::EventCount nonReclaimableSectionWait;
   std::atomic_bool memoryArbitrationWaitFlag{true};
   folly::EventCount memoryArbitrationWait;
 
+  std::atomic<uint32_t> numInitializedDrivers{0};
+  SCOPED_TESTVALUE_SET(
+      "facebook::velox::exec::Driver::runInternal",
+      std::function<void(exec::Driver*)>([&](exec::Driver* driver) {
+        numInitializedDrivers++;
+        // We need to make sure reclaimers on both build and probe side are set
+        // (in Operator::initialize) to avoid race conditions, producing
+        // consistent test results.
+        if (numInitializedDrivers.load() == 2) {
+          reclaimerInitializationWaitFlag = false;
+          nonReclaimableSectionWait.notifyAll();
+        }
+      }));
+
   std::atomic<bool> injectNonReclaimableSectionOnce{true};
   SCOPED_TESTVALUE_SET(
       "facebook::velox::common::memory::MemoryPoolImpl::allocateNonContiguous",
@@ -6809,8 +6824,11 @@ DEBUG_ONLY_TEST_F(
   });
 
   // Wait for the hash build operators to enter into non-reclaimable section.
-  nonReclaimableSectionWait.await(
-      [&]() { return !nonReclaimableSectionWaitFlag.load(); });
+  nonReclaimableSectionWait.await([&]() {
+    return (
+        !nonReclaimableSectionWaitFlag.load() &&
+        !reclaimerInitializationWaitFlag.load());
+  });
 
   // We expect capacity grow fails as we can't reclaim from hash join operators.
   memory::testingRunArbitration();