Support configure default memory usage tracking through manager optio…

…ns (#6587)

Summary:
Add to support configure default memory usage tracking through manager
options. Prestissimo can use this to monitor the default memory usage such
as disk spilling memory pool usage.
Some dead code cleanup.

Pull Request resolved: #6587

Reviewed By: spershin

Differential Revision: D49306305

Pulled By: xiaoxmeng

fbshipit-source-id: 0121c6ae1673495736d3425863dbe62012fac9d5

velox/common/memory/Memory.cpp

@@ -16,7 +16,6 @@
 
 #include "velox/common/memory/Memory.h"
 
-DECLARE_bool(velox_enable_memory_usage_track_in_default_memory_pool);
 DECLARE_int32(velox_memory_num_shared_leaf_pools);
 
 namespace facebook::velox::memory {
@@ -35,8 +34,7 @@ MemoryManager::MemoryManager(const MemoryManagerOptions& options)
           {.kind = options.arbitratorKind,
            .capacity = std::min(options.queryMemoryCapacity, options.capacity),
            .memoryPoolInitCapacity = options.memoryPoolInitCapacity,
-           .memoryPoolTransferCapacity = options.memoryPoolTransferCapacity,
-           .retryArbitrationFailure = options.retryArbitrationFailure})),
+           .memoryPoolTransferCapacity = options.memoryPoolTransferCapacity})),
       alignment_(std::max(MemoryAllocator::kMinAlignment, options.alignment)),
       checkUsageLeak_(options.checkUsageLeak),
       debugEnabled_(options.debugEnabled),
@@ -53,8 +51,7 @@ MemoryManager::MemoryManager(const MemoryManagerOptions& options)
           MemoryPool::Options{
               .alignment = alignment_,
               .maxCapacity = kMaxMemory,
-              .trackUsage =
-                  FLAGS_velox_enable_memory_usage_track_in_default_memory_pool,
+              .trackUsage = options.trackDefaultUsage,
               .checkUsageLeak = options.checkUsageLeak,
               .debugEnabled = options.debugEnabled})} {
   VELOX_CHECK_NOT_NULL(allocator_);
@@ -88,16 +85,6 @@ MemoryManager::~MemoryManager() {
   }
 }
 
