ceph/src/rocksdb/memtable/alloc_tracker.cc

   1 //  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
   2 //  This source code is licensed under both the GPLv2 (found in the
   3 //  COPYING file in the root directory) and Apache 2.0 License
   4 //  (found in the LICENSE.Apache file in the root 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 <assert.h>
  11 #include "rocksdb/write_buffer_manager.h"
  12 #include "util/allocator.h"
  13 #include "util/arena.h"
  14
  15 namespace rocksdb {
  16
  17 AllocTracker::AllocTracker(WriteBufferManager* write_buffer_manager)
  18     : write_buffer_manager_(write_buffer_manager),
  19       bytes_allocated_(0),
  20       done_allocating_(false),
  21       freed_(false) {}
  22
  23 AllocTracker::~AllocTracker() { FreeMem(); }
  24
  25 void AllocTracker::Allocate(size_t bytes) {
  26   assert(write_buffer_manager_ != nullptr);
  27   if (write_buffer_manager_->enabled()) {
  28     bytes_allocated_.fetch_add(bytes, std::memory_order_relaxed);
  29     write_buffer_manager_->ReserveMem(bytes);
  30   }
  31 }
  32
  33 void AllocTracker::DoneAllocating() {
  34   if (write_buffer_manager_ != nullptr && !done_allocating_) {
  35     if (write_buffer_manager_->enabled()) {
  36       write_buffer_manager_->ScheduleFreeMem(
  37           bytes_allocated_.load(std::memory_order_relaxed));
  38     } else {
  39       assert(bytes_allocated_.load(std::memory_order_relaxed) == 0);
  40     }
  41     done_allocating_ = true;
  42   }
  43 }
  44
  45 void AllocTracker::FreeMem() {
  46   if (!done_allocating_) {
  47     DoneAllocating();
  48   }
  49   if (write_buffer_manager_ != nullptr && !freed_) {
  50     if (write_buffer_manager_->enabled()) {
  51       write_buffer_manager_->FreeMem(
  52           bytes_allocated_.load(std::memory_order_relaxed));
  53     } else {
  54       assert(bytes_allocated_.load(std::memory_order_relaxed) == 0);
  55     }
  56     freed_ = true;
  57   }
  58 }
  59 }  // namespace rocksdb