126 lines
4.0 KiB
C++
126 lines
4.0 KiB
C++
/*
|
|
* Copyright (C) 2018 The Android Open Source Project
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#ifndef NETUTILS_OPERATIONLIMITER_H
|
|
#define NETUTILS_OPERATIONLIMITER_H
|
|
|
|
#include <mutex>
|
|
#include <unordered_map>
|
|
|
|
#include <android-base/logging.h>
|
|
#include <android-base/thread_annotations.h>
|
|
|
|
#include "Experiments.h"
|
|
|
|
namespace android {
|
|
namespace netdutils {
|
|
|
|
// Tracks the number of operations in progress on behalf of a particular key or
|
|
// ID, rejecting further attempts to start new operations after a configurable
|
|
// limit has been reached.
|
|
//
|
|
// The intended usage pattern is:
|
|
// OperationLimiter<UserId> connections_per_user;
|
|
// ...
|
|
// int connectToSomeResource(int user) {
|
|
// if (!connections_per_user.start(user)) return TRY_AGAIN_LATER;
|
|
// // ...do expensive work here...
|
|
// connections_per_user.finish(user);
|
|
// }
|
|
//
|
|
// This class is thread-safe.
|
|
template <typename KeyType>
|
|
class OperationLimiter {
|
|
public:
|
|
OperationLimiter(int limitPerKey) : mLimitPerKey(limitPerKey) {}
|
|
|
|
~OperationLimiter() {
|
|
DCHECK(mCounters.empty()) << "Destroying OperationLimiter with active operations";
|
|
}
|
|
|
|
// Returns false if |key| has reached the maximum number of concurrent operations,
|
|
// or if the global limit has been reached. Otherwise, increments the counter and returns true.
|
|
//
|
|
// Note: each successful start(key) must be matched by exactly one call to
|
|
// finish(key).
|
|
bool start(KeyType key) EXCLUDES(mMutex) {
|
|
std::lock_guard lock(mMutex);
|
|
int globalLimit =
|
|
android::net::Experiments::getInstance()->getFlag("max_queries_global", INT_MAX);
|
|
if (globalLimit < mLimitPerKey) {
|
|
LOG(ERROR) << "Misconfiguration on max_queries_global " << globalLimit;
|
|
globalLimit = INT_MAX;
|
|
}
|
|
if (mGlobalCounter >= globalLimit) {
|
|
// Oh, no!
|
|
LOG(ERROR) << "Query from " << key << " denied due to global limit: " << globalLimit;
|
|
return false;
|
|
}
|
|
|
|
auto& cnt = mCounters[key]; // operator[] creates new entries as needed.
|
|
if (cnt >= mLimitPerKey) {
|
|
// Oh, no!
|
|
LOG(ERROR) << "Query from " << key << " denied due to limit: " << mLimitPerKey;
|
|
return false;
|
|
}
|
|
|
|
++cnt;
|
|
++mGlobalCounter;
|
|
return true;
|
|
}
|
|
|
|
// Decrements the number of operations in progress accounted to |key|.
|
|
// See usage notes on start().
|
|
void finish(KeyType key) EXCLUDES(mMutex) {
|
|
std::lock_guard lock(mMutex);
|
|
|
|
--mGlobalCounter;
|
|
if (mGlobalCounter < 0) {
|
|
LOG(FATAL_WITHOUT_ABORT) << "Global operations counter going negative, this is a bug.";
|
|
return;
|
|
}
|
|
|
|
auto it = mCounters.find(key);
|
|
if (it == mCounters.end()) {
|
|
LOG(FATAL_WITHOUT_ABORT) << "Decremented non-existent counter for key=" << key;
|
|
return;
|
|
}
|
|
auto& cnt = it->second;
|
|
--cnt;
|
|
if (cnt <= 0) {
|
|
// Cleanup counters once they drop down to zero.
|
|
mCounters.erase(it);
|
|
}
|
|
}
|
|
|
|
private:
|
|
// Protects access to the map below.
|
|
std::mutex mMutex;
|
|
|
|
// Tracks the number of outstanding queries by key.
|
|
std::unordered_map<KeyType, int> mCounters GUARDED_BY(mMutex);
|
|
|
|
int mGlobalCounter GUARDED_BY(mMutex) = 0;
|
|
|
|
// Maximum number of outstanding queries from a single key.
|
|
const int mLimitPerKey;
|
|
};
|
|
|
|
} // namespace netdutils
|
|
} // namespace android
|
|
|
|
#endif // NETUTILS_OPERATIONLIMITER_H
|