Fix memory pool dtor race in Substrait2VeloxValuesNodeConversionTest

facebookincubator · Oct 13, 2022 · d1d4520 · d1d4520
1 parent 6116963
commit d1d4520
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 14 deletions.
diff --git a/velox/common/memory/Memory.h b/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 check in release build after fix the memory
+    // counting issue discovered by Prestissmo in 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 byte allocated: {}, number of allocations: {}",
+          remainingBytes,
+          tracker->getCumulativeBytes(),
+          tracker->getNumAllocs());
+    }
+#endif
   }
 
   // Actual memory allocation operations. Can be delegated.

diff --git a/velox/substrait/tests/Substrait2VeloxValuesNodeConversionTest.cpp b/velox/substrait/tests/Substrait2VeloxValuesNodeConversionTest.cpp
@@ -32,6 +32,26 @@ using namespace facebook::velox::substrait;
 
 class Substrait2VeloxValuesNodeConversionTest : public OperatorTestBase {
  public:
+  Substrait2VeloxValuesNodeConversionTest::
+      ~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_ =