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