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 | #ifndef ARROW_MEMORY_POOL_H |
19 | #define ARROW_MEMORY_POOL_H |
20 | |
21 | #include <atomic> |
22 | #include <cstdint> |
23 | #include <memory> |
24 | |
25 | #include "arrow/util/visibility.h" |
26 | |
27 | namespace arrow { |
28 | |
29 | namespace internal { |
30 | |
31 | /////////////////////////////////////////////////////////////////////// |
32 | // Helper tracking memory statistics |
33 | |
34 | class MemoryPoolStats { |
35 | public: |
36 | MemoryPoolStats() : bytes_allocated_(0), max_memory_(0) {} |
37 | |
38 | int64_t max_memory() const { return max_memory_.load(); } |
39 | |
40 | int64_t bytes_allocated() const { return bytes_allocated_.load(); } |
41 | |
42 | inline void UpdateAllocatedBytes(int64_t diff) { |
43 | auto allocated = bytes_allocated_.fetch_add(diff) + diff; |
44 | // "maximum" allocated memory is ill-defined in multi-threaded code, |
45 | // so don't try to be too rigorous here |
46 | if (diff > 0 && allocated > max_memory_) { |
47 | max_memory_ = allocated; |
48 | } |
49 | } |
50 | |
51 | protected: |
52 | std::atomic<int64_t> bytes_allocated_; |
53 | std::atomic<int64_t> max_memory_; |
54 | }; |
55 | |
56 | } // namespace internal |
57 | |
58 | class Status; |
59 | |
60 | /// Base class for memory allocation. |
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(); |
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 | /// The number of bytes that were allocated and not yet free'd through |
91 | /// this allocator. |
92 | virtual int64_t bytes_allocated() const = 0; |
93 | |
94 | /// Return peak memory allocation in this memory pool |
95 | /// |
96 | /// \return Maximum bytes allocated. If not known (or not implemented), |
97 | /// returns -1 |
98 | virtual int64_t max_memory() const; |
99 | |
100 | protected: |
101 | MemoryPool(); |
102 | }; |
103 | |
104 | class ARROW_EXPORT LoggingMemoryPool : public MemoryPool { |
105 | public: |
106 | explicit LoggingMemoryPool(MemoryPool* pool); |
107 | ~LoggingMemoryPool() override = default; |
108 | |
109 | Status Allocate(int64_t size, uint8_t** out) override; |
110 | Status Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) override; |
111 | |
112 | void Free(uint8_t* buffer, int64_t size) override; |
113 | |
114 | int64_t bytes_allocated() const override; |
115 | |
116 | int64_t max_memory() const override; |
117 | |
118 | private: |
119 | MemoryPool* pool_; |
120 | }; |
121 | |
122 | /// Derived class for memory allocation. |
123 | /// |
124 | /// Tracks the number of bytes and maximum memory allocated through its direct |
125 | /// calls. Actual allocation is delegated to MemoryPool class. |
126 | class ARROW_EXPORT ProxyMemoryPool : public MemoryPool { |
127 | public: |
128 | explicit ProxyMemoryPool(MemoryPool* pool); |
129 | ~ProxyMemoryPool() override; |
130 | |
131 | Status Allocate(int64_t size, uint8_t** out) override; |
132 | Status Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) override; |
133 | |
134 | void Free(uint8_t* buffer, int64_t size) override; |
135 | |
136 | int64_t bytes_allocated() const override; |
137 | |
138 | int64_t max_memory() const override; |
139 | |
140 | private: |
141 | class ProxyMemoryPoolImpl; |
142 | std::unique_ptr<ProxyMemoryPoolImpl> impl_; |
143 | }; |
144 | |
145 | /// Return the process-wide default memory pool. |
146 | ARROW_EXPORT MemoryPool* default_memory_pool(); |
147 | |
148 | #ifdef ARROW_NO_DEFAULT_MEMORY_POOL |
149 | #define ARROW_MEMORY_POOL_DEFAULT |
150 | #else |
151 | #define ARROW_MEMORY_POOL_DEFAULT = default_memory_pool() |
152 | #endif |
153 | |
154 | } // namespace arrow |
155 | |
156 | #endif // ARROW_MEMORY_POOL_H |
157 | |