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 "memory/allocator.h"
12 #include "memory/arena.h"
13 #include "rocksdb/write_buffer_manager.h"
14 
15 namespace ROCKSDB_NAMESPACE {
16 
AllocTracker(WriteBufferManager * write_buffer_manager)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 
~AllocTracker()23 AllocTracker::~AllocTracker() { FreeMem(); }
24 
Allocate(size_t bytes)25 void AllocTracker::Allocate(size_t bytes) {
26   assert(write_buffer_manager_ != nullptr);
27   if (write_buffer_manager_->enabled() ||
28       write_buffer_manager_->cost_to_cache()) {
29     bytes_allocated_.fetch_add(bytes, std::memory_order_relaxed);
30     write_buffer_manager_->ReserveMem(bytes);
31   }
32 }
33 
DoneAllocating()34 void AllocTracker::DoneAllocating() {
35   if (write_buffer_manager_ != nullptr && !done_allocating_) {
36     if (write_buffer_manager_->enabled() ||
37         write_buffer_manager_->cost_to_cache()) {
38       write_buffer_manager_->ScheduleFreeMem(
39           bytes_allocated_.load(std::memory_order_relaxed));
40     } else {
41       assert(bytes_allocated_.load(std::memory_order_relaxed) == 0);
42     }
43     done_allocating_ = true;
44   }
45 }
46 
FreeMem()47 void AllocTracker::FreeMem() {
48   if (!done_allocating_) {
49     DoneAllocating();
50   }
51   if (write_buffer_manager_ != nullptr && !freed_) {
52     if (write_buffer_manager_->enabled() ||
53         write_buffer_manager_->cost_to_cache()) {
54       write_buffer_manager_->FreeMem(
55           bytes_allocated_.load(std::memory_order_relaxed));
56     } else {
57       assert(bytes_allocated_.load(std::memory_order_relaxed) == 0);
58     }
59     freed_ = true;
60   }
61 }
62 }  // namespace ROCKSDB_NAMESPACE
63