stream_tests.cpp

#include "test_utils.hpp"

#include "universal_connection_pool/connection_pool.hpp"

#include <atomic>
#include <chrono>
#include <string>

namespace {

struct stream_connection {
    int id = 0;
    std::string endpoint_id;
    bool open = true;
    int reset_count = 0;
};

ucp::connection_factory<stream_connection> make_stream_factory(std::atomic<int>& created,
                                                               std::atomic<int>& closed) {
    ucp::connection_factory<stream_connection> factory;
    factory.create_for_endpoint = [&](const ucp::endpoint_config& endpoint) {
        return stream_connection{++created, endpoint.id, true, 0};
    };
    factory.reset = [](stream_connection& connection) {
        ++connection.reset_count;
    };
    factory.close = [&](stream_connection& connection) {
        connection.open = false;
        ++closed;
    };
    return factory;
}

ucp::pool_options stream_options(std::size_t max_size, std::size_t max_streams) {
    ucp::pool_options options;
    options.max_size = max_size;
    options.reaper_interval = std::chrono::milliseconds(0);
    options.borrow_timeout = std::chrono::milliseconds(50);
    options.streams.max_streams_per_connection = max_streams;
    options.endpoints = {ucp::endpoint_config{.id = "grpc-a"}};
    return options;
}

} // namespace

UCP_TEST(StreamBorrow_AllowsMultipleLeasesOnOnePhysicalConnection) {
    std::atomic<int> created{0};
    std::atomic<int> closed{0};
    ucp::connection_pool<stream_connection> pool(make_stream_factory(created, closed),
                                                 stream_options(1, 2));

    auto first = pool.borrow_stream();
    auto second = pool.borrow_stream();

    UCP_EXPECT_EQ(first->id, second->id);
    UCP_EXPECT_EQ(created.load(), 1);
    UCP_EXPECT_EQ(pool.stats().active_connections, 1u);
    UCP_EXPECT_EQ(pool.stats().active_streams, 2u);
    UCP_EXPECT_EQ(pool.stats().borrowed_stream_total, 2u);
    UCP_EXPECT_EQ(pool.stats().wait_time_sample_total, 2u);
    UCP_EXPECT_EQ(pool.stats().wait_time_histogram.count, 2u);
    UCP_EXPECT_TRUE(pool.dump_json().find("\"active_streams\":2") != std::string::npos);
    UCP_EXPECT_TRUE(pool.dump_json().find("\"wait_time_sample_total\":2") != std::string::npos);
    UCP_EXPECT_TRUE(pool.to_prometheus("streams").find("ucp_active_streams") != std::string::npos);
    UCP_EXPECT_TRUE(pool.to_prometheus("streams").find("ucp_wait_time_sample_total") !=
                    std::string::npos);

    auto third = pool.borrow_stream_for(std::chrono::milliseconds(10));
    UCP_EXPECT_FALSE(third.has_value());

    first.reset();
    UCP_EXPECT_EQ(pool.stats().active_connections, 1u);
    UCP_EXPECT_EQ(pool.stats().active_streams, 1u);
    UCP_EXPECT_EQ(pool.stats().idle_connections, 0u);

    second.reset();
    UCP_EXPECT_EQ(pool.stats().active_connections, 0u);
    UCP_EXPECT_EQ(pool.stats().active_streams, 0u);
    UCP_EXPECT_EQ(pool.stats().idle_connections, 1u);
    UCP_EXPECT_EQ(pool.stats().returned_stream_total, 2u);
}

UCP_TEST(StreamBorrowForZeroTimeout_UsesExistingStreamSlot) {
    std::atomic<int> created{0};
    std::atomic<int> closed{0};
    ucp::connection_pool<stream_connection> pool(make_stream_factory(created, closed),
                                                 stream_options(1, 2));

    auto first = pool.borrow_stream();
    auto second = pool.borrow_stream_for(std::chrono::milliseconds(0));

    UCP_EXPECT_TRUE(second.has_value());
    UCP_EXPECT_EQ((*second)->id, first->id);
    UCP_EXPECT_EQ(created.load(), 1);
    UCP_EXPECT_EQ(pool.stats().active_streams, 2u);
}

UCP_TEST(StreamBorrowForZeroTimeout_UsesIdlePhysicalConnectionWithoutCreating) {
    std::atomic<int> created{0};
    std::atomic<int> closed{0};
    ucp::connection_pool<stream_connection> pool(make_stream_factory(created, closed),
                                                 stream_options(1, 2));

    int physical_id = 0;
    {
        auto first = pool.borrow_stream();
        physical_id = first->id;
    }
    UCP_EXPECT_EQ(pool.stats().idle_connections, 1u);

    auto second = pool.borrow_stream_for(std::chrono::milliseconds(0));

    UCP_EXPECT_TRUE(second.has_value());
    UCP_EXPECT_EQ((*second)->id, physical_id);
    UCP_EXPECT_EQ(created.load(), 1);
    UCP_EXPECT_EQ(pool.stats().active_streams, 1u);
}

