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_ALLOCATOR_H |
19 | #define ARROW_ALLOCATOR_H |
20 | |
21 | #include <algorithm> |
22 | #include <cstddef> |
23 | #include <memory> |
24 | #include <utility> |
25 | |
26 | #include "arrow/memory_pool.h" |
27 | #include "arrow/status.h" |
28 | #include "arrow/util/macros.h" |
29 | |
30 | namespace arrow { |
31 | |
32 | /// \brief A STL allocator delegating allocations to a Arrow MemoryPool |
33 | template <class T> |
34 | class stl_allocator { |
35 | public: |
36 | using value_type = T; |
37 | using pointer = T*; |
38 | using const_pointer = const T*; |
39 | using reference = T&; |
40 | using const_reference = const T&; |
41 | using size_type = std::size_t; |
42 | using difference_type = std::ptrdiff_t; |
43 | |
44 | template <class U> |
45 | struct rebind { |
46 | using other = stl_allocator<U>; |
47 | }; |
48 | |
49 | /// \brief Construct an allocator from the default MemoryPool |
50 | stl_allocator() noexcept : pool_(default_memory_pool()) {} |
51 | /// \brief Construct an allocator from the given MemoryPool |
52 | explicit stl_allocator(MemoryPool* pool) noexcept : pool_(pool) {} |
53 | |
54 | template <class U> |
55 | stl_allocator(const stl_allocator<U>& rhs) noexcept : pool_(rhs.pool_) {} |
56 | |
57 | ~stl_allocator() { pool_ = NULLPTR; } |
58 | |
59 | pointer address(reference r) const noexcept { return std::addressof(r); } |
60 | |
61 | const_pointer address(const_reference r) const noexcept { return std::addressof(r); } |
62 | |
63 | pointer allocate(size_type n, const void* /*hint*/ = NULLPTR) { |
64 | uint8_t* data; |
65 | Status s = pool_->Allocate(n * sizeof(T), &data); |
66 | if (!s.ok()) throw std::bad_alloc(); |
67 | return reinterpret_cast<pointer>(data); |
68 | } |
69 | |
70 | void deallocate(pointer p, size_type n) { |
71 | pool_->Free(reinterpret_cast<uint8_t*>(p), n * sizeof(T)); |
72 | } |
73 | |
74 | size_type size_max() const noexcept { return size_type(-1) / sizeof(T); } |
75 | |
76 | template <class U, class... Args> |
77 | void construct(U* p, Args&&... args) { |
78 | new (reinterpret_cast<void*>(p)) U(std::forward<Args>(args)...); |
79 | } |
80 | |
81 | template <class U> |
82 | void destroy(U* p) { |
83 | p->~U(); |
84 | } |
85 | |
86 | MemoryPool* pool() const noexcept { return pool_; } |
87 | |
88 | private: |
89 | MemoryPool* pool_; |
90 | }; |
91 | |
92 | /// \brief A MemoryPool implementation delegating allocations to a STL allocator |
93 | /// |
94 | /// Note that STL allocators don't provide a resizing operation, and therefore |
95 | /// any buffer resizes will do a full reallocation and copy. |
96 | template <typename Allocator = std::allocator<uint8_t>> |
97 | class STLMemoryPool : public MemoryPool { |
98 | public: |
99 | /// \brief Construct a memory pool from the given allocator |
100 | explicit STLMemoryPool(const Allocator& alloc) : alloc_(alloc) {} |
101 | |
102 | Status Allocate(int64_t size, uint8_t** out) override { |
103 | try { |
104 | *out = alloc_.allocate(size); |
105 | } catch (std::bad_alloc& e) { |
106 | return Status::OutOfMemory(e.what()); |
107 | } |
108 | stats_.UpdateAllocatedBytes(size); |
109 | return Status::OK(); |
110 | } |
111 | |
112 | Status Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) override { |
113 | uint8_t* old_ptr = *ptr; |
114 | try { |
115 | *ptr = alloc_.allocate(new_size); |
116 | } catch (std::bad_alloc& e) { |
117 | return Status::OutOfMemory(e.what()); |
118 | } |
119 | memcpy(*ptr, old_ptr, std::min(old_size, new_size)); |
120 | alloc_.deallocate(old_ptr, old_size); |
121 | stats_.UpdateAllocatedBytes(new_size - old_size); |
122 | return Status::OK(); |
123 | } |
124 | |
125 | void Free(uint8_t* buffer, int64_t size) override { |
126 | alloc_.deallocate(buffer, size); |
127 | stats_.UpdateAllocatedBytes(-size); |
128 | } |
129 | |
130 | int64_t bytes_allocated() const override { return stats_.bytes_allocated(); } |
131 | |
132 | int64_t max_memory() const override { return stats_.max_memory(); } |
133 | |
134 | private: |
135 | Allocator alloc_; |
136 | internal::MemoryPoolStats stats_; |
137 | }; |
138 | |
139 | template <class T1, class T2> |
140 | bool operator==(const stl_allocator<T1>& lhs, const stl_allocator<T2>& rhs) noexcept { |
141 | return lhs.pool() == rhs.pool(); |
142 | } |
143 | |
144 | template <class T1, class T2> |
145 | bool operator!=(const stl_allocator<T1>& lhs, const stl_allocator<T2>& rhs) noexcept { |
146 | return !(lhs == rhs); |
147 | } |
148 | |
149 | } // namespace arrow |
150 | |
151 | #endif // ARROW_ALLOCATOR_H |
152 | |