/*
 * 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 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 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