[native] Make PeriodicMemoryChecker member of PrestoServer

Author: spershin

presto-native-execution/presto_cpp/main/LinuxMemoryChecker.cpp

@@ -205,7 +205,9 @@ class LinuxMemoryChecker : public PeriodicMemoryChecker {
       };
 
       // Unit is in bytes.
-      return inactiveAnon + activeAnon;
+      const auto memBytes = inactiveAnon + activeAnon;
+      cachedSystemUsedMemoryBytes_ = memBytes;
+      return memBytes;
     }
 
     // Last resort use host machine info.
@@ -226,7 +228,10 @@ class LinuxMemoryChecker : public PeriodicMemoryChecker {
       }
     };
     // Unit is in bytes.
-    return (memAvailable && memTotal) ? memTotal - memAvailable : 0;
+    const auto memBytes =
+        (memAvailable && memTotal) ? memTotal - memAvailable : 0;
+    cachedSystemUsedMemoryBytes_ = memBytes;
+    return memBytes;
   }
 
   int64_t mallocBytes() const override {
@@ -274,20 +279,18 @@ class LinuxMemoryChecker : public PeriodicMemoryChecker {
   }
 };
 
-folly::Singleton<facebook::presto::PeriodicMemoryChecker> checker(
-    []() -> facebook::presto::PeriodicMemoryChecker* {
-      auto* systemConfig = SystemConfig::instance();
-      if (systemConfig->systemMemPushbackEnabled()) {
-        PeriodicMemoryChecker::Config config;
-        config.systemMemPushbackEnabled =
-            systemConfig->systemMemPushbackEnabled();
-        config.systemMemLimitBytes =
-            static_cast<uint64_t>(systemConfig->systemMemLimitGb()) << 30;
-        config.systemMemShrinkBytes =
-            static_cast<uint64_t>(systemConfig->systemMemShrinkGb()) << 30;
-        return std::make_unique<LinuxMemoryChecker>(config).release();
-      }
-      return nullptr;
-    });
+std::unique_ptr<PeriodicMemoryChecker> createMemoryChecker() {
+  auto* systemConfig = SystemConfig::instance();
+  if (systemConfig->systemMemPushbackEnabled()) {
+    PeriodicMemoryChecker::Config config;
+    config.systemMemPushbackEnabled = systemConfig->systemMemPushbackEnabled();
+    config.systemMemLimitBytes =
+        static_cast<uint64_t>(systemConfig->systemMemLimitGb()) << 30;
+    config.systemMemShrinkBytes =
+        static_cast<uint64_t>(systemConfig->systemMemShrinkGb()) << 30;
+    return std::make_unique<LinuxMemoryChecker>(config);
+  }
+  return nullptr;
+}
 
 } // namespace facebook::presto

presto-native-execution/presto_cpp/main/PeriodicMemoryChecker.cpp

@@ -223,8 +223,8 @@ void PeriodicMemoryChecker::pushbackMemory() {
 #ifndef PRESTO_MEMORY_CHECKER_TYPE
 // Initialize singleton for the checker to be nullptr if
 // PRESTO_MEMORY_CHECKER_TYPE is not defined.
-folly::Singleton<facebook::presto::PeriodicMemoryChecker> checker([]() {
+std::unique_ptr<PeriodicMemoryChecker> createMemoryChecker() {
   return nullptr;
-});
+}
 #endif
 } // namespace facebook::presto

presto-native-execution/presto_cpp/main/PeriodicMemoryChecker.h

