vdm/core/include/vdm/rate/token_bucket.hpp

// vdm/rate/token_bucket.hpp — a lazily-refilled token bucket, and the global -> queue ->
// task limiter hierarchy built on it (docs/04 §6).
//
// A segment worker calls RateLimiter::acquire(task, n) after receiving n body bytes. If
// every applicable level (task, its queue, global) has n tokens, it consumes n from each
// and returns 0. Otherwise it consumes nothing and returns how long to wait before
// retrying — the worker returns CURL_WRITEFUNC_PAUSE and schedules a curl_easy_pause
// resume after that delay (the "precision" layer on top of CURLOPT_MAX_RECV_SPEED_LARGE).
//
// This header compiles standalone.

#ifndef VDM_RATE_TOKEN_BUCKET_HPP
#define VDM_RATE_TOKEN_BUCKET_HPP

#include <algorithm>
#include <chrono>
#include <cstdint>
#include <mutex>
#include <optional>
#include <unordered_map>

#include "vdm/ids.hpp"

namespace vdm::rate {

// rate_bps == 0 means unlimited: consume() always succeeds and never waits.
class TokenBucket {
  public:
    TokenBucket() = default;
    // `burst` caps how many tokens accumulate while idle; 0 => 1 second's worth.
    explicit TokenBucket(std::uint64_t rate_bps, std::uint64_t burst = 0) {
        set_rate(rate_bps, burst);
    }

    void set_rate(std::uint64_t rate_bps, std::uint64_t burst = 0) {
        std::lock_guard lk(mu_);
        const bool was_unlimited = rate_ == 0;
        rate_ = rate_bps;
        cap_ = burst ? burst : rate_bps;  // 1 s of burst by default
        // A freshly-limited bucket starts full: you may transfer at burst speed
        // immediately, then it throttles (classic token bucket / IDM behaviour). Lowering
        // an existing limit only clamps down — it never hands out a fresh burst.
        if (rate_bps > 0 && was_unlimited)
            tokens_ = cap_;
        else if (tokens_ > cap_)
            tokens_ = cap_;
        last_ = clock::now();
    }
    [[nodiscard]] std::uint64_t rate() const {
        std::lock_guard lk(mu_);
        return rate_;
    }

    // Consume `n` if available; otherwise consume nothing. Returns the wait until `n`
    // tokens *would* be available (0 when it consumed).
    [[nodiscard]] std::chrono::nanoseconds consume(std::uint64_t n) {
        std::lock_guard lk(mu_);
        if (rate_ == 0)
            return {};
        refill_locked();
        if (tokens_ >= n) {
            tokens_ -= n;
            return {};
        }
        const std::uint64_t deficit = n - tokens_;
        // ns to earn `deficit` tokens at rate_ bytes/s
        return std::chrono::nanoseconds{
            static_cast<std::int64_t>((deficit * 1'000'000'000ull + rate_ - 1) / rate_)};
    }

    // Two-phase for the hierarchy: check every level, then commit on all or none.
    [[nodiscard]] std::chrono::nanoseconds peek(std::uint64_t n) {
        std::lock_guard lk(mu_);
        if (rate_ == 0)
            return {};
        refill_locked();
        if (tokens_ >= n)
            return {};
        const std::uint64_t deficit = n - tokens_;
        return std::chrono::nanoseconds{
            static_cast<std::int64_t>((deficit * 1'000'000'000ull + rate_ - 1) / rate_)};
    }
    void commit(std::uint64_t n) {
        std::lock_guard lk(mu_);
        if (rate_ == 0)
            return;
        tokens_ = tokens_ >= n ? tokens_ - n : 0;
    }

  private:
    using clock = std::chrono::steady_clock;
    void refill_locked() {
        auto now = clock::now();
        auto dt = std::chrono::duration_cast<std::chrono::nanoseconds>(now - last_).count();
        if (dt <= 0)
            return;
        last_ = now;
        // added = rate_ * dt / 1e9, guarding overflow for very long idle gaps
        long double added =
            static_cast<long double>(rate_) * static_cast<long double>(dt) / 1'000'000'000.0L;
        std::uint64_t add =
            added >= static_cast<long double>(cap_) ? cap_ : static_cast<std::uint64_t>(added);
        tokens_ = tokens_ + add > cap_ ? cap_ : tokens_ + add;
    }

    mutable std::mutex mu_;
    std::uint64_t rate_ = 0;
    std::uint64_t cap_ = 0;
    std::uint64_t tokens_ = 0;
    clock::time_point last_ = clock::now();
};

// The hierarchy. All limits default to 0 (unlimited). A task with no queue is limited by
// task + global only.
class RateLimiter {
  public:
    void set_global_limit(std::uint64_t bps) { global_.set_rate(bps); }
    [[nodiscard]] std::uint64_t global_limit() const { return global_.rate(); }

    void set_queue_limit(QueueId q, std::uint64_t bps) {
        std::lock_guard lk(mu_);
        queues_[q].set_rate(bps);
    }
    void set_task_limit(TaskId t, std::uint64_t bps) {
        std::lock_guard lk(mu_);
        tasks_[t].set_rate(bps);
    }

    void attach_task(TaskId t, std::optional<QueueId> q) {
        std::lock_guard lk(mu_);
        tasks_.try_emplace(t);
        if (q) {
            task_queue_[t] = *q;
            queues_.try_emplace(*q);
        } else {
            task_queue_.erase(t);
        }
    }
    void detach_task(TaskId t) {
        std::lock_guard lk(mu_);
        tasks_.erase(t);
        task_queue_.erase(t);
    }

    // Consume `n` bytes against task, queue and global. 0 => consumed everywhere. > 0 =>
    // consumed nowhere; wait that long and retry. Held under mu_ for its whole duration
    // so a concurrent detach_task() can't invalidate the bucket it is using.
    [[nodiscard]] std::chrono::nanoseconds acquire(TaskId t, std::uint64_t n) {
        std::lock_guard lk(mu_);
        TokenBucket *tb = nullptr;
        TokenBucket *qb = nullptr;
        if (auto it = tasks_.find(t); it != tasks_.end())
            tb = &it->second;
        if (auto qit = task_queue_.find(t); qit != task_queue_.end())
            if (auto q = queues_.find(qit->second); q != queues_.end())
                qb = &q->second;

        // peek all, then commit all or none — no level "leaks" tokens on a partial miss.
        std::chrono::nanoseconds wait{};
        if (tb)
            wait = std::max(wait, tb->peek(n));
        if (qb)
            wait = std::max(wait, qb->peek(n));
        wait = std::max(wait, global_.peek(n));
        if (wait.count() > 0)
            return wait;

        if (tb)
            tb->commit(n);
        if (qb)
            qb->commit(n);
        global_.commit(n);
        return {};
    }

  private:
    mutable std::mutex mu_;  // guards the maps AND serialises acquire()
    TokenBucket global_;
    std::unordered_map<QueueId, TokenBucket> queues_;
    std::unordered_map<TaskId, TokenBucket> tasks_;
    std::unordered_map<TaskId, QueueId> task_queue_;
};

}  // namespace vdm::rate

#endif  // VDM_RATE_TOKEN_BUCKET_HPP