ceph/src/rocksdb/memtable/memtable_allocator.cc

   1 //  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
   2 //  This source code is licensed under the BSD-style license found in the
   3 //  LICENSE file in the root directory of this source tree. An additional grant
   4 //  of patent rights can be found in the PATENTS file in the same directory.
   5 //
   6 // Copyright (c) 2011 The LevelDB Authors. All rights reserved.
   7 // Use of this source code is governed by a BSD-style license that can be
   8 // found in the LICENSE file. See the AUTHORS file for names of contributors.
   9
  10 #include "memtable/memtable_allocator.h"
  11
  12 #include <assert.h>
  13 #include "rocksdb/write_buffer_manager.h"
  14 #include "util/arena.h"
  15
  16 namespace rocksdb {
  17
  18 MemTableAllocator::MemTableAllocator(Allocator* allocator,
  19                                      WriteBufferManager* write_buffer_manager)
  20     : allocator_(allocator),
  21       write_buffer_manager_(write_buffer_manager),
  22       bytes_allocated_(0) {}
  23
  24 MemTableAllocator::~MemTableAllocator() { DoneAllocating(); }
  25
  26 char* MemTableAllocator::Allocate(size_t bytes) {
  27   assert(write_buffer_manager_ != nullptr);
  28   if (write_buffer_manager_->enabled()) {
  29     bytes_allocated_.fetch_add(bytes, std::memory_order_relaxed);
  30     write_buffer_manager_->ReserveMem(bytes);
  31   }
  32   return allocator_->Allocate(bytes);
  33 }
  34
  35 char* MemTableAllocator::AllocateAligned(size_t bytes, size_t huge_page_size,
  36                                          Logger* logger) {
  37   assert(write_buffer_manager_ != nullptr);
  38   if (write_buffer_manager_->enabled()) {
  39     bytes_allocated_.fetch_add(bytes, std::memory_order_relaxed);
  40     write_buffer_manager_->ReserveMem(bytes);
  41   }
  42   return allocator_->AllocateAligned(bytes, huge_page_size, logger);
  43 }
  44
  45 void MemTableAllocator::DoneAllocating() {
  46   if (write_buffer_manager_ != nullptr) {
  47     if (write_buffer_manager_->enabled()) {
  48       write_buffer_manager_->FreeMem(
  49           bytes_allocated_.load(std::memory_order_relaxed));
  50     } else {
  51       assert(bytes_allocated_.load(std::memory_order_relaxed) == 0);
  52     }
  53     write_buffer_manager_ = nullptr;
  54   }
  55 }
  56
  57 size_t MemTableAllocator::BlockSize() const { return allocator_->BlockSize(); }
  58
  59 }  // namespace rocksdb