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). start(KeyType key)59 bool start(KeyType key) EXCLUDES(mMutex) { 60 std::lock_guard lock(mMutex); 61 int globalLimit = 62 android::net::Experiments::getInstance()->getFlag("max_queries_global", INT_MAX); 63 if (globalLimit < mLimitPerKey) { 64 LOG(ERROR) << "Misconfiguration on max_queries_global " << globalLimit; 65 globalLimit = INT_MAX; 66 } 67 if (mGlobalCounter >= globalLimit) { 68 // Oh, no! 69 LOG(ERROR) << "Query from " << key << " denied due to global limit: " << globalLimit; 70 return false; 71 } 72 73 auto& cnt = mCounters[key]; // operator[] creates new entries as needed. 74 if (cnt >= mLimitPerKey) { 75 // Oh, no! 76 LOG(ERROR) << "Query from " << key << " denied due to limit: " << mLimitPerKey; 77 return false; 78 } 79 80 ++cnt; 81 ++mGlobalCounter; 82 return true; 83 } 84 85 // Decrements the number of operations in progress accounted to |key|. 86 // See usage notes on start(). finish(KeyType key)87 void finish(KeyType key) EXCLUDES(mMutex) { 88 std::lock_guard lock(mMutex); 89 90 --mGlobalCounter; 91 if (mGlobalCounter < 0) { 92 LOG(FATAL_WITHOUT_ABORT) << "Global operations counter going negative, this is a bug."; 93 return; 94 } 95 96 auto it = mCounters.find(key); 97 if (it == mCounters.end()) { 98 LOG(FATAL_WITHOUT_ABORT) << "Decremented non-existent counter for key=" << key; 99 return; 100 } 101 auto& cnt = it->second; 102 --cnt; 103 if (cnt <= 0) { 104 // Cleanup counters once they drop down to zero. 105 mCounters.erase(it); 106 } 107 } 108 109 private: 110 // Protects access to the map below. 111 std::mutex mMutex; 112 113 // Tracks the number of outstanding queries by key. 114 std::unordered_map<KeyType, int> mCounters GUARDED_BY(mMutex); 115 116 int mGlobalCounter GUARDED_BY(mMutex) = 0; 117 118 // Maximum number of outstanding queries from a single key. 119 const int mLimitPerKey; 120 }; 121 122 } // namespace netdutils 123 } // namespace android 124 125 #endif // NETUTILS_OPERATIONLIMITER_H 126