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