Home | History | Annotate | Download | only in netdutils 
      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:
     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 
     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