Add flag to suppress memory capacity exceeded error message (#4940)

Summary:
This to control the test output of some fuzzer memory tests.

Pull Request resolved: #4940

Reviewed By: tanjialiang

Differential Revision: D45877935

Pulled By: xiaoxmeng

fbshipit-source-id: 239aac33defd53bbe89745690070fdbfbd437d9c

velox/common/memory/MemoryPool.cpp

@@ -22,6 +22,8 @@
 #include "velox/common/memory/Memory.h"
 #include "velox/common/testutil/TestValue.h"
 
+DECLARE_bool(velox_suppress_memory_capacity_exceeding_error_message);
+
 using facebook::velox::common::testutil::TestValue;
 
 namespace facebook::velox::memory {
@@ -839,9 +841,12 @@ std::string MemoryPoolImpl::capExceedingMessage(
   VELOX_CHECK_NULL(parent_);
 
   std::stringstream out;
+  out << "\n" << errorMessage << "\n";
+  if (FLAGS_velox_suppress_memory_capacity_exceeding_error_message) {
+    return out.str();
+  }
   {
     std::lock_guard<std::mutex> l(mutex_);
-    out << "\n" << errorMessage << "\n";
     const Stats stats = statsLocked();
     const MemoryUsage usage{
         .name = name(),

velox/common/memory/tests/CMakeLists.txt

@@ -20,6 +20,7 @@ add_executable(
   HashStringAllocatorTest.cpp
   MemoryAllocatorTest.cpp
   MemoryArbitratorTest.cpp
+  MemoryCapExceededTest.cpp
   MemoryManagerTest.cpp
   MemoryPoolTest.cpp
   SharedArbitratorTest.cpp)

velox/exec/tests/MemoryCapExceededTest.cpp → velox/common/memory/tests/MemoryCapExceededTest.cpp

@@ -19,12 +19,27 @@
 
 #include <gmock/gmock.h>
 
+DECLARE_bool(velox_suppress_memory_capacity_exceeding_error_message);
+
 namespace facebook::velox::exec::test {
 namespace {
 
-class MemoryCapExceededTest : public OperatorTestBase {};
+class MemoryCapExceededTest : public OperatorTestBase,
+                              public testing::WithParamInterface<bool> {
+  void SetUp() override {
+    OperatorTestBase::SetUp();
+    // NOTE: if 'GetParam()' is true, then suppress the verbose error message in
+    // memory capacity exceeded exception.
+    FLAGS_velox_suppress_memory_capacity_exceeding_error_message = GetParam();
+  }
 
-TEST_F(MemoryCapExceededTest, singleDriver) {
+  void TearDown() override {
+    OperatorTestBase::TearDown();
+    FLAGS_velox_suppress_memory_capacity_exceeding_error_message = false;
+  }
+};
+
+TEST_P(MemoryCapExceededTest, singleDriver) {
   // Executes a plan with a single driver thread and query memory limit that
   // forces it to throw MEM_CAP_EXCEEDED exception. Verifies that the error
   // message contains all the details expected.
@@ -34,8 +49,9 @@ TEST_F(MemoryCapExceededTest, singleDriver) {
   constexpr int64_t kMaxBytes = 5LL << 20; // 5MB
   // We look for these lines separately, since their order can change (not sure
   // why).
-  std::array<std::string, 14> expectedTexts = {
-      "Exceeded memory pool cap of 5.00MB when requesting 2.00MB",
+  std::vector<std::string> expectedTexts = {
+      "Exceeded memory pool cap of 5.00MB when requesting 2.00MB"};
+  std::vector<std::string> expectedDetailedTexts = {
       "node.1 usage 1.00MB peak 1.00MB",
       "op.1.0.0.FilterProject usage 12.00KB peak 12.00KB",
       "node.2 usage 4.00MB peak 4.00MB",
@@ -78,10 +94,22 @@ TEST_F(MemoryCapExceededTest, singleDriver) {
           << "Expected error message to contain '" << expectedText
           << "', but received '" << errorMessage << "'.";
     }
+    for (const auto& expectedText : expectedDetailedTexts) {
+      LOG(ERROR) << expectedText;
+      if (!GetParam()) {
+        ASSERT_TRUE(errorMessage.find(expectedText) != std::string::npos)
+            << "Expected error message to contain '" << expectedText
+            << "', but received '" << errorMessage << "'.";
+      } else {
+        ASSERT_TRUE(errorMessage.find(expectedText) == std::string::npos)
+            << "Unexpected error message to contain '" << expectedText
+            << "', but received '" << errorMessage << "'.";
+      }
+    }
   }
 }
 
-TEST_F(MemoryCapExceededTest, multipleDrivers) {
+TEST_P(MemoryCapExceededTest, multipleDrivers) {
   // Executes a plan that runs with ten drivers and query memory limit that
   // forces it to throw MEM_CAP_EXCEEDED exception. Verifies that the error
   // message contains information that acknowledges the existence of N
@@ -101,7 +129,7 @@ TEST_F(MemoryCapExceededTest, multipleDrivers) {
     data.push_back(rowVector);
   }
 
-  std::array<std::string, 10> expectedTexts = {
+  std::vector<std::string> expectedTexts = {
       "op.1.0.9.Aggregation usage",
       "op.1.0.8.Aggregation usage",
       "op.1.0.7.Aggregation usage",
@@ -134,14 +162,20 @@ TEST_F(MemoryCapExceededTest, multipleDrivers) {
   } catch (const VeloxException& e) {
     const auto errorMessage = e.message();
     for (const auto& expectedText : expectedTexts) {
-      ASSERT_TRUE(errorMessage.find(expectedText) != std::string::npos)
-          << "Expected error message to contain '" << expectedText
-          << "', but received '" << errorMessage << "'.";
+      if (!GetParam()) {
+        ASSERT_TRUE(errorMessage.find(expectedText) != std::string::npos)
+            << "Expected error message to contain '" << expectedText
+            << "', but received '" << errorMessage << "'.";
+      } else {
+        ASSERT_TRUE(errorMessage.find(expectedText) == std::string::npos)
+            << "Unexpected error message to contain '" << expectedText
+            << "', but received '" << errorMessage << "'.";
+      }
     }
   }
 }
 
-TEST_F(MemoryCapExceededTest, memoryManagerCapacityExeededError) {
+TEST_P(MemoryCapExceededTest, memoryManagerCapacityExeededError) {
   // Executes a plan with no memory pool capacity limit but very small memory
   // manager's limit.
   memory::IMemoryManager::Options options{.capacity = 1 << 20};
@@ -152,8 +186,9 @@ TEST_F(MemoryCapExceededTest, memoryManagerCapacityExeededError) {
   constexpr int64_t kMaxBytes = 5LL << 20; // 5MB
   // We look for these lines separately, since their order can change (not sure
   // why).
-  std::array<std::string, 14> expectedTexts = {
-      "Exceeded memory manager cap of 1.00MB when requesting 368.00KB, memory pool cap is 5.00MB",
+  std::vector<std::string> expectedTexts = {
+      "Exceeded memory manager cap of 1.00MB when requesting 368.00KB, memory pool cap is 5.00MB"};
+  std::vector<std::string> expectedDetailedTexts = {
       "node.2 usage 1.00MB peak 2.00MB",
       "op.2.0.0.Aggregation usage 1012.00KB peak 1.35MB",
       "node.1 usage 1.00MB peak 1.00MB",
@@ -197,9 +232,25 @@ TEST_F(MemoryCapExceededTest, memoryManagerCapacityExeededError) {
           << "Expected error message to contain '" << expectedText
           << "', but received '" << errorMessage << "'.";
     }
+    for (const auto& expectedText : expectedDetailedTexts) {
+      if (!GetParam()) {
+        ASSERT_TRUE(errorMessage.find(expectedText) != std::string::npos)
+            << "Expected error message to contain '" << expectedText
+            << "', but received '" << errorMessage << "'.";
+      } else {
+        ASSERT_TRUE(errorMessage.find(expectedText) == std::string::npos)
+            << "Unexpected error message to contain '" << expectedText
+            << "', but received '" << errorMessage << "'.";
+      }
+    }
   }
   Task::testingWaitForAllTasksToBeDeleted();
 }
 
+VELOX_INSTANTIATE_TEST_SUITE_P(
+    MemoryCapExceededTest,
+    MemoryCapExceededTest,
+    testing::ValuesIn({false, true}));
+
 } // namespace
 } // namespace facebook::velox::exec::test

velox/common/memory/tests/SharedArbitratorTest.cpp

@@ -34,6 +34,7 @@
 #include "velox/vector/fuzzer/VectorFuzzer.h"
 
 DECLARE_bool(velox_memory_leak_check_enabled);
+DECLARE_bool(velox_suppress_memory_capacity_exceeding_error_message);
 
 using namespace ::testing;
 using namespace facebook::velox::common::testutil;
@@ -334,6 +335,7 @@ class MockSharedArbitrationTest : public testing::Test {
 
   void TearDown() override {
     clearQueries();
+    FLAGS_velox_suppress_memory_capacity_exceeding_error_message = false;
   }
 
   void setupMemory(
@@ -846,6 +848,7 @@ TEST_F(MockSharedArbitrationTest, enterArbitrationException) {
 }
 
 TEST_F(MockSharedArbitrationTest, concurrentArbitrations) {
+  FLAGS_velox_suppress_memory_capacity_exceeding_error_message = true;
   const int numQueries = 10;
   const int numOpsPerQuery = 5;
   std::vector<std::shared_ptr<MockQuery>> queries;
@@ -892,13 +895,14 @@ TEST_F(MockSharedArbitrationTest, concurrentArbitrations) {
 }
 
 TEST_F(MockSharedArbitrationTest, concurrentArbitrationWithTransientRoots) {
+  FLAGS_velox_suppress_memory_capacity_exceeding_error_message = true;
   std::mutex mutex;
   std::vector<std::shared_ptr<MockQuery>> queries;
   queries.push_back(addQuery());
   queries.back()->addMemoryOp();
 
   const int numMemThreads = 20;
-  const int numAllocationsPerQuery = 5000;
+  const int numAllocationsPerQuery = 3000;
   std::vector<std::thread> memThreads;
   for (int i = 0; i < numMemThreads; ++i) {
     memThreads.emplace_back([&, i]() {
@@ -934,7 +938,7 @@ TEST_F(MockSharedArbitrationTest, concurrentArbitrationWithTransientRoots) {
     });
   }
 
-  const int numControlOps = 2000;
+  const int numControlOps = 1000;
   const int maxNumQueries = 64;
   std::thread controlThread([&]() {
     folly::Random::DefaultGenerator rng;
@@ -2466,7 +2470,8 @@ DEBUG_ONLY_TEST_F(
   }
 }
 
-TEST_F(SharedArbitrationTest, DISABLED_concurrentArbitration) {
+TEST_F(SharedArbitrationTest, concurrentArbitration) {
+  FLAGS_velox_suppress_memory_capacity_exceeding_error_message = true;
   const int numVectors = 8;
   std::vector<RowVectorPtr> vectors;
   fuzzerOpts_.vectorSize = 32;

velox/exec/tests/CMakeLists.txt

@@ -57,7 +57,6 @@ add_executable(
   ProbeOperatorStateTest.cpp
   RoundRobinPartitionFunctionTest.cpp
   RowContainerTest.cpp
-  MemoryCapExceededTest.cpp
   SpillTest.cpp
   SpillOperatorGroupTest.cpp
   SpillerTest.cpp

velox/flag_definitions/flags.cpp

@@ -107,3 +107,9 @@ DEFINE_bool(
     velox_enable_memory_usage_track_in_default_memory_pool,
     false,
     "If true, enable memory usage tracking in the default memory pool");
+
+DEFINE_bool(
+    velox_suppress_memory_capacity_exceeding_error_message,
+    false,
+    "If true, suppress the verbose error message in memory capacity exceeded "
+    "exception. This is only used by test to control the test error output size");