1 /* 2 * Copyright (C) 2018 The Android Open Source Project 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 #ifndef NETUTILS_OPERATIONLIMITER_H 18 #define NETUTILS_OPERATIONLIMITER_H 19 20 #include <mutex> 21 #include <unordered_map> 22 23 #include <android-base/logging.h> 24 #include <android-base/thread_annotations.h> 25 26 namespace android { 27 namespace netdutils { 28 29 // Tracks the number of operations in progress on behalf of a particular key or 30 // ID, rejecting further attempts to start new operations after a configurable 31 // limit has been reached. 32 // 33 // The intended usage pattern is: 34 // OperationLimiter<UserId> connections_per_user; 35 // ... 36 // int connectToSomeResource(int user) { 37 // if (!connections_per_user.start(user)) return TRY_AGAIN_LATER; 38 // // ...do expensive work here... 39 // connections_per_user.finish(user); 40 // } 41 // 42 // This class is thread-safe. 43 template <typename KeyType> 44 class OperationLimiter { 45 public: 46 OperationLimiter(int limitPerKey, int globalLimit = INT_MAX) mLimitPerKey(limitPerKey)47 : mLimitPerKey(limitPerKey), mGlobalLimit(globalLimit) {} 48 ~OperationLimiter()49 ~OperationLimiter() { 50 DCHECK(mCounters.empty()) << "Destroying OperationLimiter with active operations"; 51 } 52 53 // Returns false if |key| has reached the maximum number of concurrent operations, 54 // or if the global limit has been reached. Otherwise, increments the counter and returns true. 55 // 56 // Note: each successful start(key) must be matched by exactly one call to 57 // finish(key). start(KeyType key)58 bool start(KeyType key) EXCLUDES(mMutex) { 59 std::lock_guard lock(mMutex); 60 61 if (mGlobalCounter >= mGlobalLimit) { 62 // Oh, no! 63 return false; 64 } 65 66 auto& cnt = mCounters[key]; // operator[] creates new entries as needed. 67 if (cnt >= mLimitPerKey) { 68 // Oh, no! 69 return false; 70 } 71 72 ++cnt; 73 ++mGlobalCounter; 74 return true; 75 } 76 77 // Decrements the number of operations in progress accounted to |key|. 78 // See usage notes on start(). finish(KeyType key)79 void finish(KeyType key) EXCLUDES(mMutex) { 80 std::lock_guard lock(mMutex); 81 82 --mGlobalCounter; 83 if (mGlobalCounter < 0) { 84 LOG(FATAL_WITHOUT_ABORT) << "Global operations counter going negative, this is a bug."; 85 return; 86 } 87 88 auto it = mCounters.find(key); 89 if (it == mCounters.end()) { 90 LOG(FATAL_WITHOUT_ABORT) << "Decremented non-existent counter for key=" << key; 91 return; 92 } 93 auto& cnt = it->second; 94 --cnt; 95 if (cnt <= 0) { 96 // Cleanup counters once they drop down to zero. 97 mCounters.erase(it); 98 } 99 } 100 101 private: 102 // Protects access to the map below. 103 std::mutex mMutex; 104 105 // Tracks the number of outstanding queries by key. 106 std::unordered_map<KeyType, int> mCounters GUARDED_BY(mMutex); 107 108 int mGlobalCounter GUARDED_BY(mMutex) = 0; 109 110 // Maximum number of outstanding queries from a single key. 111 const int mLimitPerKey; 112 113 // Maximum number of outstanding queries, globally. 114 const int mGlobalLimit; 115 }; 116 117 } // namespace netdutils 118 } // namespace android 119 120 #endif // NETUTILS_OPERATIONLIMITER_H 121