From a9b947d9ac422290cc42b4c29ba5cc1b934244b3 Mon Sep 17 00:00:00 2001
From: Antonio Andelic <antonio@clickhouse.com>
Date: Sun, 28 Jul 2024 11:51:09 +0200
Subject: [PATCH] Extract common

---
 base/base/CMakeLists.txt      |  1 +
 base/base/Numa.cpp            | 37 +++++++++++++++++++++++++++++++++++
 base/base/Numa.h              | 12 ++++++++++++
 base/base/getMemoryAmount.cpp | 29 +++------------------------
 programs/keeper/Keeper.cpp    |  7 +++++++
 programs/server/Server.cpp    | 36 ++++------------------------------
 6 files changed, 64 insertions(+), 58 deletions(-)
 create mode 100644 base/base/Numa.cpp
 create mode 100644 base/base/Numa.h

diff --git a/base/base/CMakeLists.txt b/base/base/CMakeLists.txt
index 451a6eb5e8b..341c92d3042 100644
--- a/base/base/CMakeLists.txt
+++ b/base/base/CMakeLists.txt
@@ -32,6 +32,7 @@ set (SRCS
     StringRef.cpp
     safeExit.cpp
     throwError.cpp
+    Numa.cpp
 )
 
 add_library (common ${SRCS})
diff --git a/base/base/Numa.cpp b/base/base/Numa.cpp
new file mode 100644
index 00000000000..0bf56a993b8
--- /dev/null
+++ b/base/base/Numa.cpp
@@ -0,0 +1,37 @@
+#include <base/Numa.h>
+
+#include "config.h"
+
+#if USE_NUMACTL
+#    include <numa.h>
+#endif
+
+namespace DB
+{
+
+std::optional<size_t> getNumaNodesTotalMemory()
+{
+    std::optional<size_t> total_memory;
+#if USE_NUMACTL
+    if (numa_available() != -1)
+    {
+        auto * membind = numa_get_membind();
+        if (!numa_bitmask_equal(membind, numa_all_nodes_ptr))
+        {
+            total_memory.emplace(0);
+            auto max_node = numa_max_node();
+            for (int i = 0; i <= max_node; ++i)
+            {
+                if (numa_bitmask_isbitset(membind, i))
+                    *total_memory += numa_node_size(i, nullptr);
+            }
+        }
+
+        numa_bitmask_free(membind);
+    }
+
+#endif
+    return total_memory;
+}
+
+}
diff --git a/base/base/Numa.h b/base/base/Numa.h
new file mode 100644
index 00000000000..b48ab15766b
--- /dev/null
+++ b/base/base/Numa.h
@@ -0,0 +1,12 @@
+#pragma once
+
+#include <optional>
+
+namespace DB
+{
+
+/// return total memory of NUMA nodes the process is bound to
+/// if NUMA is not supported or process can use all nodes, std::nullopt is returned
+std::optional<size_t> getNumaNodesTotalMemory();
+
+}
diff --git a/base/base/getMemoryAmount.cpp b/base/base/getMemoryAmount.cpp
index 56cddbfd628..03aab1eac72 100644
--- a/base/base/getMemoryAmount.cpp
+++ b/base/base/getMemoryAmount.cpp
@@ -2,6 +2,7 @@
 
 #include <base/cgroupsv2.h>
 #include <base/getPageSize.h>
+#include <base/Numa.h>
 
 #include <fstream>
 
@@ -9,13 +10,6 @@
 #include <sys/types.h>
 #include <sys/param.h>
 
