policy_tests.cpp

#include "test_utils.hpp"

#include "universal_connection_pool/connection_pool.hpp"

#include <atomic>
#include <chrono>
#include <limits>
#include <memory>
#include <string>

namespace {

struct policy_connection {
    int id = 0;
    std::string endpoint_id;
};

ucp::connection_factory<policy_connection> make_policy_factory(std::atomic<int>& created) {
    ucp::connection_factory<policy_connection> factory;
    factory.create_for_endpoint = [&](const ucp::endpoint_config& endpoint) {
        return policy_connection{++created, endpoint.id};
    };
    return factory;
}

ucp::pool_options make_policy_options() {
    ucp::pool_options options;
    options.max_size = 2;
    options.reaper_interval = std::chrono::milliseconds(0);
    options.histogram_buckets = {std::chrono::hours(1)};
    options.policy_generation = "policy-v1";
    options.feature_gates = {ucp::feature_gate{.name = "stable", .enabled = true}};
    options.endpoints = {
        ucp::endpoint_config{
            .id = "writer",
            .group = "primary",
            .region = "us",
            .role = ucp::endpoint_role::write,
            .weight = 1},
        ucp::endpoint_config{
            .id = "reader",
            .group = "replica",
            .region = "us",
            .role = ucp::endpoint_role::read,
            .weight = 1},
    };
    return options;
}

struct fixed_policy_provider final : ucp::policy_provider {
    ucp::policy_snapshot next;

    ucp::policy_snapshot snapshot() override {
        return next;
    }
};

} // namespace

UCP_TEST(PolicyUpdate_AppliesHotFieldsAndExportsDiagnostics) {
    std::atomic<int> created{0};
    auto options = make_policy_options();
    ucp::connection_pool<policy_connection> pool(make_policy_factory(created), options);

    ucp::performance_options performance;
    performance.profile = ucp::performance_profile::low_latency;
    performance.lock_shard_count = 4;

    ucp::policy_snapshot snapshot;
    snapshot.generation = "policy-v2";
    snapshot.max_size = 1u;
    snapshot.borrow_timeout = std::chrono::milliseconds(5);
    snapshot.routing = ucp::routing_policy::latency_aware;
    snapshot.performance = performance;
    snapshot.replace_tenant_quotas = true;
    snapshot.tenant_quotas = {
        ucp::tenant_quota{
            .tenant_id = "tenant-a",
            .max_active = 1,
            .overflow = ucp::quota_overflow_policy::reject},
    };
    snapshot.feature_gates = {
        ucp::feature_gate{.name = "latency-routing", .enabled = false},
        ucp::feature_gate{.name = "policy-diagnostics", .enabled = true},
    };

    pool.apply_policy_update(snapshot);

    ucp::borrow_options tenant_a;
    tenant_a.tenant_id = "tenant-a";
    auto first = pool.borrow(tenant_a);
    auto rejected = pool.borrow_for(std::chrono::milliseconds(5), tenant_a);

    const auto stats = pool.stats();
    UCP_EXPECT_FALSE(rejected.has_value());
    UCP_EXPECT_EQ(stats.policy_generation, std::string("policy-v2"));
    UCP_EXPECT_EQ(stats.feature_gates.size(), 2u);
    UCP_EXPECT_EQ(stats.wait_time_histogram.count, 0u);

    const auto diagnostics = pool.dump_diagnostics_json();
    UCP_EXPECT_TRUE(diagnostics.find("\"policy_generation\":\"policy-v2\"") != std::string::npos);
    UCP_EXPECT_TRUE(diagnostics.find("\"name\":\"latency-routing\"") != std::string::npos);
    UCP_EXPECT_TRUE(diagnostics.find("\"enabled\":false") != std::string::npos);
    UCP_EXPECT_TRUE(diagnostics.find("\"profile\":\"low_latency\"") != std::string::npos);

    const auto json = pool.dump_json();
    UCP_EXPECT_TRUE(json.find("\"policy_generation\":\"policy-v2\"") != std::string::npos);
    UCP_EXPECT_TRUE(json.find("\"feature_gates\"") != std::string::npos);

    const auto prometheus = pool.to_prometheus("policy");
    UCP_EXPECT_TRUE(prometheus.find("ucp_policy_generation") != std::string::npos);
    UCP_EXPECT_TRUE(prometheus.find("generation=\"policy-v2\"") != std::string::npos);
    UCP_EXPECT_TRUE(prometheus.find("ucp_feature_gate_enabled") != std::string::npos);
}

