[native] Add Linux Memory Manager

Co-authored-by: Minhan Cao <[email protected]>

.circleci/continue_config.yml

@@ -123,6 +123,7 @@ jobs:
                     -DPRESTO_ENABLE_REMOTE_FUNCTIONS=ON \
                     -DPRESTO_ENABLE_JWT=ON \
                     -DPRESTO_STATS_REPORTER_TYPE=PROMETHEUS \
+                    -DPRESTO_MEMORY_CHECKER_TYPE=LINUX_MEMORY_CHECKER \
                     -DCMAKE_PREFIX_PATH=/usr/local \
                     -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
                     -DMAX_LINK_JOBS=2

presto-docs/src/main/sphinx/presto_cpp/features.rst

@@ -153,6 +153,15 @@ is valid.
 If the time period exceeds the parameter value, the request is rejected as
 authentication failure (HTTP 401).
 
+LinuxMemoryChecker
+------------------
+
+The LinuxMemoryChecker extends from PeriodicMemoryChecker and periodically checks 
+memory usage using memory calculation from inactive_anon + active_anon in the memory stat 
+file from Linux cgroups V1 or V2. The LinuxMemoryChecker is used for Linux systems only.
+
+The LinuxMemoryChecker can be enabled by setting the CMake flag ``PRESTO_MEMORY_CHECKER_TYPE=LINUX_MEMORY_CHECKER``. 
+
 Async Data Cache and Prefetching
 --------------------------------
 

presto-docs/src/main/sphinx/presto_cpp/properties.rst

@@ -298,3 +298,38 @@ The configuration properties of Presto C++ workers are described here, in alphab
 * **Default value:** ``0.25``
 
   See description for ``shared-arbitrator.memory-pool-min-free-capacity``
+
+Memory Checker Properties
+-------------------------
+
+The LinuxMemoryChecker extends from PeriodicMemoryChecker and is used for Linux systems only. 
+The LinuxMemoryChecker can be enabled by setting the CMake flag ``PRESTO_MEMORY_CHECKER_TYPE=LINUX_MEMORY_CHECKER``. 
+The following properties for PeriodicMemoryChecker are as follows:
+
+``system-mem-pushback-enabled``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``boolean``
+* **Default value:** ``false``
+
+If set to ``true``, starts memory limit checker to trigger memory pushback when
+server is under low memory pressure.
+
+``system-mem-limit-gb``
+^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``integer``
+* **Default value:** ``55``
+
+Specifies the system memory limit that triggers the memory pushback or heap dump if
+the server memory usage is beyond this limit. A value of zero means no limit is set. 
+This only applies if ``system-mem-pushback-enabled`` is ``true``.
+
+``system-mem-shrink-gb``
+^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``integer``
+* **Default value:** ``8``
+
+Specifies the amount of memory to shrink when the memory pushback is
+triggered. This only applies if ``system-mem-pushback-enabled`` is ``true``.

presto-native-execution/Makefile

@@ -42,6 +42,9 @@ endif
 ifneq ($(PRESTO_STATS_REPORTER_TYPE),)
   EXTRA_CMAKE_FLAGS += -DPRESTO_STATS_REPORTER_TYPE=$(PRESTO_STATS_REPORTER_TYPE)
 endif
+ifneq ($(PRESTO_MEMORY_CHECKER_TYPE),)
+  EXTRA_CMAKE_FLAGS += -DPRESTO_MEMORY_CHECKER_TYPE=$(PRESTO_MEMORY_CHECKER_TYPE)
+endif
 
 CMAKE_FLAGS := -DTREAT_WARNINGS_AS_ERRORS=${TREAT_WARNINGS_AS_ERRORS}
 CMAKE_FLAGS += -DENABLE_ALL_WARNINGS=${ENABLE_WALL}

presto-native-execution/presto_cpp/main/CMakeLists.txt

@@ -122,3 +122,19 @@ if(PRESTO_STATS_REPORTER_TYPE)
         "${PRESTO_STATS_REPORTER_TYPE} is not a valid stats reporter name")
   endif()
 endif()
