]>
Commit | Line | Data |
---|---|---|
1d09f67e TL |
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: | |
38 | MemoryPoolStats() : bytes_allocated_(0), max_memory_(0) {} | |
39 | ||
40 | int64_t max_memory() const { return max_memory_.load(); } | |
41 | ||
42 | int64_t bytes_allocated() const { return bytes_allocated_.load(); } | |
43 | ||
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. | |
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 |