UCP_TEST(PolicyUpdate_RejectsInvalidPolicyAndKeepsPreviousGeneration) {
    std::atomic<int> created{0};
    auto options = make_policy_options();
    ucp::connection_pool<policy_connection> pool(make_policy_factory(created), options);

    bool rejected_capacity = false;
    try {
        ucp::policy_snapshot invalid;
        invalid.generation = "policy-bad-capacity";
        invalid.max_size = 0u;
        invalid.feature_gates = {ucp::feature_gate{.name = "bad", .enabled = true}};
        pool.apply_policy_update(invalid);
    } catch (const ucp::invalid_pool_options&) {
        rejected_capacity = true;
    }
    UCP_EXPECT_TRUE(rejected_capacity);
    UCP_EXPECT_EQ(pool.stats().policy_generation, std::string("policy-v1"));
    UCP_EXPECT_EQ(pool.stats().feature_gates[0].name, std::string("stable"));

    bool rejected_sampling = false;
    try {
        ucp::performance_options performance;
        performance.metrics_sampling.wait_time_histogram_sample_rate = 0;
        ucp::policy_snapshot invalid;
        invalid.generation = "policy-bad-sampling";
        invalid.performance = performance;
        pool.apply_policy_update(invalid);
    } catch (const ucp::invalid_pool_options&) {
        rejected_sampling = true;
    }
    UCP_EXPECT_TRUE(rejected_sampling);
    UCP_EXPECT_EQ(pool.stats().policy_generation, std::string("policy-v1"));

    auto first = pool.borrow();
    auto second = pool.borrow_for(std::chrono::milliseconds(50));
    UCP_EXPECT_TRUE(second.has_value());
}

UCP_TEST(ConnectionPoolConstructor_UsesSharedOptionValidation) {
    std::atomic<int> created{0};
    auto options = make_policy_options();
    options.max_lifetime = std::chrono::milliseconds(-1);

    bool rejected_lifetime = false;
    try {
        ucp::connection_pool<policy_connection> pool(make_policy_factory(created), options);
        (void)pool;
    } catch (const ucp::invalid_pool_options&) {
        rejected_lifetime = true;
    }
    UCP_EXPECT_TRUE(rejected_lifetime);

    options = make_policy_options();
    options.latency_routing.ewma_alpha = 0.0;
    bool rejected_alpha = false;
    try {
        ucp::connection_pool<policy_connection> pool(make_policy_factory(created), options);
        (void)pool;
    } catch (const ucp::invalid_pool_options&) {
        rejected_alpha = true;
    }
    UCP_EXPECT_TRUE(rejected_alpha);

    options = make_policy_options();
    options.latency_routing.ewma_alpha = std::numeric_limits<double>::quiet_NaN();
    bool rejected_nan_alpha = false;
    try {
        ucp::connection_pool<policy_connection> pool(make_policy_factory(created), options);
        (void)pool;
    } catch (const ucp::invalid_pool_options&) {
        rejected_nan_alpha = true;
    }
    UCP_EXPECT_TRUE(rejected_nan_alpha);

    options = make_policy_options();
    options.latency_routing.error_penalty = -1.0;
    bool rejected_negative_penalty = false;
    try {
        ucp::connection_pool<policy_connection> pool(make_policy_factory(created), options);
        (void)pool;
    } catch (const ucp::invalid_pool_options&) {
        rejected_negative_penalty = true;
    }
    UCP_EXPECT_TRUE(rejected_negative_penalty);

    ucp::policy_snapshot invalid_policy;
    invalid_policy.latency_routing = ucp::latency_routing_options{};
    invalid_policy.latency_routing->inflight_penalty =
        std::numeric_limits<double>::infinity();
    bool rejected_hot_policy = false;
    try {
        ucp::connection_pool<policy_connection> pool(make_policy_factory(created), make_policy_options());
        pool.apply_policy_update(invalid_policy);
    } catch (const ucp::invalid_pool_options&) {
        rejected_hot_policy = true;
    }
    UCP_EXPECT_TRUE(rejected_hot_policy);
}

UCP_TEST(PolicyProvider_AppliesSnapshotOnDemand) {
    std::atomic<int> created{0};
    auto provider = std::make_shared<fixed_policy_provider>();
    provider->next.generation = "policy-provider-v2";
    provider->next.feature_gates = {
        ucp::feature_gate{.name = "provider-gate", .enabled = true},
    };

    auto options = make_policy_options();
    options.policy_provider_source = provider;
    ucp::connection_pool<policy_connection> pool(make_policy_factory(created), options);

    pool.update_policy_from_provider();

    const auto stats = pool.stats();
    UCP_EXPECT_EQ(stats.policy_generation, std::string("policy-provider-v2"));
    UCP_EXPECT_EQ(stats.feature_gates.size(), 1u);
    UCP_EXPECT_EQ(stats.feature_gates[0].name, std::string("provider-gate"));

    fixed_policy_provider direct;
    direct.next.generation = "policy-provider-v3";
    direct.next.feature_gates = {
        ucp::feature_gate{.name = "direct-gate", .enabled = false},
    };
    pool.update_policy_from_provider(direct);

    const auto updated = pool.stats();
    UCP_EXPECT_EQ(updated.policy_generation, std::string("policy-provider-v3"));
    UCP_EXPECT_EQ(updated.feature_gates[0].name, std::string("direct-gate"));
    UCP_EXPECT_FALSE(updated.feature_gates[0].enabled);
}