1 // Licensed to the Apache Software Foundation (ASF) under one 2 // or more contributor license agreements. See the NOTICE file 3 // distributed with this work for additional information 4 // regarding copyright ownership. The ASF licenses this file 5 // to you under the Apache License, Version 2.0 (the 6 // "License"); you may not use this file except in compliance 7 // with the License. You may obtain a copy of the License at 8 // 9 // http://www.apache.org/licenses/LICENSE-2.0 10 // 11 // Unless required by applicable law or agreed to in writing, 12 // software distributed under the License is distributed on an 13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 // KIND, either express or implied. See the License for the 15 // specific language governing permissions and limitations 16 // under the License. 17 18 #pragma once 19 20 #include <atomic> 21 #include <cstdint> 22 #include <memory> 23 #include <string> 24 25 #include "arrow/status.h" 26 #include "arrow/type_fwd.h" 27 #include "arrow/util/visibility.h" 28 29 namespace arrow { 30 31 namespace internal { 32 33 /////////////////////////////////////////////////////////////////////// 34 // Helper tracking memory statistics 35 36 class MemoryPoolStats { 37 public: MemoryPoolStats()38 MemoryPoolStats() : bytes_allocated_(0), max_memory_(0) {} 39 max_memory()40 int64_t max_memory() const { return max_memory_.load(); } 41 bytes_allocated()42 int64_t bytes_allocated() const { return bytes_allocated_.load(); } 43 UpdateAllocatedBytes(int64_t diff)44 inline void UpdateAllocatedBytes(int64_t diff) { 45 auto allocated = bytes_allocated_.fetch_add(diff) + diff; 46 // "maximum" allocated memory is ill-defined in multi-threaded code, 47 // so don't try to be too rigorous here 48 if (diff > 0 && allocated > max_memory_) { 49 max_memory_ = allocated; 50 } 51 } 52 53 protected: 54 std::atomic<int64_t> bytes_allocated_; 55 std::atomic<int64_t> max_memory_; 56 }; 57 58 } // namespace internal 59 60 /// Base class for memory allocation on the CPU. 61 /// 62 /// Besides tracking the number of allocated bytes, the allocator also should 63 /// take care of the required 64-byte alignment. 64 class ARROW_EXPORT MemoryPool { 65 public: 66 virtual ~MemoryPool() = default; 67 68 /// \brief EXPERIMENTAL. Create a new instance of the default MemoryPool 69 static std::unique_ptr<MemoryPool> CreateDefault(); 70 71 /// Allocate a new memory region of at least size bytes. 72 /// 73 /// The allocated region shall be 64-byte aligned. 74 virtual Status Allocate(int64_t size, uint8_t** out) = 0; 75 76 /// Resize an already allocated memory section. 77 /// 78 /// As by default most default allocators on a platform don't support aligned 79 /// reallocation, this function can involve a copy of the underlying data. 80 virtual Status Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) = 0; 81 82 /// Free an allocated region. 83 /// 84 /// @param buffer Pointer to the start of the allocated memory region 85 /// @param size Allocated size located at buffer. An allocator implementation 86 /// may use this for tracking the amount of allocated bytes as well as for 87 /// faster deallocation if supported by its backend. 88 virtual void Free(uint8_t* buffer, int64_t size) = 0; 89 90 /// Return unused memory to the OS 91 /// 92 /// Only applies to allocators that hold onto unused memory. This will be 93 /// best effort, a memory pool may not implement this feature or may be 94 /// unable to fulfill the request due to fragmentation. ReleaseUnused()95 virtual void ReleaseUnused() {} 96 97 /// The number of bytes that were allocated and not yet free'd through 98 /// this allocator. 99 virtual int64_t bytes_allocated() const = 0; 100 101 /// Return peak memory allocation in this memory pool 102 /// 103 /// \return Maximum bytes allocated. If not known (or not implemented), 104 /// returns -1 105 virtual int64_t max_memory() const; 106 107 /// The name of the backend used by this MemoryPool (e.g. "system" or "jemalloc"). 108 virtual std::string backend_name() const = 0; 109 110 protected: 111 MemoryPool() = default; 112 }; 113 114 class ARROW_EXPORT LoggingMemoryPool : public MemoryPool { 115 public: 116 explicit LoggingMemoryPool(MemoryPool* pool); 117 ~LoggingMemoryPool() override = default; 118 119 Status Allocate(int64_t size, uint8_t** out) override; 120 Status Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) override; 121 122 void Free(uint8_t* buffer, int64_t size) override; 123 124 int64_t bytes_allocated() const override; 125 126 int64_t max_memory() const override; 127 128 std::string backend_name() const override; 129 130 private: 131 MemoryPool* pool_; 132 }; 133 134 /// Derived class for memory allocation. 135 /// 136 /// Tracks the number of bytes and maximum memory allocated through its direct 137 /// calls. Actual allocation is delegated to MemoryPool class. 138 class ARROW_EXPORT ProxyMemoryPool : public MemoryPool { 139 public: 140 explicit ProxyMemoryPool(MemoryPool* pool); 141 ~ProxyMemoryPool() override; 142 143 Status Allocate(int64_t size, uint8_t** out) override; 144 Status Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) override; 145 146 void Free(uint8_t* buffer, int64_t size) override; 147 148 int64_t bytes_allocated() const override; 149 150 int64_t max_memory() const override; 151 152 std::string backend_name() const override; 153 154 private: 155 class ProxyMemoryPoolImpl; 156 std::unique_ptr<ProxyMemoryPoolImpl> impl_; 157 }; 158 159 /// \brief Return a process-wide memory pool based on the system allocator. 160 ARROW_EXPORT MemoryPool* system_memory_pool(); 161 162 /// \brief Return a process-wide memory pool based on jemalloc. 163 /// 164 /// May return NotImplemented if jemalloc is not available. 165 ARROW_EXPORT Status jemalloc_memory_pool(MemoryPool** out); 166 167 /// \brief Set jemalloc memory page purging behavior for future-created arenas 168 /// to the indicated number of milliseconds. See dirty_decay_ms and 169 /// muzzy_decay_ms options in jemalloc for a description of what these do. The 170 /// default is configured to 1000 (1 second) which releases memory more 171 /// aggressively to the operating system than the jemalloc default of 10 172 /// seconds. If you set the value to 0, dirty / muzzy pages will be released 173 /// immediately rather than with a time decay, but this may reduce application 174 /// performance. 175 ARROW_EXPORT 176 Status jemalloc_set_decay_ms(int ms); 177 178 /// \brief Return a process-wide memory pool based on mimalloc. 179 /// 180 /// May return NotImplemented if mimalloc is not available. 181 ARROW_EXPORT Status mimalloc_memory_pool(MemoryPool** out); 182 183 ARROW_EXPORT std::vector<std::string> SupportedMemoryBackendNames(); 184 185 } // namespace arrow 186