track IO layer read buffers via MemTrackerLimiter using PODArray

Author: sollhui

be/src/io/fs/buffered_reader.cpp

@@ -449,6 +449,14 @@ void PrefetchBuffer::prefetch_buffer() {
         _prefetched.notify_all();
     }
 
+    // Lazy-allocate the backing buffer on first actual prefetch, avoiding the cost of
+    // pre-allocating memory for readers that are initialized but never read (e.g. when
+    // many file readers are created concurrently for a TVF scan over many small S3 files).
+    if (_buf.empty()) {
+        _buf.resize(_size);
+    }
+
+
     int read_range_index = search_read_range(_offset);
     size_t buf_size;
     if (read_range_index == -1) {
@@ -463,7 +471,7 @@ void PrefetchBuffer::prefetch_buffer() {
 
     {
         SCOPED_RAW_TIMER(&_statis.read_time);
-        s = _reader->read_at(_offset, Slice {_buf.get(), buf_size}, &_len, _io_ctx);
+        s = _reader->read_at(_offset, Slice {_buf.data(), buf_size}, &_len, _io_ctx);
     }
     if (UNLIKELY(s.ok() && buf_size != _len)) {
         // This indicates that the data size returned by S3 object storage is smaller than what we requested,
@@ -602,7 +610,7 @@ Status PrefetchBuffer::read_buffer(size_t off, const char* out, size_t buf_len,
         size_t read_len = std::min({buf_len, _offset + _size - off, _offset + _len - off});
         {
             SCOPED_RAW_TIMER(&_statis.copy_time);
-            memcpy((void*)out, _buf.get() + (off - _offset), read_len);
+            memcpy((void*)out, _buf.data() + (off - _offset), read_len);
         }
         *bytes_read = read_len;
         _statis.request_io += 1;

be/src/io/fs/buffered_reader.h

@@ -29,6 +29,7 @@
 
 #include "common/status.h"
 #include "core/custom_allocator.h"
+#include "core/pod_array.h"
 #include "core/typeid_cast.h"
 #include "io/cache/cached_remote_file_reader.h"
 #include "io/file_factory.h"
@@ -436,7 +437,6 @@ struct PrefetchBuffer : std::enable_shared_from_this<PrefetchBuffer>, public Pro
               _reader(reader),
               _io_ctx_holder(std::move(io_ctx)),
               _io_ctx(_io_ctx_holder.get()),
-              _buf(new char[buffer_size]),
               _sync_profile(std::move(sync_profile)) {}
 
     PrefetchBuffer(PrefetchBuffer&& other)
@@ -465,7 +465,7 @@ struct PrefetchBuffer : std::enable_shared_from_this<PrefetchBuffer>, public Pro
     io::FileReader* _reader = nullptr;
     std::shared_ptr<const IOContext> _io_ctx_holder;
     const IOContext* _io_ctx = nullptr;
-    std::unique_ptr<char[]> _buf;
+    PODArray<char> _buf;
     BufferStatus _buffer_status {BufferStatus::RESET};
     std::mutex _lock;
     std::condition_variable _prefetched;

be/src/io/fs/hdfs_file_system.cpp

@@ -31,6 +31,7 @@
 
 #include "common/config.h"
 #include "common/status.h"
+#include "core/pod_array.h"
 #include "io/fs/err_utils.h"
 #include "io/fs/hdfs/hdfs_mgr.h"
 #include "io/fs/hdfs_file_reader.h"
@@ -296,18 +297,19 @@ Status HdfsFileSystem::download_impl(const Path& remote_file, const Path& local_
     // 4. read remote and write to local
     LOG(INFO) << "read remote file: " << remote_file << " to local: " << local_file;
     constexpr size_t buf_sz = 1024 * 1024;
-    std::unique_ptr<char[]> read_buf(new char[buf_sz]);
+    PODArray<char> read_buf;
+    read_buf.resize(buf_sz);
     size_t cur_offset = 0;
     while (true) {
         size_t read_len = 0;
-        Slice file_slice(read_buf.get(), buf_sz);
+        Slice file_slice(read_buf.data(), buf_sz);
         RETURN_IF_ERROR(hdfs_reader->read_at(cur_offset, file_slice, &read_len));
         cur_offset += read_len;
         if (read_len == 0) {
             break;
         }
 
-        RETURN_IF_ERROR(local_writer->append({read_buf.get(), read_len}));
+        RETURN_IF_ERROR(local_writer->append({read_buf.data(), read_len}));
     }
     return local_writer->close();
 }

be/src/io/fs/http_file_reader.cpp

@@ -84,7 +84,7 @@ HttpFileReader::HttpFileReader(const OpenFileInfo& fileInfo, std::string url, in
         }
     }
 
-    _read_buffer = std::make_unique<char[]>(READ_BUFFER_SIZE);
+    _read_buffer.resize(READ_BUFFER_SIZE);
 }
 
 HttpFileReader::~HttpFileReader() {
@@ -154,8 +154,8 @@ Status HttpFileReader::read_at_impl(size_t offset, Slice result, size_t* bytes_r
     VLOG(2) << "HttpFileReader::read_at_impl offset=" << offset << " size=" << result.size
             << " url=" << _url << " range_supported=" << _range_supported;
 
-    if (!_read_buffer) {
-        _read_buffer = std::make_unique<char[]>(READ_BUFFER_SIZE);
+    if (_read_buffer.empty()) {
+        _read_buffer.resize(READ_BUFFER_SIZE);
     }
 
     size_t to_read = result.size;
@@ -192,7 +192,7 @@ Status HttpFileReader::read_at_impl(size_t offset, Slice result, size_t* bytes_r
                 << "Buffer overflow: buffer_idx=" << buffer_idx << " copy_len=" << copy_len
                 << " READ_BUFFER_SIZE=" << READ_BUFFER_SIZE;
 
-        std::memcpy(result.data, _read_buffer.get() + buffer_idx, copy_len);
+        std::memcpy(result.data, _read_buffer.data() + buffer_idx, copy_len);
         buffer_offset += copy_len;
         to_read -= copy_len;
         offset += copy_len;
@@ -319,12 +319,12 @@ Status HttpFileReader::read_at_impl(size_t offset, Slice result, size_t* bytes_r
         buffer_offset += buf.size();
     } else {
         size_t cached = std::min(buf.size(), (size_t)READ_BUFFER_SIZE);
-        std::memcpy(_read_buffer.get(), buf.data(), cached);
+        std::memcpy(_read_buffer.data(), buf.data(), cached);
         _buffer_start = offset;
         _buffer_end = offset + cached;
 
         size_t copy_len = std::min(remaining, cached);
-        std::memcpy(result.data + buffer_offset, _read_buffer.get(), copy_len);
+        std::memcpy(result.data + buffer_offset, _read_buffer.data(), copy_len);
         buffer_offset += copy_len;
     }
 

be/src/io/fs/http_file_reader.h

@@ -23,6 +23,7 @@
 #include <string>
 
 #include "common/status.h"
+#include "core/pod_array.h"
 #include "io/fs/file_handle_cache.h"
 #include "io/fs/file_reader.h"
 #include "io/fs/file_system.h"
@@ -62,7 +63,7 @@ class HttpFileReader final : public FileReader {
     // Returns OK on success with _range_supported set appropriately
     Status detect_range_support();
 
-    std::unique_ptr<char[]> _read_buffer;
+    PODArray<char> _read_buffer;
     static constexpr size_t READ_BUFFER_SIZE = 1 << 20; // 1MB
     // Default maximum file size for servers that don't support Range requests
     static constexpr size_t DEFAULT_MAX_REQUEST_SIZE = 100 << 20; // 100MB