1 /*
2  * Copyright (c) Facebook, Inc. and its affiliates.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *     http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #pragma once
18 
19 #include <cstddef>
20 #include <cstdint>
21 
22 #include <folly/Utility.h>
23 #include <folly/experimental/observer/Observer.h>
24 
25 namespace apache {
26 namespace thrift {
27 
28 /**
29  * Tracks memory consumption.
30  */
31 class MemoryTracker : folly::MoveOnly {
32  public:
33   using Observer = folly::observer::Observer<size_t>;
34   using AtomicObserver = folly::observer::ReadMostlyAtomicObserver<size_t>;
35 
36   // A basic tracker which does not check for limits
MemoryTracker()37   MemoryTracker()
38       : MemoryTracker(
39             folly::observer::makeStaticObserver<size_t>(0),
40             folly::observer::makeStaticObserver<size_t>(0)) {}
41 
42   // A tracker with a memory limit
MemoryTracker(Observer limitObserver)43   explicit MemoryTracker(Observer limitObserver)
44       : MemoryTracker(
45             std::move(limitObserver),
46             folly::observer::makeStaticObserver<size_t>(0)) {}
47 
48   // A tracker with a memory limit which is only checked for memory increments
49   // of a minimum size.
MemoryTracker(Observer limitObserver,Observer minObserver)50   MemoryTracker(Observer limitObserver, Observer minObserver)
51       : limitObserver_(std::move(limitObserver)),
52         minObserver_(std::move(minObserver)) {}
53 
54   // Increases the memory consumption counter by `size`. Returns `true` if usage
55   // is below limit or `size` is less than minimum.
increment(size_t size)56   bool increment(size_t size) {
57     usage_ += size;
58     auto limit = *limitObserver_;
59     auto min = *minObserver_;
60     return !limit || size < min || usage_ <= limit;
61   }
62 
decrement(size_t size)63   void decrement(size_t size) {
64     DCHECK_GE(usage_, size);
65     usage_ -= std::min(size, usage_);
66   }
67 
68  private:
69   size_t usage_{0};
70   AtomicObserver limitObserver_;
71   AtomicObserver minObserver_;
72 };
73 
74 } // namespace thrift
75 } // namespace apache
76