vdm/core/tests/rate/token_bucket_test.cpp

#include "vdm/rate/token_bucket.hpp"

#include <atomic>
#include <chrono>
#include <thread>
#include <vector>

#include "vtest.hpp"

using namespace vdm;
using namespace vdm::rate;
using namespace std::chrono_literals;

namespace {
TaskId tid(std::uint64_t v) {
    return TaskId{v};
}
QueueId qid(std::uint64_t v) {
    return QueueId{v};
}
}  // namespace

VT_TEST(tb_unlimited_never_waits) {
    TokenBucket b(0);
    for (int i = 0; i < 1000; ++i)
        VT_CHECK_EQ(b.consume(1'000'000).count(), 0);
}

VT_TEST(tb_burst_then_throttle) {
    // 1000 B/s, default burst = 1 s = 1000 tokens.
    TokenBucket b(1000);
    VT_CHECK_EQ(b.consume(1000).count(), 0);  // drains the burst
    auto w = b.consume(1000);                 // empty now: must wait ~1 s
    VT_CHECK(w >= 900ms && w <= 1100ms);
}

VT_TEST(tb_refills_over_time) {
    TokenBucket b(10'000, /*burst=*/10'000);
    VT_CHECK_EQ(b.consume(10'000).count(), 0);
    std::this_thread::sleep_for(120ms);  // ~1200 tokens back
    auto w = b.consume(1000);
    VT_CHECK_EQ(w.count(), 0);  // affordable from the refill
    auto w2 = b.consume(5000);
    VT_CHECK(w2.count() > 0);  // not that much yet
}

VT_TEST(tb_burst_caps_accumulation) {
    TokenBucket b(1000, /*burst=*/2000);
    std::this_thread::sleep_for(100ms);  // idle far longer than burst/rate would fill
    std::this_thread::sleep_for(100ms);
    VT_CHECK_EQ(b.consume(2000).count(), 0);  // at most the 2000 cap accumulated
    VT_CHECK(b.consume(1).count() > 0);       // and no more
}

VT_TEST(tb_set_rate_zero_makes_unlimited) {
    TokenBucket b(1000);
    VT_CHECK_EQ(b.consume(1000).count(), 0);
    VT_CHECK(b.consume(1000).count() > 0);
    b.set_rate(0);
    VT_CHECK_EQ(b.consume(1'000'000).count(), 0);
}

// --- the hierarchy --------------------------------------------------------------------

VT_TEST(rl_all_unlimited_by_default) {
    RateLimiter rl;
    rl.attach_task(tid(1), std::nullopt);
    for (int i = 0; i < 100; ++i)
        VT_CHECK_EQ(rl.acquire(tid(1), 1'000'000).count(), 0);
}

VT_TEST(rl_tightest_level_binds) {
    RateLimiter rl;
    rl.set_global_limit(100'000);
    rl.set_queue_limit(qid(9), 20'000);
    rl.set_task_limit(tid(1), 50'000);
    rl.attach_task(tid(1), qid(9));

    // burst: task 50k, queue 20k, global 100k -> the queue's 20k is the ceiling
    VT_CHECK_EQ(rl.acquire(tid(1), 20'000).count(), 0);
    auto w = rl.acquire(tid(1), 5'000);
    VT_CHECK(w.count() > 0);  // queue bucket is dry even though task & global aren't
}

VT_TEST(rl_no_partial_consumption_on_miss) {
    RateLimiter rl;
    rl.set_global_limit(1'000'000);   // plenty
    rl.set_task_limit(tid(1), 1000);  // 1 s burst
    rl.attach_task(tid(1), std::nullopt);

    VT_CHECK_EQ(rl.acquire(tid(1), 1000).count(), 0);  // drain the task bucket
    for (int i = 0; i < 5; ++i)
        VT_CHECK(rl.acquire(tid(1), 1000).count() > 0);  // task bucket blocks, repeatedly

    // global must NOT have been charged for any of those blocked attempts: a fresh task
    // limited only by the global bucket can still spend nearly its whole burst (only the
    // one *successful* 1000-byte acquire above was charged).
    rl.attach_task(tid(2), std::nullopt);
    VT_CHECK_EQ(rl.acquire(tid(2), 990'000).count(), 0);
}

VT_TEST(rl_detach_then_acquire_is_safe_and_unlimited) {
    RateLimiter rl;
    rl.set_task_limit(tid(1), 1000);
    rl.attach_task(tid(1), std::nullopt);
    VT_CHECK_EQ(rl.acquire(tid(1), 1000).count(), 0);
    rl.detach_task(tid(1));
    // unknown task -> no task/queue bucket, only global (unlimited here)
    VT_CHECK_EQ(rl.acquire(tid(1), 1'000'000).count(), 0);
}

VT_TEST(rl_enforces_aggregate_rate_under_load) {
    RateLimiter rl;
    const std::uint64_t rate = 4'000'000;  // 4 MB/s global
    rl.set_global_limit(rate);
    for (std::uint64_t i = 1; i <= 8; ++i)
        rl.attach_task(tid(i), std::nullopt);

    std::atomic<std::uint64_t> moved{0};
    auto t0 = std::chrono::steady_clock::now();
    std::vector<std::jthread> ws;
    for (std::uint64_t i = 1; i <= 8; ++i) {
        ws.emplace_back([&, id = tid(i)] {
            for (int k = 0; k < 400; ++k) {
                std::uint64_t chunk = 16 * 1024;
                for (;;) {
                    auto w = rl.acquire(id, chunk);
                    if (w.count() == 0)
                        break;
                    std::this_thread::sleep_for(
                        std::min<std::chrono::nanoseconds>(w, std::chrono::milliseconds(20)));
                }
                moved.fetch_add(chunk);
            }
        });
    }
    ws.clear();  // join
    auto secs = std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();

    double effective = moved.load() / secs;
    // Allow one burst's worth of slop plus scheduling noise: effective rate should sit
    // within ~2x of the configured limit, never wildly above.
    VT_CHECK(effective <= rate * 2.5);
    VT_CHECK(moved.load() == 8u * 400u * 16u * 1024u);
}