@@ -76,9 +76,14 @@ class PeriodicMemoryChecker {
   /// Stops the 'PeriodicMemoryChecker'.
   void stop();
 
+  /// Returns the last known cached 'current' system memory usage in bytes.
+  int64_t cachedSystemUsedMemoryBytes() const {
+    return cachedSystemUsedMemoryBytes_;
+  }
+
  protected:
-  /// Returns current system memory usage. The returned value is used to compare
-  /// with 'Config::systemMemLimitBytes'.
+  /// Fetches and returns current system memory usage in bytes.
+  /// The returned value is used to compare with 'Config::systemMemLimitBytes'.
   virtual int64_t systemUsedMemoryBytes() = 0;
 
   /// Returns current bytes allocated by malloc. The returned value is used to
@@ -102,6 +107,7 @@ class PeriodicMemoryChecker {
   virtual void pushbackMemory();
 
   const Config config_;
+  std::atomic<int64_t> cachedSystemUsedMemoryBytes_{0};
 
  private:
   // Struct that stores the file names of the heap profiles dumped and the
@@ -137,4 +143,6 @@ class PeriodicMemoryChecker {
       std::greater<DumpFileInfo>>
       dumpFilesByHeapMemUsageMinPq_;
 };
+
+std::unique_ptr<PeriodicMemoryChecker> createMemoryChecker();
 } // namespace facebook::presto

presto-native-execution/presto_cpp/main/PrestoServer.cpp

@@ -527,8 +527,10 @@ void PrestoServer::run() {
   addServerPeriodicTasks();
   addAdditionalPeriodicTasks();
   periodicTaskManager_->start();
-
-  addMemoryCheckerPeriodicTask();
+  createPeriodicMemoryChecker();
+  if (memoryChecker_ != nullptr) {
+    memoryChecker_->start();
+  }
 
   auto setTaskUriCb = [&](bool useHttps, int port) {
     std::string taskUri;
@@ -618,12 +620,13 @@ void PrestoServer::run() {
   }
 
   PRESTO_SHUTDOWN_LOG(INFO) << "Stopping all periodic tasks";
-  periodicTaskManager_->stop();
 
+  if (memoryChecker_ != nullptr) {
+    memoryChecker_->stop();
+  }
+  periodicTaskManager_->stop();
   stopAdditionalPeriodicTasks();
 
-  stopMemoryCheckerPeriodicTask();
-
   // Destroy entities here to ensure we won't get any messages after Server
   // object is gone and to have nice log in case shutdown gets stuck.
   PRESTO_SHUTDOWN_LOG(INFO) << "Destroying Task Resource";
@@ -1032,18 +1035,6 @@ void PrestoServer::updateAnnouncerDetails() {
   }
 }
 
-void PrestoServer::addMemoryCheckerPeriodicTask() {
-  if (folly::Singleton<PeriodicMemoryChecker>::try_get()) {
-    folly::Singleton<PeriodicMemoryChecker>::try_get()->start();
-  }
-}
-
-void PrestoServer::stopMemoryCheckerPeriodicTask() {
-  if (folly::Singleton<PeriodicMemoryChecker>::try_get()) {
-    folly::Singleton<PeriodicMemoryChecker>::try_get()->stop();
-  }
-}
-
 void PrestoServer::addServerPeriodicTasks() {
   periodicTaskManager_->addTask(
       [server = this]() { server->populateMemAndCPUInfo(); },
@@ -1107,6 +1098,12 @@ void PrestoServer::addServerPeriodicTasks() {
   }
 }
 
+void PrestoServer::createPeriodicMemoryChecker() {
+  // The call below will either produce nullptr or unique pointer to an instance
+  // of LinuxMemoryChecker.
+  memoryChecker_ = createMemoryChecker();
+}
+
 std::shared_ptr<velox::exec::TaskListener> PrestoServer::getTaskListener() {
   return nullptr;
 }

presto-native-execution/presto_cpp/main/PrestoServer.h

@@ -67,6 +67,7 @@ class Announcer;
 class SignalHandler;
 class TaskManager;
 class TaskResource;
+class PeriodicMemoryChecker;
 class PeriodicTaskManager;
 class SystemConfig;
 
@@ -115,16 +116,14 @@ class PrestoServer {
   void enableAnnouncer(bool enable);
 
  protected:
+  virtual void createPeriodicMemoryChecker();
+
   /// Hook for derived PrestoServer implementations to add/stop additional
   /// periodic tasks.
   virtual void addAdditionalPeriodicTasks(){};
 
   virtual void stopAdditionalPeriodicTasks(){};
 
-  virtual void addMemoryCheckerPeriodicTask();
-
-  virtual void stopMemoryCheckerPeriodicTask();
-
   virtual void initializeCoordinatorDiscoverer();
 
   virtual std::shared_ptr<velox::exec::TaskListener> getTaskListener();
@@ -273,6 +272,7 @@ class PrestoServer {
   std::chrono::steady_clock::time_point start_;
   std::unique_ptr<PeriodicTaskManager> periodicTaskManager_;
   std::unique_ptr<PrestoServerOperations> prestoServerOperations_;
+  std::unique_ptr<PeriodicMemoryChecker> memoryChecker_;
 
   // We update these members asynchronously and return in http requests w/o
   // delay.