#pragma once
#include <atomic>
#include "rocksdb/memory_allocator.h"
namespace ROCKSDB_NAMESPACE {
class DefaultMemoryAllocator : public MemoryAllocator {
public:
static const char* kClassName() { return "DefaultMemoryAllocator"; }
const char* Name() const override { return kClassName(); }
void* Allocate(size_t size) override {
return static_cast<void*>(new char[size]);
}
void Deallocate(void* p) override { delete[] static_cast<char*>(p); }
};
class BaseMemoryAllocator : public MemoryAllocator {
public:
void* Allocate(size_t ) override {
assert(false);
return nullptr;
}
void Deallocate(void* ) override { assert(false); }
};
class MemoryAllocatorWrapper : public MemoryAllocator {
public:
explicit MemoryAllocatorWrapper(const std::shared_ptr<MemoryAllocator>& t);
~MemoryAllocatorWrapper() override {}
MemoryAllocator* target() const { return target_.get(); }
void* Allocate(size_t size) override { return target_->Allocate(size); }
void Deallocate(void* p) override { return target_->Deallocate(p); }
size_t UsableSize(void* p, size_t allocation_size) const override {
return target_->UsableSize(p, allocation_size);
}
const Customizable* Inner() const override { return target_.get(); }
protected:
std::shared_ptr<MemoryAllocator> target_;
};
class CountedMemoryAllocator : public MemoryAllocatorWrapper {
public:
CountedMemoryAllocator()
: MemoryAllocatorWrapper(std::make_shared<DefaultMemoryAllocator>()),
allocations_(0),
deallocations_(0) {}
explicit CountedMemoryAllocator(const std::shared_ptr<MemoryAllocator>& t)
: MemoryAllocatorWrapper(t), allocations_(0), deallocations_(0) {}
static const char* kClassName() { return "CountedMemoryAllocator"; }
const char* Name() const override { return kClassName(); }
std::string GetId() const override { return std::string(Name()); }
void* Allocate(size_t size) override {
allocations_++;
return MemoryAllocatorWrapper::Allocate(size);
}
void Deallocate(void* p) override {
deallocations_++;
MemoryAllocatorWrapper::Deallocate(p);
}
uint64_t GetNumAllocations() const { return allocations_; }
uint64_t GetNumDeallocations() const { return deallocations_; }
private:
std::atomic<uint64_t> allocations_;
std::atomic<uint64_t> deallocations_;
};
}