UCP_TEST(StreamBorrow_DoesNotShareActivePhysicalConnectionAcrossAffinityKeys) {
    std::atomic<int> created{0};
    std::atomic<int> closed{0};
    auto options = stream_options(1, 2);
    options.locality.policy = ucp::locality_policy::prefer_warm;
    options.locality.migration = ucp::migration_policy::allow;
    ucp::connection_pool<stream_connection> pool(make_stream_factory(created, closed), options);

    ucp::borrow_options key_a;
    key_a.affinity_key = "tenant-a";
    auto first = pool.borrow_stream(key_a);

    ucp::borrow_options key_b;
    key_b.affinity_key = "tenant-b";
    auto different_key = pool.try_borrow_stream(key_b);
    UCP_EXPECT_FALSE(different_key.has_value());

    auto same_key = pool.try_borrow_stream(key_a);
    UCP_EXPECT_TRUE(same_key.has_value());
    UCP_EXPECT_EQ((*same_key)->id, first->id);
    UCP_EXPECT_EQ(pool.stats().warm_reuse_total, 1u);
    UCP_EXPECT_EQ(pool.stats().locality_rejected_total, 1u);
}

UCP_TEST(StreamBorrow_PinTokenDoesNotJoinUnpinnedActivePhysicalConnection) {
    std::atomic<int> created{0};
    std::atomic<int> closed{0};
    ucp::connection_pool<stream_connection> pool(make_stream_factory(created, closed),
                                                 stream_options(1, 2));

    auto first = pool.borrow_stream();

    ucp::borrow_options pinned_options;
    pinned_options.pin_token = "session-1";
    auto pinned = pool.try_borrow_stream(pinned_options);

    UCP_EXPECT_FALSE(pinned.has_value());
    UCP_EXPECT_EQ(pool.stats().active_streams, 1u);
    UCP_EXPECT_EQ(created.load(), 1);
}

UCP_TEST(StreamBorrow_CrossTenantReuseDisabledDoesNotShareOrReturnIdle) {
    std::atomic<int> created{0};
    std::atomic<int> closed{0};
    auto options = stream_options(1, 2);
    options.protocol.allow_cross_tenant_reuse = false;
    ucp::connection_pool<stream_connection> pool(make_stream_factory(created, closed), options);

    ucp::borrow_options tenant_a;
    tenant_a.tenant_id = "tenant-a";
    auto first = pool.borrow_stream(tenant_a);

    ucp::borrow_options tenant_b;
    tenant_b.tenant_id = "tenant-b";
    auto second = pool.try_borrow_stream(tenant_b);
    UCP_EXPECT_FALSE(second.has_value());

    first.reset();
    UCP_EXPECT_EQ(pool.stats().active_streams, 0u);
    UCP_EXPECT_EQ(pool.stats().idle_connections, 0u);
    UCP_EXPECT_EQ(closed.load(), 1);
}

UCP_TEST(StreamBorrow_AtStreamCapacityCreatesSecondPhysicalConnection) {
    std::atomic<int> created{0};
    std::atomic<int> closed{0};
    ucp::connection_pool<stream_connection> pool(make_stream_factory(created, closed),
                                                 stream_options(2, 1));

    auto first = pool.borrow_stream();
    auto second = pool.borrow_stream();

    UCP_EXPECT_TRUE(first->id != second->id);
    UCP_EXPECT_EQ(created.load(), 2);
    UCP_EXPECT_EQ(pool.stats().active_connections, 2u);
    UCP_EXPECT_EQ(pool.stats().active_streams, 2u);
}

UCP_TEST(StreamBorrow_EndpointCapacityOverridesPoolDefault) {
    std::atomic<int> created{0};
    std::atomic<int> closed{0};
    auto options = stream_options(1, 4);
    options.endpoints = {ucp::endpoint_config{.id = "grpc-a", .stream_capacity = 1}};
    ucp::connection_pool<stream_connection> pool(make_stream_factory(created, closed), options);

    auto first = pool.borrow_stream();
    auto second = pool.try_borrow_stream();

    UCP_EXPECT_FALSE(second.has_value());
    UCP_EXPECT_EQ(pool.stats().active_streams, 1u);
}

UCP_TEST(StreamBorrow_DrainConnectionPreventsNewStreamsAndRetiresOnReturn) {
    std::atomic<int> created{0};
    std::atomic<int> closed{0};
    ucp::connection_pool<stream_connection> pool(make_stream_factory(created, closed),
                                                 stream_options(1, 2));

    auto first = pool.borrow_stream();
    const auto first_id = first->id;
    first.drain_connection();

    auto second = pool.try_borrow_stream();
    UCP_EXPECT_FALSE(second.has_value());

    first.reset();
    UCP_EXPECT_EQ(pool.stats().active_connections, 0u);
    UCP_EXPECT_EQ(pool.stats().active_streams, 0u);
    UCP_EXPECT_EQ(pool.stats().idle_connections, 0u);
    UCP_EXPECT_EQ(closed.load(), 1);

    auto replacement = pool.borrow_stream();
    UCP_EXPECT_TRUE(replacement->id != first_id);
    UCP_EXPECT_EQ(created.load(), 2);
}