rate_limiter.hpp

Go to the documentation of this file.

#ifndef LIBFILEZILLA_RATE_LIMITER_HEADER
#define LIBFILEZILLA_RATE_LIMITER_HEADER
 
#include "event_handler.hpp"
 
#include <atomic>
#include <vector>
 
namespace fz {
 
namespace rate {
using type = uint64_t;
enum : type {
    unlimited = static_cast<type>(-1)
};
}
 
namespace direction {
enum type : size_t {
    inbound,
    outbound
};
}
 
class rate_limiter;
 
class FZ_PUBLIC_SYMBOL rate_limit_manager final : public event_handler
{
public:
    explicit rate_limit_manager(event_loop & loop);
    virtual ~rate_limit_manager();
 
    void add(rate_limiter* limiter);
 
    void set_burst_tolerance(rate::type tolerance);
 
private:
    friend class rate_limiter;
    friend class bucket_base;
    friend class bucket;
 
    void FZ_PRIVATE_SYMBOL record_activity();
 
    void FZ_PRIVATE_SYMBOL operator()(event_base const& ev);
    void FZ_PRIVATE_SYMBOL on_timer(timer_id const&);
 
    void FZ_PRIVATE_SYMBOL process(rate_limiter* limiter, bool locked);
 
    std::atomic<int> activity_{2};
    mutex mtx_{false};
    std::vector<rate_limiter*> limiters_;
 
    std::atomic<timer_id> timer_{};
 
    std::atomic<rate::type> burst_tolerance_{1};
};
 
class FZ_PUBLIC_SYMBOL bucket_base
{
public:
    virtual ~bucket_base() noexcept = default;
 
    virtual void remove_bucket();
 
protected:
    friend class rate_limiter;
 
    virtual void lock_tree() { mtx_.lock(); }
 
    virtual void update_stats(bool & active) = 0;
 
    virtual size_t weight() const { return 1; }
 
    virtual size_t unsaturated(direction::type const /*d*/) const { return 0; }
 
    virtual void set_mgr_recursive(rate_limit_manager * mgr);
 
    virtual rate::type add_tokens(direction::type const /*d*/, rate::type /*tokens*/, rate::type /*limit*/) = 0;
 
    virtual rate::type distribute_overflow(direction::type const /*d*/, rate::type /*tokens*/) { return 0; }
 
    virtual void unlock_tree() { mtx_.unlock(); }
 
    virtual std::array<rate::type, 2> gather_unspent_for_removal() = 0;
 
    mutex mtx_{false};
    rate_limit_manager * mgr_{};
    void * parent_{};
    size_t idx_{static_cast<size_t>(-1)};
};
 
class FZ_PUBLIC_SYMBOL rate_limiter final : public bucket_base
{
public:
    rate_limiter() = default;
    explicit rate_limiter(rate_limit_manager * mgr);
    virtual ~rate_limiter();
 
    void add(bucket_base* bucket);
 
    void set_limits(rate::type download_limit, rate::type upload_limit);
 
    rate::type limit(direction::type const d);
 
private:
    friend class bucket_base;
    friend class rate_limit_manager;
 
    virtual void FZ_PRIVATE_SYMBOL lock_tree() override;
 
    bool FZ_PRIVATE_SYMBOL do_set_limit(direction::type const d, rate::type limit);
 
    virtual void FZ_PRIVATE_SYMBOL update_stats(bool & active) override;
    virtual size_t FZ_PRIVATE_SYMBOL weight() const override { return weight_; }
    virtual size_t FZ_PRIVATE_SYMBOL unsaturated(direction::type const d) const override { return data_[d].unused_capacity_ ? data_[d].unsaturated_ : 0; }
    virtual void FZ_PRIVATE_SYMBOL set_mgr_recursive(rate_limit_manager * mgr) override;
 
    virtual rate::type FZ_PRIVATE_SYMBOL add_tokens(direction::type const d, rate::type tokens, rate::type limit) override;
    virtual rate::type FZ_PRIVATE_SYMBOL distribute_overflow(direction::type const d, rate::type tokens) override;
 
    virtual void FZ_PRIVATE_SYMBOL unlock_tree() override;
 
    void FZ_PRIVATE_SYMBOL pay_debt(direction::type const d);
 
    virtual std::array<rate::type, 2> FZ_PRIVATE_SYMBOL gather_unspent_for_removal() override;
 
    std::vector<bucket_base*> buckets_;
    std::vector<size_t> scratch_buffer_;
    size_t weight_{};
 
    struct FZ_PRIVATE_SYMBOL data_t {
        rate::type limit_{rate::unlimited};
        rate::type merged_tokens_{};
        rate::type overflow_{};
        rate::type debt_{};
        rate::type unused_capacity_{};
        rate::type carry_{};
        size_t unsaturated_{};
    };
    data_t data_[2];
};
 
class FZ_PUBLIC_SYMBOL bucket : public bucket_base
{
public:
    virtual ~bucket();
 
    virtual void remove_bucket() override;
 
    rate::type available(direction::type const d);
 
    void consume(direction::type const d, rate::type amount);
 
protected:
    virtual void wakeup(direction::type /*d*/) {}
 
    bool waiting(scoped_lock & l, direction::type d);
 
private:
    virtual void update_stats(bool & active) override final;
    virtual size_t unsaturated(direction::type const d) const override final { return data_[d].unsaturated_ ? 1 : 0; }
 
    virtual rate::type add_tokens(direction::type const d, rate::type tokens, rate::type limit) override final;
    virtual rate::type distribute_overflow(direction::type const d, rate::type tokens) override final;
 
    virtual void unlock_tree() override final;
 
    virtual std::array<rate::type, 2> gather_unspent_for_removal() override final;
 
    struct FZ_PRIVATE_SYMBOL data_t {
        rate::type available_{rate::unlimited};
        rate::type overflow_multiplier_{1};
        rate::type bucket_size_{rate::unlimited};
        bool waiting_{};
        bool unsaturated_{};
    } data_[2];
};
 
}
 
#endif