-#ifdef VELOX_ENABLE_BACKWARD_COMPATIBILITY
-// static
-MemoryManager& MemoryManager::getInstance(
-    const MemoryManagerOptions& options,
-    bool ensureCapacity) {
-  static MemoryManager manager{options};
-  return manager;
-}
-#endif
-
 // static
 MemoryManager& MemoryManager::getInstance(const MemoryManagerOptions& options) {
   static MemoryManager manager{options};

velox/common/memory/Memory.h

@@ -42,6 +42,7 @@
 
 DECLARE_bool(velox_memory_leak_check_enabled);
 DECLARE_bool(velox_memory_pool_debug_enabled);
+DECLARE_bool(velox_enable_memory_usage_track_in_default_memory_pool);
 
 namespace facebook::velox::memory {
 #define VELOX_MEM_LOG_PREFIX "[MEM] "
@@ -77,6 +78,10 @@ struct MemoryManagerOptions {
   /// capacity for system usage.
   int64_t queryMemoryCapacity{kMaxMemory};
 
+  /// If true, enable memory usage tracking in the default memory pool.
+  bool trackDefaultUsage{
+      FLAGS_velox_enable_memory_usage_track_in_default_memory_pool};
+
   /// If true, check the memory pool and usage leaks on destruction.
   ///
   /// TODO: deprecate this flag after all the existing memory leak use cases
@@ -105,14 +110,6 @@ struct MemoryManagerOptions {
   /// The minimal memory capacity to transfer out of or into a memory pool
   /// during the memory arbitration.
   uint64_t memoryPoolTransferCapacity{32 << 20};
-
-  /// If true, handle the memory arbitration failure by aborting the memory
-  /// pool with most capacity and retry the memory arbitration, otherwise we
-  /// simply fails the memory arbitration requestor itself. This helps the
-  /// distributed query execution use case such as Prestissimo that fail the
-  /// same query on all the workers instead of a random victim query which
-  /// happens to trigger the failed memory arbitration.
-  bool retryArbitrationFailure{true};
 };
 
 /// 'MemoryManager' is responsible for managing the memory pools. For now, users
@@ -130,12 +127,6 @@ class MemoryManager {
   FOLLY_EXPORT static MemoryManager& getInstance(
       const MemoryManagerOptions& options = MemoryManagerOptions{});
 
-#ifdef VELOX_ENABLE_BACKWARD_COMPATIBILITY
-  FOLLY_EXPORT static MemoryManager& getInstance(
-      const MemoryManagerOptions& options,
-      bool ensureCapacity);
-#endif
-
   /// Returns the memory capacity of this memory manager which puts a hard cap
   /// on memory usage, and any allocation that exceeds this capacity throws.
   int64_t capacity() const;

velox/common/memory/MemoryArbitrator.h

@@ -59,14 +59,6 @@ class MemoryArbitrator {
     /// The minimal memory capacity to transfer out of or into a memory pool
     /// during the memory arbitration.
     uint64_t memoryPoolTransferCapacity{32 << 20};
-
-    /// If true, handle the memory arbitration failure by aborting the memory
-    /// pool with most capacity and retry the memory arbitration, otherwise we
-    /// simply fails the memory arbitration requestor itself. This helps the
-    /// distributed query execution use case such as Prestissimo that fail the
-    /// same query on all the workers instead of a random victim query which
-    /// happens to trigger the failed memory arbitration.
-    bool retryArbitrationFailure{true};
   };
 
   using Factory = std::function<std::unique_ptr<MemoryArbitrator>(
@@ -198,6 +190,10 @@ class MemoryArbitrator {
     bool operator>=(const Stats& other) const;
     bool operator<=(const Stats& other) const;
 
+    bool empty() const {
+      return numRequests == 0;
+    }
+
     /// Returns the debug string of this stats.
     std::string toString() const;
   };

velox/common/memory/tests/MemoryArbitratorTest.cpp

@@ -119,8 +119,12 @@ TEST_F(MemoryArbitrationTest, queryMemoryCapacity) {
 }
 
 TEST_F(MemoryArbitrationTest, arbitratorStats) {
+  const MemoryArbitrator::Stats emptyStats;
+  ASSERT_TRUE(emptyStats.empty());
   const MemoryArbitrator::Stats anchorStats(5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5);
+  ASSERT_FALSE(anchorStats.empty());
   const MemoryArbitrator::Stats largeStats(8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8);
+  ASSERT_FALSE(largeStats.empty());
   ASSERT_TRUE(!(anchorStats == largeStats));
   ASSERT_TRUE(anchorStats != largeStats);
   ASSERT_TRUE(anchorStats < largeStats);
@@ -155,8 +159,8 @@ class FakeTestArbitrator : public MemoryArbitrator {
             {.kind = config.kind,
              .capacity = config.capacity,
              .memoryPoolInitCapacity = config.memoryPoolInitCapacity,
-             .memoryPoolTransferCapacity = config.memoryPoolTransferCapacity,
-             .retryArbitrationFailure = config.retryArbitrationFailure}) {}
+             .memoryPoolTransferCapacity = config.memoryPoolTransferCapacity}) {
+  }
 
   void reserveMemory(MemoryPool* pool, uint64_t bytes) override {
     VELOX_NYI();

velox/common/memory/tests/MemoryManagerTest.cpp

@@ -22,6 +22,7 @@
 #include "velox/common/memory/Memory.h"
 
 DECLARE_int32(velox_memory_num_shared_leaf_pools);
+DECLARE_bool(velox_enable_memory_usage_track_in_default_memory_pool);
 
 using namespace ::testing;
 
@@ -113,13 +114,13 @@ TEST_F(MemoryManagerTest, Ctor) {
 namespace {
 class FakeTestArbitrator : public MemoryArbitrator {
  public:
-  FakeTestArbitrator(const Config& config)
+  explicit FakeTestArbitrator(const Config& config)
       : MemoryArbitrator(
             {.kind = config.kind,
              .capacity = config.capacity,
              .memoryPoolInitCapacity = config.memoryPoolInitCapacity,
-             .memoryPoolTransferCapacity = config.memoryPoolTransferCapacity,
-             .retryArbitrationFailure = config.retryArbitrationFailure}) {}
+             .memoryPoolTransferCapacity = config.memoryPoolTransferCapacity}) {
+  }
 
   void reserveMemory(MemoryPool* pool, uint64_t bytes) override {
     VELOX_NYI();
@@ -327,6 +328,24 @@ TEST(MemoryHeaderTest, addDefaultLeafMemoryPool) {
   ASSERT_EQ(namedPool->name(), "namedPool");
 }
 
+TEST_F(MemoryManagerTest, defaultMemoryUsageTracking) {
+  for (bool trackDefaultMemoryUsage : {false, true}) {
+    MemoryManagerOptions options;
+    options.trackDefaultUsage = trackDefaultMemoryUsage;
+    MemoryManager manager{options};
+    auto defaultPool = manager.addLeafPool("defaultMemoryUsageTracking");
+    ASSERT_EQ(defaultPool->trackUsage(), trackDefaultMemoryUsage);
+  }
+
+  for (bool trackDefaultMemoryUsage : {false, true}) {
+    FLAGS_velox_enable_memory_usage_track_in_default_memory_pool =
+        trackDefaultMemoryUsage;
+    MemoryManager manager{};
+    auto defaultPool = manager.addLeafPool("defaultMemoryUsageTracking");
+    ASSERT_EQ(defaultPool->trackUsage(), trackDefaultMemoryUsage);
+  }
+}
+
 TEST_F(MemoryManagerTest, memoryPoolManagement) {
   const int alignment = 32;
   MemoryManagerOptions options;