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
27namespace arrow {
28
29namespace internal {
30
31///////////////////////////////////////////////////////////////////////
32// Helper tracking memory statistics
33
34class 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
58class 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.
64class 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
104class 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.
126class 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.
146ARROW_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