Use Arena allocator

This commit is contained in:
2024-01-18 17:26:55 -08:00
parent 5f3313a3bf
commit fc3cc98d64
2 changed files with 201 additions and 25 deletions

View File

@@ -20,6 +20,172 @@ static auto operator<=>(const Key &lhs, const Key &rhs) {
return c != 0 ? c <=> 0 : lhs.len <=> rhs.len;
}
// ==================== BEGIN ARENA IMPL ====================
/// Group allocations with similar lifetimes to amortize the cost of malloc/free
struct Arena {
explicit Arena(int initialSize = 0);
/// O(log n) in the number of allocations
~Arena();
struct ArenaImpl;
Arena(const Arena &) = delete;
Arena &operator=(const Arena &) = delete;
Arena(Arena &&other) noexcept;
Arena &operator=(Arena &&other) noexcept;
private:
ArenaImpl *impl = nullptr;
friend void *operator new(size_t size, std::align_val_t align, Arena &arena);
};
inline void operator delete(void *, std::align_val_t, Arena &) {}
void *operator new(size_t size, std::align_val_t align, Arena &arena);
void *operator new(size_t size, std::align_val_t align, Arena *arena) = delete;
inline void operator delete(void *, Arena &) {}
inline void *operator new(size_t size, Arena &arena) {
return operator new(size, std::align_val_t(alignof(std::max_align_t)), arena);
}
inline void *operator new(size_t size, Arena *arena) = delete;
inline void operator delete[](void *, Arena &) {}
inline void *operator new[](size_t size, Arena &arena) {
return operator new(size, arena);
}
inline void *operator new[](size_t size, Arena *arena) = delete;
inline void operator delete[](void *, std::align_val_t, Arena &) {}
inline void *operator new[](size_t size, std::align_val_t align, Arena &arena) {
return operator new(size, align, arena);
}
inline void *operator new[](size_t size, std::align_val_t align,
Arena *arena) = delete;
/// align must be a power of two
template <class T> T *align_up(T *t, size_t align) {
auto unaligned = uintptr_t(t);
auto aligned = (unaligned + align - 1) & ~(align - 1);
return reinterpret_cast<T *>(reinterpret_cast<char *>(t) + aligned -
unaligned);
}
/// align must be a power of two
constexpr inline int align_up(uint32_t unaligned, uint32_t align) {
return (unaligned + align - 1) & ~(align - 1);
}
/// Returns the smallest power of two >= x
constexpr inline uint32_t nextPowerOfTwo(uint32_t x) {
return x <= 1 ? 1 : 1 << (32 - __builtin_clz(x - 1));
}
/// \private
struct Arena::ArenaImpl {
Arena::ArenaImpl *prev;
int capacity;
int used;
uint8_t *begin() { return reinterpret_cast<uint8_t *>(this + 1); }
};
static_assert(sizeof(Arena::ArenaImpl) == 16);
static_assert(alignof(Arena::ArenaImpl) == 8);
Arena::Arena(int initialSize) : impl(nullptr) {
if (initialSize > 0) {
auto allocationSize = align_up(initialSize + sizeof(ArenaImpl), 16);
impl = (Arena::ArenaImpl *)malloc(allocationSize);
impl->prev = nullptr;
impl->capacity = allocationSize - sizeof(ArenaImpl);
impl->used = 0;
}
}
namespace {
void onDestroy(Arena::ArenaImpl *impl) {
while (impl) {
auto *prev = impl->prev;
free(impl);
impl = prev;
}
}
} // namespace
Arena::Arena(Arena &&other) noexcept
: impl(std::exchange(other.impl, nullptr)) {}
Arena &Arena::operator=(Arena &&other) noexcept {
onDestroy(impl);
impl = std::exchange(other.impl, nullptr);
return *this;
}
Arena::~Arena() { onDestroy(impl); }
void *operator new(size_t size, std::align_val_t align, Arena &arena) {
int64_t aligned_size = size + size_t(align) - 1;
if (arena.impl == nullptr ||
(arena.impl->capacity - arena.impl->used) < aligned_size) {
auto allocationSize = align_up(
sizeof(Arena::ArenaImpl) +
std::max<int>(aligned_size,
(arena.impl ? std::max<int>(sizeof(Arena::ArenaImpl),
arena.impl->capacity * 2)
: 0)),
16);
auto *impl = (Arena::ArenaImpl *)malloc(allocationSize);
impl->prev = arena.impl;
impl->capacity = allocationSize - sizeof(Arena::ArenaImpl);
impl->used = 0;
arena.impl = impl;
}
auto *result =
align_up(arena.impl->begin() + arena.impl->used, size_t(align));
auto usedDelta = (result - arena.impl->begin()) + size - arena.impl->used;
arena.impl->used += usedDelta;
return result;
}
/// STL-friendly allocator using an arena
template <class T> struct ArenaAlloc {
typedef T value_type;
ArenaAlloc() = delete;
explicit ArenaAlloc(Arena *arena) : arena(arena) {}
Arena *arena;
template <class U> constexpr ArenaAlloc(const ArenaAlloc<U> &other) noexcept {
arena = other.arena;
}
[[nodiscard]] T *allocate(size_t n) {
if (n > 0xfffffffffffffffful / sizeof(T)) { // NOLINT
fprintf(stderr, "Requested bad alloc! sizeof(T): %zu, n: %zu\n",
sizeof(T), n); // NOLINT
fflush(stderr);
abort();
}
return static_cast<T *>((void *)new (std::align_val_t(alignof(T)), *arena)
uint8_t[n * sizeof(T)]); // NOLINT
}
void deallocate(T *, size_t) noexcept {}
private:
};
template <class T, class U>
bool operator==(const ArenaAlloc<T> &lhs, const ArenaAlloc<U> &rhs) {
return lhs.arena == rhs.arena;
}
template <class T, class U>
bool operator!=(const ArenaAlloc<T> &lhs, const ArenaAlloc<U> &rhs) {
return !(lhs == rhs);
}
// ==================== END ARENA IMPL ====================
namespace {
// A node in the tree representing write conflict history. This tree maintains
// several invariants:
@@ -171,15 +337,15 @@ struct StepwiseLastLeq {
}
};
void lastLeqMulti(Node *root, std::span<Key> keys, Iterator *results) {
void lastLeqMulti(Arena &arena, Node *root, std::span<Key> keys,
Iterator *results) {
assert(std::is_sorted(keys.begin(), keys.end()));
if (keys.size() == 0) {
return;
}
auto stepwiseLastLeqs =
std::unique_ptr<StepwiseLastLeq[]>{new StepwiseLastLeq[keys.size()]};
auto *stepwiseLastLeqs = new (arena) StepwiseLastLeq[keys.size()];
// Descend until queries for front and back diverge
Node *current = root;
@@ -201,13 +367,12 @@ void lastLeqMulti(Node *root, std::span<Key> keys, Iterator *results) {
int index = 0;
{
auto iter = stepwiseLastLeqs.get();
auto iter = stepwiseLastLeqs;
for (const auto &k : keys) {
*iter++ = StepwiseLastLeq(current, resultP, k, index++);
}
}
auto stepwiseSpan =
std::span<StepwiseLastLeq>(stepwiseLastLeqs.get(), keys.size());
auto stepwiseSpan = std::span<StepwiseLastLeq>(stepwiseLastLeqs, keys.size());
runInterleaved(stepwiseSpan);
for (const auto &stepwise : stepwiseSpan) {
results[stepwise.index] = Iterator{stepwise.result, stepwise.resultC};
@@ -300,14 +465,12 @@ void lastLeqMulti(Node *root, std::span<Key> keys, Iterator *results) {
fprintf(file, "}\n");
}
[[maybe_unused]] Key toKey(int n) {
[[maybe_unused]] Key toKey(Arena &arena, int n) {
constexpr int kMaxLength = 4;
// TODO use arena allocation
static std::vector<std::vector<uint8_t>> *results =
new std::vector<std::vector<uint8_t>>;
int i = kMaxLength;
results->push_back(std::vector<uint8_t>(kMaxLength, '0'));
uint8_t *itoaBuf = results->back().data();
uint8_t *itoaBuf = new (arena) uint8_t[kMaxLength];
memset(itoaBuf, '0', kMaxLength);
do {
itoaBuf[--i] = "0123456789abcdef"[n % 16];
n /= 16;
@@ -401,7 +564,9 @@ int64_t checkMaxVersion(Node *node, bool &success) {
bool checkInvariants(Node *node) {
bool success = true;
// Check bst invariant
std::vector<std::string_view> keys;
Arena arena;
std::vector<std::string_view, ArenaAlloc<std::string_view>> keys{
ArenaAlloc<std::string_view>(&arena)};
for (auto iter = extrema(node, false); iter != nullptr;
iter = next(iter, true)) {
keys.push_back(std::string_view((char *)(iter + 1), iter->len));
@@ -438,12 +603,13 @@ struct ConflictSet::Impl {
}
void check(const ReadRange *reads, Result *results, int count) const {
auto iters = std::unique_ptr<Iterator[]>{new Iterator[count]};
auto begins = std::unique_ptr<Key[]>{new Key[count]};
Arena arena;
auto *iters = new (arena) Iterator[count];
auto *begins = new (arena) Key[count];
for (int i = 0; i < count; ++i) {
begins.get()[i] = reads[i].begin;
begins[i] = reads[i].begin;
}
lastLeqMulti(root, std::span<Key>(begins.get(), count), iters.get());
lastLeqMulti(arena, root, std::span<Key>(begins, count), iters);
// TODO check non-singleton reads lol
for (int i = 0; i < count; ++i) {
assert(reads[i].end.len == 0);
@@ -501,8 +667,8 @@ struct ConflictSet::Impl {
};
void addWrites(const WriteRange *writes, int count) {
auto stepwiseInserts =
std::unique_ptr<StepwiseInsert[]>{new StepwiseInsert[count]};
Arena arena;
auto *stepwiseInserts = new (arena) StepwiseInsert[count];
for (int i = 0; i < count; ++i) {
// TODO handle non-singleton writes lol
assert(writes[i].end.len == 0);
@@ -518,11 +684,12 @@ struct ConflictSet::Impl {
// probably fine.
// TODO better/faster RNG?
std::mt19937 g(fastRand());
std::shuffle(stepwiseInserts.get(), stepwiseInserts.get() + count, g);
std::shuffle(stepwiseInserts, stepwiseInserts + count, g);
runInterleaved(std::span<StepwiseInsert>(stepwiseInserts.get(), count));
runInterleaved(std::span<StepwiseInsert>(stepwiseInserts, count));
std::vector<Node *> workList;
std::vector<Node *, ArenaAlloc<Node *>> workList{
ArenaAlloc<Node *>(&arena)};
workList.reserve(count);
for (int i = 0; i < count; ++i) {
workList.push_back(*stepwiseInserts[i].current);
@@ -568,7 +735,8 @@ struct ConflictSet::Impl {
}
~Impl() {
std::vector<Node *> toFree;
Arena arena;
std::vector<Node *, ArenaAlloc<Node *>> toFree{ArenaAlloc<Node *>(&arena)};
if (root != nullptr) {
toFree.push_back(root);
}
@@ -600,9 +768,12 @@ void ConflictSet::setOldestVersion(int64_t oldestVersion) {
}
ConflictSet::ConflictSet(int64_t oldestVersion)
: impl(new Impl{oldestVersion}) {}
: impl(new(malloc(sizeof(Impl))) Impl{oldestVersion}) {}
ConflictSet::~ConflictSet() { delete impl; }
ConflictSet::~ConflictSet() {
impl->~Impl();
free(impl);
}
ConflictSet::ConflictSet(ConflictSet &&other) noexcept
: impl(std::exchange(other.impl, nullptr)) {}
@@ -618,8 +789,9 @@ int main(void) {
ConflictSet::Impl cs{writeVersion};
constexpr int kNumKeys = 5;
ConflictSet::WriteRange write[kNumKeys];
Arena arena;
for (int i = 0; i < kNumKeys; ++i) {
write[i].begin = toKey(i);
write[i].begin = toKey(arena, i);
write[i].end.len = 0;
write[i].writeVersion = ++writeVersion;
}