1
0
Fork 0
mirror of https://github.com/dragonflydb/dragonfly.git synced 2024-12-14 11:58:02 +00:00
dragonflydb-dragonfly/core/segment_allocator.h

76 lines
2 KiB
C++

// Copyright 2022, Roman Gershman. All rights reserved.
// See LICENSE for licensing terms.
//
#pragma once
#include <absl/container/flat_hash_map.h>
#include <mimalloc.h>
/***
* This class is tightly coupled with mimalloc segment allocation logic and is designed to provide
* a compact pointer representation (4bytes ptr) over 64bit address space that gives you
* 32GB of allocations with option to extend it to 32*256GB if needed.
*
*/
namespace dfly {
/**
* @brief Tightly coupled with mi_malloc 2.x implementation.
* Fetches 8MB segment pointers from the allocated pointers.
* Provides own indexing of small pointers to real address space using the segment ptrs/
*/
class SegmentAllocator {
static constexpr uint32_t kSegmentIdBits = 12;
static constexpr uint32_t kSegmentIdMask = (1 << kSegmentIdBits) - 1;
static constexpr uint64_t kSegmentAlignMask = ~((1 << 23) - 1);
public:
using Ptr = uint32_t;
SegmentAllocator(mi_heap_t* heap);
uint8_t* Translate(Ptr p) const {
return address_table_[p & kSegmentIdMask] + Offset(p);
}
std::pair<Ptr, uint8_t*> Allocate(uint32_t size);
void Free(Ptr ptr) {
mi_free(Translate(ptr));
}
mi_heap_t* heap() {
return heap_;
}
private:
static uint32_t Offset(Ptr p) {
return (p >> kSegmentIdBits) * 8;
}
void ValidateMapSize();
std::vector<uint8_t*> address_table_;
absl::flat_hash_map<uint64_t, uint16_t> rev_indx_;
mi_heap_t* heap_;
};
inline auto SegmentAllocator::Allocate(uint32_t size) -> std::pair<Ptr, uint8_t*> {
uint64_t ptr = (uint64_t)mi_heap_malloc(heap_, size);
uint64_t seg_ptr = ptr & kSegmentAlignMask;
// could be speed up using last used seg_ptr.
auto [it, inserted] = rev_indx_.emplace(seg_ptr, address_table_.size());
if (inserted) {
ValidateMapSize();
address_table_.push_back((uint8_t*)seg_ptr);
}
Ptr res = (((ptr - seg_ptr) / 8) << kSegmentIdBits) | it->second;
return std::make_pair(res, (uint8_t*)ptr);
}
} // namespace dfly