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 26namespace android { 27namespace 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. 47template<typename KeyType> 48class OperationLimiter { 49public: 50 explicit OperationLimiter(int limit) : mLimitPerKey(limit) {} 51 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). 62 bool start(KeyType key) EXCLUDES(mMutex) { 63 std::lock_guard<std::mutex> 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(). 75 void finish(KeyType key) EXCLUDES(mMutex) { 76 std::lock_guard<std::mutex> 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 90private: 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