+
+if(PRESTO_MEMORY_CHECKER_TYPE)
+  add_compile_definitions(PRESTO_MEMORY_CHECKER_TYPE)
+  # Check if the current platform is Linux and the memory checker type is
+  # LINUX_MEMORY_CHECKER.
+  if(UNIX
+     AND NOT APPLE
+     AND (PRESTO_MEMORY_CHECKER_TYPE STREQUAL "LINUX_MEMORY_CHECKER"))
+    add_library(presto_linux_memory_checker OBJECT LinuxMemoryChecker.cpp)
+    target_link_libraries(presto_server presto_linux_memory_checker)
+  else()
+    message(
+      FATAL_ERROR
+        "${PRESTO_MEMORY_CHECKER_TYPE} is not a valid memory checker name")
+  endif()
+endif()

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

@@ -0,0 +1,178 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <boost/regex.hpp>
+#include <folly/Singleton.h>
+#include <folly/gen/File.h>
+#include <sys/stat.h>
+#include "presto_cpp/main/PeriodicMemoryChecker.h"
+#include "presto_cpp/main/common/Configs.h"
+
+namespace facebook::presto {
+
+class LinuxMemoryChecker : public PeriodicMemoryChecker {
+ public:
+  explicit LinuxMemoryChecker(const PeriodicMemoryChecker::Config& config)
+      : PeriodicMemoryChecker(config) {
+    // Find out what cgroup version (v1 or v2) we have based on the directory
+    // it's mounted.
+    struct stat buffer;
+    if ((stat(kCgroupV1Path, &buffer) == 0)) {
+      statFile_ = kCgroupV1Path;
+    } else if ((stat(kCgroupV2Path, &buffer) == 0)) {
+      statFile_ = kCgroupV2Path;
+    } else {
+      statFile_ = "None";
+    }
+    LOG(INFO) << fmt::format("Using memory stat file {}", statFile_);
+  }
+
+  ~LinuxMemoryChecker() override {}
+
+  int64_t getUsedMemory() {
+    return systemUsedMemoryBytes();
+  }
+
+  void setStatFile(std::string statFile) {
+    statFile_ = statFile;
+    LOG(INFO) << fmt::format("Changed to using memory stat file {}", statFile_);
+  }
+
+ protected:
+  // Current memory calculation used is inactive_anon + active_anon
+  // Our first attempt was using memInfo memTotal - memAvailable.
+  // However memInfo is not containerized so we reserve this as a
+  // last resort.
+  //
+  // Next we tried to use what docker/kubernetes uses for their
+  // calculation. cgroup usage_in_bytes - total_inactive_files.
+  // However we found out that usage_in_bytes is a fuzz value
+  // and has a chance for the sync to occur after the shrink
+  // polling interval. This would result in double shrinks.
+  //
+  // Therefore we decided on values from the memory.stat file
+  // that are real time statistics. At first we tried to use
+  // the calculation suggested by the kernel team RSS+CACHE(+SWAP)
+  // However we noticed that this value was not closely related to the
+  // value in usage_in_bytes which is used to OOMKill. We then looked
+  // at all of the values in the stat file and decided that
+  // inactive_anon + active_anon moves closest to that of
+  // usage_in_bytes
+  //
+  // NOTE: We do not know if cgroup V2 memory.current is a fuzz
+  // value. It may be better than what we currently use. For
+  // consistency we will match cgroup V1 and change if
+  // necessary.
+  int64_t systemUsedMemoryBytes() override {
+    size_t memAvailable = 0;
+    size_t memTotal = 0;
+    size_t inactiveAnon = 0;
+    size_t activeAnon = 0;
+
+    if (statFile_ != "None") {
+      folly::gen::byLine(statFile_.c_str()) |
+          [&](const folly::StringPiece& line) -> void {
+        if (inactiveAnon == 0) {
+          inactiveAnon =
+              extractNumericConfigValueWithRegex(line, kInactiveAnonRegex);
+        }
+
+        if (activeAnon == 0) {
+          activeAnon =
+              extractNumericConfigValueWithRegex(line, kActiveAnonRegex);
+        }
+
+        if (activeAnon != 0 && inactiveAnon != 0) {
+          return;
+        }
+      };
+
+      // Unit is in bytes.
+      return inactiveAnon + activeAnon;
+    }
+
+    // Last resort use host machine info.
+    folly::gen::byLine("/proc/meminfo") |
+        [&](const folly::StringPiece& line) -> void {
+      if (memAvailable == 0) {
+        memAvailable =
+            extractNumericConfigValueWithRegex(line, kMemAvailableRegex) * 1024;
+      }
+
+      if (memTotal == 0) {
+        memTotal =
+            extractNumericConfigValueWithRegex(line, kMemTotalRegex) * 1024;
+      }
+
+      if (memAvailable != 0 && memTotal != 0) {
+        return;
+      }
+    };
+    // Unit is in bytes.
+    return (memAvailable && memTotal) ? memTotal - memAvailable : 0;
+  }
+
+  int64_t mallocBytes() const override {
+    VELOX_UNSUPPORTED();
+  }
+
+  void periodicCb() override {
+    return;
+  }
+
+  bool heapDumpCb(const std::string& filePath) const override {
+    VELOX_UNSUPPORTED();
+  }
+
+  void removeDumpFile(const std::string& filePath) const override {
+    VELOX_UNSUPPORTED();
+  }
+
+ private:
+  const boost::regex kInactiveAnonRegex{R"!(inactive_anon\s*(\d+)\s*)!"};
+  const boost::regex kActiveAnonRegex{R"!(active_anon\s*(\d+)\s*)!"};
+  const boost::regex kMemAvailableRegex{R"!(MemAvailable:\s*(\d+)\s*kB)!"};
+  const boost::regex kMemTotalRegex{R"!(MemTotal:\s*(\d+)\s*kB)!"};
+  const char* kCgroupV1Path = "/sys/fs/cgroup/memory/memory.stat";
+  const char* kCgroupV2Path = "/sys/fs/cgroup/memory.stat";
+  std::string statFile_;
+
+  size_t extractNumericConfigValueWithRegex(
+      const folly::StringPiece& line,
+      const boost::regex& regex) {
+    boost::cmatch match;
+    if (boost::regex_match(line.begin(), line.end(), match, regex)) {
+      if (match.size() > 1) {
+        std::string numStr(match[1].str());
+        return folly::to<size_t>(numStr);
+      }
+    }
+    return 0;
+  }
+};
+
+folly::Singleton<facebook::presto::PeriodicMemoryChecker> checker(
+    []() -> facebook::presto::PeriodicMemoryChecker* {
+      PeriodicMemoryChecker::Config config;
+      auto* systemConfig = SystemConfig::instance();
+      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();
+    });
+
+} // namespace facebook::presto

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

