Merge pull request #23211 from azat/mmap-accounting

[RFC] Fix memory tracking with min_bytes_to_use_mmap_io

src/IO/MMapReadBufferFromFileWithCache.cpp

@@ -1,4 +1,5 @@
 #include <IO/MMapReadBufferFromFileWithCache.h>
+#include <Common/CurrentMemoryTracker.h>
 
 
 namespace DB
@@ -74,4 +75,16 @@ off_t MMapReadBufferFromFileWithCache::seek(off_t offset, int whence)
     return new_pos;
 }
 
+void MMapReadBufferFromFileWithCache::track(size_t bytes_)
+{
+    CurrentMemoryTracker::alloc(bytes_);
+    tracked_bytes = bytes_;
+}
+
+MMapReadBufferFromFileWithCache::~MMapReadBufferFromFileWithCache()
+{
+    if (tracked_bytes)
+        CurrentMemoryTracker::free(tracked_bytes);
+}
+
 }

src/IO/MMapReadBufferFromFileWithCache.h

@@ -16,13 +16,27 @@ public:
     /// Map till end of file.
     MMapReadBufferFromFileWithCache(MMappedFileCache & cache, const std::string & file_name, size_t offset);
 
+    ~MMapReadBufferFromFileWithCache() override;
+
     off_t getPosition() override;
     std::string getFileName() const override;
     off_t seek(off_t offset, int whence) override;
 
+    /// Track memory with MemoryTracker.
+    ///
+    /// NOTE: that this is not the size of mmap() region, but estimated size of
+    /// data that will be read (see MergeTreeReaderStream).
+    /// And from one point of view it should not be accounted here, since the
+    /// kernel may unmap physical pages, but in practice firstly it will mmap it,
+    /// RSS will grow, total_memory_tracker will be synced with RSS and the
+    /// allocations will start to fail.
+    void track(size_t bytes_);
+
 private:
     MMappedFileCache::MappedPtr mapped;
 
+    size_t tracked_bytes = 0;
+
     void init();
 };
 

src/IO/createReadBufferFromFileBase.cpp

@@ -50,6 +50,7 @@ std::unique_ptr<ReadBufferFromFileBase> createReadBufferFromFileBase(
         try
         {
             auto res = std::make_unique<MMapReadBufferFromFileWithCache>(*mmap_cache, filename_, 0);
+            res->track(estimated_size);
             ProfileEvents::increment(ProfileEvents::CreatedReadBufferMMap);
             return res;
         }

tests/queries/0_stateless/01605_adaptive_granularity_block_borders_long.sql

@@ -23,6 +23,6 @@ SELECT marks FROM system.parts WHERE table = 'adaptive_table' and database=curre
 -- If we have computed granularity incorrectly than we will exceed this limit.
 SET max_memory_usage='30M';
 
-SELECT max(length(value)) FROM adaptive_table;
+SELECT max(length(value)) FROM adaptive_table SETTINGS min_bytes_to_use_mmap_io=0;
 
 DROP TABLE IF EXISTS adaptive_table;

tests/queries/0_stateless/01818_min_bytes_to_use_mmap_io_memory_tracking.sql

@@ -0,0 +1,5 @@
+drop table if exists data_01818;
+
+create table data_01818 (key Int, value String) engine=MergeTree() order by key settings min_bytes_for_wide_part=0 as select number, randomPrintableASCII(100) from numbers(1e6);
+select * from data_01818 format Null settings min_bytes_to_use_mmap_io=1, max_memory_usage='20Mi', max_threads=1; -- { serverError 241 }
+select * from data_01818 format Null settings min_bytes_to_use_mmap_io=1e9, max_memory_usage='20Mi', max_threads=1;

tests/queries/1_stateful/00084_external_aggregation.sql

@@ -2,6 +2,8 @@ SET max_bytes_before_external_group_by = 200000000;
 
 SET max_memory_usage = 1500000000;
 SET max_threads = 12;
+SET min_bytes_to_use_mmap_io = 0;
+
 SELECT URL, uniq(SearchPhrase) AS u FROM test.hits GROUP BY URL ORDER BY u DESC, URL LIMIT 10;
 
 SET max_memory_usage = 300000000;