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