@@ -219,4 +219,12 @@ void PeriodicMemoryChecker::pushbackMemory() {
             << " Effective bytes shrunk: "
             << velox::succinctBytes(actualFreedBytes);
 }
+
+#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([]() {
+  return nullptr;
+});
+#endif
 } // namespace facebook::presto

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

@@ -18,6 +18,7 @@
 #include <glog/logging.h>
 #include "CoordinatorDiscoverer.h"
 #include "presto_cpp/main/Announcer.h"
+#include "presto_cpp/main/PeriodicMemoryChecker.h"
 #include "presto_cpp/main/PeriodicTaskManager.h"
 #include "presto_cpp/main/SignalHandler.h"
 #include "presto_cpp/main/SystemConnector.h"
@@ -565,6 +566,8 @@ void PrestoServer::run() {
   addAdditionalPeriodicTasks();
   periodicTaskManager_->start();
 
+  addMemoryCheckerPeriodicTask();
+
   // Start everything. After the return from the following call we are shutting
   // down.
   httpServer_->start(getHttpServerFilters());
@@ -584,6 +587,8 @@ void PrestoServer::run() {
 
   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";
@@ -966,6 +971,18 @@ 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(); },

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

@@ -121,6 +121,10 @@ class PrestoServer {
 
   virtual void stopAdditionalPeriodicTasks(){};
 
+  virtual void addMemoryCheckerPeriodicTask();
+
+  virtual void stopMemoryCheckerPeriodicTask();
+
   virtual void initializeCoordinatorDiscoverer();
 
   virtual std::shared_ptr<velox::exec::TaskListener> getTaskListener();

presto-native-execution/presto_cpp/main/tests/CMakeLists.txt

@@ -24,6 +24,11 @@ add_executable(
   TaskManagerTest.cpp
   QueryContextManagerTest.cpp)
 
+if(DEFINED PRESTO_MEMORY_CHECKER_TYPE AND PRESTO_MEMORY_CHECKER_TYPE STREQUAL
+                                          "LINUX_MEMORY_CHECKER")
+  target_sources(presto_server_test PRIVATE LinuxMemoryCheckerTest.cpp)
+endif()
+
 add_test(
   NAME presto_server_test
   COMMAND presto_server_test