Fix memory pool dtor race in Substrait2VeloxValuesNodeConversionTest (f…

…acebookincubator#2828)

Summary:
This is caused by the async task deletion (task dtor triggered by
async driver removal). The fix is to add memory counting wait in
Substrait2VeloxValuesNodeConversionTest dtor.

Pull Request resolved: facebookincubator#2828

Differential Revision: D40332245

Pulled By: xiaoxmeng

fbshipit-source-id: 7b433274385b9ed0e08cec06cd658e2a5f012a62

velox/common/memory/Memory.h

@@ -424,20 +424,20 @@ class MemoryPoolImpl : public MemoryPoolBase {
       int64_t cap = kMaxMemory);
 
   ~MemoryPoolImpl() {
-    // TODO(xiaoxmeng): Inconsistency discovered by following check when
-    // running with Prestissmo. Re-enable this check after the issue gets fixed.
-    //
-    // if (const auto& tracker = getMemoryUsageTracker()) {
-    //   auto remainingBytes = tracker->getCurrentUserBytes();
-    //   VELOX_CHECK_EQ(
-    //       0,
-    //       remainingBytes,
-    //       "Memory pool should be destroyed only after all allocated memory "
-    //       "has been freed. Remaining bytes allocated: {}, cumulative bytes
-    //       allocated: {}, number of allocations: {}", remainingBytes,
-    //       tracker->getCumulativeBytes(),
-    //       tracker->getNumAllocs());
-    // }
+    // TODO(xiaoxmeng): enable this in release build after fix the issue exposed
+    // by Prestissmo in real cluster.
+#ifndef NDEBUG
+    if (const auto& tracker = getMemoryUsageTracker()) {
+      auto remainingBytes = tracker->getCurrentUserBytes();
+      VELOX_CHECK_EQ(
+          0,
+          remainingBytes,
+          "Memory pool should be destroyed only after all allocated memory has been freed. Remaining bytes allocated: {}, cumulative bytes allocated: {}, number of allocations: {}",
+          remainingBytes,
+          tracker->getCumulativeBytes(),
+          tracker->getNumAllocs());
+    }
+#endif
   }
 
   // Actual memory allocation operations. Can be delegated.

velox/substrait/tests/Substrait2VeloxValuesNodeConversionTest.cpp

@@ -32,6 +32,25 @@ using namespace facebook::velox::substrait;
 
 class Substrait2VeloxValuesNodeConversionTest : public OperatorTestBase {
  public:
+  ~Substrait2VeloxValuesNodeConversionTest() {
+    auto tracker = pool_->getMemoryUsageTracker();
+    if (tracker == nullptr) {
+      return;
+    }
+    // Wait for current user bytes count dropping to zero as there might be
+    // async event still running at the background like the last task reference
+    // dropped at the background.
+    const uint64_t kMaxWaitUs = 3'000'000;
+    while (tracker->getCurrentUserBytes() != 0) {
+      const uint64_t kWaitIntervalUs = 100'000;
+      LOG(WARNING) << "Memory pool has " << tracker->getCurrentUserBytes()
+                   << " current user bytes, wait " << kWaitIntervalUs
+                   << " us for the pending usage to be released";
+      std::this_thread::sleep_for(std::chrono::microseconds(kWaitIntervalUs));
+    }
+  }
+
+ protected:
   std::unique_ptr<memory::ScopedMemoryPool> pool_{
       memory::getDefaultScopedMemoryPool()};
   std::shared_ptr<SubstraitVeloxPlanConverter> planConverter_ =