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 // // Before opening a new connection 37 // if (!limiter.start(user)) { 38 // return error; 39 // } else { 40 // // open the connection 41 // // ...do some work... 42 // // close the connection 43 // limiter.finish(user); 44 // } 45 // 46 // This class is thread-safe. 47 template<typename KeyType> 48 class OperationLimiter { 49 public: OperationLimiter(int limit)50 explicit OperationLimiter(int limit) : mLimitPerKey(limit) {} 51 ~OperationLimiter()52 ~OperationLimiter() { 53 DCHECK(mCounters.empty()) 54 << "Destroying OperationLimiter with active operations"; 55 } 56 57 // Returns false if |key| has reached the maximum number of concurrent 58 // operations, otherwise increments the counter and returns true. 59 // 60 // Note: each successful start(key) must be matched by exactly one call to 61 // finish(key). start(KeyType key)62 bool start(KeyType key) EXCLUDES(mMutex) { 63 std::lock_guard lock(mMutex); 64 auto& cnt = mCounters[key]; // operator[] creates new entries as needed. 65 if (cnt >= mLimitPerKey) { 66 // Oh, no! 67 return false; 68 } 69 ++cnt; 70 return true; 71 } 72 73 // Decrements the number of operations in progress accounted to |key|. 74 // See usage notes on start(). finish(KeyType key)75 void finish(KeyType key) EXCLUDES(mMutex) { 76 std::lock_guard lock(mMutex); 77 auto it = mCounters.find(key); 78 if (it == mCounters.end()) { 79 LOG(FATAL_WITHOUT_ABORT) << "Decremented non-existent counter for key=" << key; 80 return; 81 } 82 auto& cnt = it->second; 83 --cnt; 84 if (cnt <= 0) { 85 // Cleanup counters once they drop down to zero. 86 mCounters.erase(it); 87 } 88 } 89 90 private: 91 // Protects access to the map below. 92 std::mutex mMutex; 93 94 // Tracks the number of outstanding queries by key. 95 std::unordered_map<KeyType, int> mCounters GUARDED_BY(mMutex); 96 97 // Maximum number of outstanding queries from a single key. 98 const int mLimitPerKey; 99 }; 100 101 } // namespace netdutils 102 } // namespace android 103 104 #endif // NETUTILS_OPERATIONLIMITER_H 105