MemoryWorker: use sanitizer info for allocated counter

src/Common/MemoryWorker.cpp

@@ -15,13 +15,19 @@
 #include <Common/logger_useful.h>
 #include <Common/setThreadName.h>
 
+#include <base/defines.h>
+
 #include <fmt/ranges.h>
 
 #include <filesystem>
 #include <optional>
 
 #include <unistd.h>
 
+#if defined(ADDRESS_SANITIZER) || defined(THREAD_SANITIZER) || defined(MEMORY_SANITIZER)
+#include <sanitizer/allocator_interface.h>
+#endif
+
 namespace fs = std::filesystem;
 
 namespace ProfileEvents
@@ -539,30 +545,46 @@ MemoryWorker::~MemoryWorker()
 #endif
 }
 
-uint64_t MemoryWorker::getMemoryUsage(bool log_error)
+MemoryWorker::MemoryUsage MemoryWorker::getMemoryUsage(bool log_error)
 {
+    MemoryUsage usage;
+
     switch (source)
     {
         case MemoryUsageSource::Cgroups:
         {
             if (cgroups_reader != nullptr)
-                return cgroups_reader->readMemoryUsage();
+            {
+                usage.resident = cgroups_reader->readMemoryUsage();
+                break;
+            }
             [[fallthrough]];
         }
         case MemoryUsageSource::Jemalloc:
 #if USE_JEMALLOC
             epoch_mib.setValue(0);
-            return resident_mib.getValue();
+            usage.resident = resident_mib.getValue();
+            break;
 #else
             [[fallthrough]];
 #endif
         case MemoryUsageSource::None:
         {
             if (log_error)
                 LOG_ERROR(log, "Trying to fetch memory usage while no memory source can be used");
-            return 0;
+            break;
         }
     }
+
+#if defined(ADDRESS_SANITIZER) || defined(THREAD_SANITIZER) || defined(MEMORY_SANITIZER)
+    /// Sanitizer memory overhead (redzones, ...) makes RSS exceed application allocations;
+    // use allocator bytes for MEMORY_LIMIT_EXCEEDED, resident for RSS/cgroup limit sizing.
+    usage.allocated = __sanitizer_get_current_allocated_bytes();
+#else
+    usage.allocated = usage.resident;
+#endif
+
+    return usage;
 }
 
 namespace
@@ -839,19 +861,19 @@ void MemoryWorker::updateResidentMemoryThread()
 
             Stopwatch total_watch;
 
-            Int64 resident = getMemoryUsage(first_run);
-            MemoryTracker::updateRSS(resident);
+            const MemoryUsage memory_usage = getMemoryUsage(first_run);
+            MemoryTracker::updateRSS(memory_usage.resident);
 
             if (page_cache)
-                page_cache->autoResize(std::max(resident, total_memory_tracker.get()), total_memory_tracker.getHardLimit());
+                page_cache->autoResize(std::max(memory_usage.resident, total_memory_tracker.get()), total_memory_tracker.getHardLimit());
 
 #if USE_JEMALLOC
             const auto memory_tracker_limit = total_memory_tracker.getHardLimit();
             const auto purge_total_memory_threshold = static_cast<double>(memory_tracker_limit) * purge_total_memory_threshold_ratio;
             const auto purge_dirty_pages_threshold = static_cast<double>(memory_tracker_limit) * purge_dirty_pages_threshold_ratio;
 
             const bool needs_purge
-                = (purge_total_memory_threshold_ratio > 0 && static_cast<double>(resident) > purge_total_memory_threshold)
+                = (purge_total_memory_threshold_ratio > 0 && static_cast<double>(memory_usage.resident) > purge_total_memory_threshold)
                 || (purge_dirty_pages_threshold_ratio > 0
                     && static_cast<double>(pdirty_mib.getValue() * page_size) > purge_dirty_pages_threshold);
 
@@ -907,22 +929,23 @@ void MemoryWorker::updateResidentMemoryThread()
                 }
             }
 
-            /// update MemoryTracker with `allocated` information from jemalloc when:
+            /// update MemoryTracker with resident memory information (cgroup or jemalloc) when:
             ///  - it's a first run of MemoryWorker (MemoryTracker could've missed some allocation before its initialization)
             ///  - MemoryTracker stores a negative value
             ///  - `correct_tracker` is set to true
             if (first_run || total_memory_tracker.get() < 0) [[unlikely]]
-                MemoryTracker::updateAllocated(resident, /*log_change=*/true);
+                MemoryTracker::updateAllocated(memory_usage.allocated, /*log_change=*/true);
             else if (correct_tracker)
-                MemoryTracker::updateAllocated(resident, /*log_change=*/false);
+                MemoryTracker::updateAllocated(memory_usage.allocated, /*log_change=*/false);
 #else
             /// we don't update in the first run if we don't have jemalloc
-            /// because we can only use resident memory information
+            /// because without a sanitizer we can only use resident memory information
             /// resident memory can be much larger than the actual allocated memory
             /// so we rather ignore the potential difference caused by allocated memory
             /// before MemoryTracker initialization
+            /// sanitizer builds provide allocated memory, but keep the same behavior
             if (total_memory_tracker.get() < 0 || correct_tracker) [[unlikely]]
-                MemoryTracker::updateAllocated(resident, /*log_change=*/false);
+                MemoryTracker::updateAllocated(memory_usage.allocated, /*log_change=*/false);
 #endif
 
             /// Capture the settings generation before reading ratio/ceiling. We re-read
@@ -955,7 +978,7 @@ void MemoryWorker::updateResidentMemoryThread()
                         /// are excluded. Under load `tracked` can be orders of magnitude smaller
                         /// than the actual RSS, which makes `(tracked + available) * ratio` compute
                         /// a hard limit close to current RSS and reject every subsequent allocation.
-                        Int64 used = std::max<Int64>(0, resident);
+                        Int64 used = std::max<Int64>(0, memory_usage.resident);
                         /// `used + available` is the upper bound of memory we could potentially own:
                         /// what we already use plus what is still free in our cgroup (or on the host).
                         /// Scaling by `ratio < 1` leaves headroom for other processes on the host.

src/Common/MemoryWorker.h

@@ -132,7 +132,13 @@ class MemoryWorker
 
     ~MemoryWorker();
 private:
-    uint64_t getMemoryUsage(bool log_error);
+    struct MemoryUsage
+    {
+        Int64 resident = 0;
+        Int64 allocated = 0;
+    };
+
+    MemoryUsage getMemoryUsage(bool log_error);
 
     void updateResidentMemoryThread();