add test for scheduler queue budget

This commit is contained in:
serxa 2024-06-14 13:12:23 +00:00
parent a5eeeb3422
commit ec71d35aa6

View File

@ -1,11 +1,13 @@
#include <gtest/gtest.h>
#include <Common/Scheduler/SchedulerRoot.h>
#include <Common/Scheduler/Nodes/tests/ResourceTest.h>
#include <Common/Scheduler/SchedulerRoot.h>
#include <Common/randomSeed.h>
#include <barrier>
#include <future>
#include <pcg_random.hpp>
using namespace DB;
@ -22,6 +24,17 @@ struct ResourceTest : public ResourceTestBase
{
scheduler.stop(true);
}
std::mutex rng_mutex;
pcg64 rng{randomSeed()};
template <typename T>
T randomInt(T from, T to)
{
std::uniform_int_distribution<T> distribution(from, to);
std::lock_guard lock(rng_mutex);
return distribution(rng);
}
};
struct ResourceHolder
@ -133,6 +146,31 @@ TEST(SchedulerRoot, Smoke)
}
}
TEST(SchedulerRoot, Budget)
{
ResourceTest t;
ResourceHolder r1(t);
r1.add<ConstraintTest>("/", "<max_requests>1</max_requests>");
r1.add<PriorityPolicy>("/prio");
auto a = r1.addQueue("/prio/A", "");
r1.registerResource();
ResourceCost total_real_cost = 0;
int total_requests = 10;
for (int i = 0 ; i < total_requests; i++)
{
ResourceCost est_cost = t.randomInt(1, 10);
ResourceCost real_cost = t.randomInt(0, 10);
ResourceGuard rg(ResourceGuard::Metrics::getIOWrite(), a, est_cost);
rg.consume(real_cost);
total_real_cost += real_cost;
}
EXPECT_EQ(total_requests, a.queue->dequeued_requests);
EXPECT_EQ(total_real_cost, a.queue->dequeued_cost - a.queue->getBudget());
}
TEST(SchedulerRoot, Cancel)
{
ResourceTest t;