Limit the number of outstanding DNS queries by UID
Test: system/netd/tests/runtests.sh
Test: run netdutils_test on marlin:
[----------] 3 tests from OperationLimiter
[ RUN ] OperationLimiter.limits
[ OK ] OperationLimiter.limits (0 ms)
[ RUN ] OperationLimiter.finishWithoutStart
[ OK ] OperationLimiter.finishWithoutStart (3 ms)
[ RUN ] OperationLimiter.destroyWithActiveOperations
[ OK ] OperationLimiter.destroyWithActiveOperations (1 ms)
[----------] 3 tests from OperationLimiter (6 ms total)
Bug: 79674503
Change-Id: I5f11f0ed6b6f2479921d90a919d17dfd7b7f5788
diff --git a/libnetdutils/include/netdutils/OperationLimiter.h b/libnetdutils/include/netdutils/OperationLimiter.h
new file mode 100644
index 0000000..633536b
--- /dev/null
+++ b/libnetdutils/include/netdutils/OperationLimiter.h
@@ -0,0 +1,104 @@
+/*
+ * 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>
+
+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;
+// ...
+// // Before opening a new connection
+// if (!limiter.start(user)) {
+// return error;
+// } else {
+// // open the connection
+// // ...do some work...
+// // close the connection
+// limiter.finish(user);
+// }
+//
+// This class is thread-safe.
+template<typename KeyType>
+class OperationLimiter {
+public:
+ explicit OperationLimiter(int limit) : mLimitPerKey(limit) {}
+
+ ~OperationLimiter() {
+ DCHECK(mCounters.empty())
+ << "Destroying OperationLimiter with active operations";
+ }
+
+ // Returns false if |key| has reached the maximum number of concurrent
+ // operations, 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<std::mutex> lock(mMutex);
+ auto& cnt = mCounters[key]; // operator[] creates new entries as needed.
+ if (cnt >= mLimitPerKey) {
+ // Oh, no!
+ return false;
+ }
+ ++cnt;
+ 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<std::mutex> lock(mMutex);
+ 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);
+
+ // Maximum number of outstanding queries from a single key.
+ const int mLimitPerKey;
+};
+
+} // namespace netdutils
+} // namespace android
+
+#endif // NETUTILS_OPERATIONLIMITER_H