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
30namespace arrow {
31
32/// \brief A STL allocator delegating allocations to a Arrow MemoryPool
33template <class T>
34class 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.
96template <typename Allocator = std::allocator<uint8_t>>
97class 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
139template <class T1, class T2>
140bool operator==(const stl_allocator<T1>& lhs, const stl_allocator<T2>& rhs) noexcept {
141 return lhs.pool() == rhs.pool();
142}
143
144template <class T1, class T2>
145bool 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