Fix memory pool dtor race in VeloxSubstraitRoundTripPlanConverterTest

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

Hide file tree

Showing 3 changed files with 33 additions and 15 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 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.

diff --git a/velox/substrait/tests/Substrait2VeloxValuesNodeConversionTest.cpp b/velox/substrait/tests/Substrait2VeloxValuesNodeConversionTest.cpp
@@ -31,7 +31,7 @@ using namespace facebook::velox::exec::test;
 using namespace facebook::velox::substrait;
 
 class Substrait2VeloxValuesNodeConversionTest : public OperatorTestBase {
- public:
+ protected:
   std::unique_ptr<memory::ScopedMemoryPool> pool_{
       memory::getDefaultScopedMemoryPool()};
   std::shared_ptr<SubstraitVeloxPlanConverter> planConverter_ =

diff --git a/velox/substrait/tests/VeloxSubstraitRoundTripPlanConverterTest.cpp b/velox/substrait/tests/VeloxSubstraitRoundTripPlanConverterTest.cpp
@@ -30,6 +30,24 @@ using namespace facebook::velox::substrait;
 
 class VeloxSubstraitRoundTripPlanConverterTest : public OperatorTestBase {
  protected:
+  ~VeloxSubstraitRoundTripPlanConverterTest() {
+    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));
+    }
+  }
+
   /// Makes a vector of INTEGER type with 'size' RowVectorPtr.
   /// @param size The number of RowVectorPtr.
   /// @param childSize The number of columns for each row.