-#include "config.h"
-
-#if USE_NUMACTL
-#include <numa.h>
-#endif
-
-
 namespace
 {
 
@@ -68,25 +62,8 @@ uint64_t getMemoryAmountOrZero()
 
     uint64_t memory_amount = num_pages * page_size;
 
-#if USE_NUMACTL
-    if (numa_available() != -1)
-    {
-        auto * membind = numa_get_membind();
-        if (!numa_bitmask_equal(membind, numa_all_nodes_ptr))
-        {
-            uint64_t total_numa_memory = 0;
-            auto max_node = numa_max_node();
-            for (int i = 0; i <= max_node; ++i)
-            {
-                if (numa_bitmask_isbitset(membind, i))
-                    total_numa_memory += numa_node_size(i, nullptr);
-            }
-
-            memory_amount = total_numa_memory;
-        }
-        numa_bitmask_free(membind);
-    }
-#endif
+    if (auto total_numa_memory = DB::getNumaNodesTotalMemory(); total_numa_memory.has_value())
+        memory_amount = *total_numa_memory;
 
     /// Respect the memory limit set by cgroups v2.
     auto limit_v2 = getCgroupsV2MemoryLimit();
diff --git a/programs/keeper/Keeper.cpp b/programs/keeper/Keeper.cpp
index 8cf1a4d1999..a1793b5d603 100644
--- a/programs/keeper/Keeper.cpp
+++ b/programs/keeper/Keeper.cpp
@@ -19,6 +19,7 @@
 #include <base/getMemoryAmount.h>
 #include <base/scope_guard.h>
 #include <base/safeExit.h>
+#include <base/Numa.h>
 #include <Poco/Net/NetException.h>
 #include <Poco/Net/TCPServerParams.h>
 #include <Poco/Net/TCPServer.h>
@@ -311,6 +312,12 @@ try
 
     MainThreadStatus::getInstance();
 
+    if (auto total_numa_memory = getNumaNodesTotalMemory(); total_numa_memory.has_value())
+    {
+        LOG_INFO(
+            log, "ClickHouse is bound to a subset of NUMA nodes. Total memory of all available nodes: {}", ReadableSize(*total_numa_memory));
+    }
+
 #if !defined(NDEBUG) || !defined(__OPTIMIZE__)
     LOG_WARNING(log, "Keeper was built in debug mode. It will work slowly.");
 #endif
diff --git a/programs/server/Server.cpp b/programs/server/Server.cpp
index b818ff1f3a2..86ebd3a3225 100644
--- a/programs/server/Server.cpp
+++ b/programs/server/Server.cpp
@@ -22,6 +22,7 @@
 #include <base/coverage.h>
 #include <base/getFQDNOrHostName.h>
 #include <base/safeExit.h>
+#include <base/Numa.h>
 #include <Common/PoolId.h>
 #include <Common/MemoryTracker.h>
 #include <Common/ClickHouseRevision.h>
@@ -140,10 +141,6 @@
 #   include <azure/core/diagnostics/logger.hpp>
 #endif
 
-#if USE_NUMACTL
-#include <numa.h>
-#endif
-
 
 #include <incbin.h>
 /// A minimal file used when the server is run without installation
@@ -759,36 +756,11 @@ try
         setenv("OPENSSL_CONF", config_dir.c_str(), true); /// NOLINT
     }
 
-#if USE_NUMACTL
-    if (numa_available() != -1)
+    if (auto total_numa_memory = getNumaNodesTotalMemory(); total_numa_memory.has_value())
     {
-        auto * membind = numa_get_membind();
-        if (!numa_bitmask_equal(membind, numa_all_nodes_ptr))
-        {
-            uint64_t total_numa_memory = 0;
-            auto max_node = numa_max_node();
-            for (int i = 0; i <= max_node; ++i)
-            {
-                if (numa_bitmask_isbitset(membind, i))
-                    total_numa_memory += numa_node_size(i, nullptr);
-            }
-
-            LOG_INFO(
-                log,
-                "ClickHouse is bound to a subset of NUMA nodes. Total memory of all available nodes: {}",
-                ReadableSize(total_numa_memory));
-        }
-        else
-        {
-            LOG_TRACE(
-                log,
-                "All NUMA nodes are used. Detected NUMA nodes: {}",
-                numa_num_configured_nodes());
-        }
-
-        numa_bitmask_free(membind);
+        LOG_INFO(
+            log, "ClickHouse is bound to a subset of NUMA nodes. Total memory of all available nodes: {}", ReadableSize(*total_numa_memory));
     }
-#endif
 
     registerInterpreters();
     registerFunctions();