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