winsock2_wrapper/test_pool.cpp

/**
 * @file test_pool.cpp
 * @brief ThreadPool + LockFreeQueue 综合测试
 *
 * 测试场景：
 * 1. 基本提交与返回值验证
 * 2. 多生产者并发提交
 * 3. 网络 I/O 场景模拟
 * 4. 关闭安全性验证
 */

#include "pool.hpp"

#include <cassert>
#include <chrono>
#include <iostream>
#include <numeric>
#include <string>
#include <vector>

// ============================================================================
// 测试辅助宏
// ============================================================================

// 轻量断言，失败时打印行号
#define TEST_ASSERT(cond, msg) \
    do { \
        if (!(cond)) { \
            std::cerr << "FAIL [" << __LINE__ << "]: " << msg << std::endl; \
            return false; \
        } \
    } while (0)

// 运行一个测试用例并记录结果
static int g_passed = 0;
static int g_failed = 0;

#define RUN_TEST(name) \
    do { \
        std::cout << "  " << #name << "... "; \
        if (test_##name()) { \
            std::cout << "PASSED" << std::endl; \
            ++g_passed; \
        } else { \
            std::cout << "FAILED" << std::endl; \
            ++g_failed; \
        } \
    } while (0)

// ============================================================================
// 测试用例
// ============================================================================

// ---------------------------------------------------------------------------
// Test 1: 基本提交 + future 返回值
// ---------------------------------------------------------------------------
static bool test_basic_submit() {
    ThreadPool pool(4);

    // 提交一个返回 int 的任务
    auto f1 = pool.submit([] { return 42; });
    TEST_ASSERT(f1.get() == 42, "simple int return");

    // 提交带参数的任务
    auto f2 = pool.submit([](int a, int b) { return a + b; }, 10, 32);
    TEST_ASSERT(f2.get() == 42, "parameterized callable");

    // 提交返回 string 的任务
    auto f3 = pool.submit([](const std::string& s) { return s + " world"; },
                          "hello");
    TEST_ASSERT(f3.get() == "hello world", "string return");

    // 提交 void 任务
    bool called = false;
    auto f4 = pool.submit([&called] { called = true; });
    f4.get();  // wait for completion
    TEST_ASSERT(called, "void task should set flag");

    return true;
}

// ---------------------------------------------------------------------------
// Test 2: 多生产者并发提交
// ---------------------------------------------------------------------------
static bool test_multi_producer() {
    constexpr int num_producers = 8;
    constexpr int tasks_per_producer = 500;

    ThreadPool pool(4);

    // 所有 producer 向同一个计数器累加
    std::atomic<long long> counter{0};

    // 启动多个 producer 线程，每个提交若干任务
    std::vector<std::thread> producers;
    producers.reserve(num_producers);

    for (int p = 0; p < num_producers; ++p) {
        producers.emplace_back([&pool, &counter, tasks_per_producer, p] {
            std::vector<std::future<void>> futures;
            for (int i = 0; i < tasks_per_producer; ++i) {
                auto f = pool.submit([&counter] {
                    // 原子加 1
                    counter.fetch_add(1, std::memory_order_relaxed);
                });
                futures.push_back(std::move(f));
            }
            // 等待此 producer 的所有任务完成
            for (auto& f : futures) {
                f.get();
            }
        });
    }

    // 加入所有 producer
    for (auto& t : producers) {
        t.join();
    }

    // 验证总计数
    long long expected = static_cast<long long>(num_producers) * tasks_per_producer;
    TEST_ASSERT(counter.load() == expected,
        "counter should match total tasks: " << counter.load()
        << " vs " << expected);

    return true;
}

// ---------------------------------------------------------------------------
// Test 3: 网络 I/O 场景模拟
// ---------------------------------------------------------------------------
static bool test_network_scenario() {
    /*
     * 模拟典型的网络服务器任务模式：
     * - 多个连接并发到达
     * - 每个连接经过：接收 → 处理 → 响应 流水线
     * - 使用 future 链接后续处理步骤
     */

    ThreadPool pool(std::thread::hardware_concurrency());

    constexpr int num_connections = 200;

    std::atomic<int> received_count{0};
    std::atomic<int> processed_count{0};
    std::atomic<int> responded_count{0};

    std::vector<std::future<void>> final_futures;
    final_futures.reserve(num_connections);

    for (int conn_id = 0; conn_id < num_connections; ++conn_id) {
        // 模拟处理一个连接：接收 → 处理 → 响应
        auto f = pool.submit([&received_count, &processed_count, &responded_count, conn_id] {
            // Phase 1: 模拟接收数据
            received_count.fetch_add(1, std::memory_order_relaxed);

            // Phase 2: 模拟业务处理
            // 计算一些东西来表示处理
            volatile int result = 0;
            for (int i = 0; i < 100; ++i) {
                result += (conn_id ^ i) & 0xFF;
            }
            (void)result;
            processed_count.fetch_add(1, std::memory_order_relaxed);

            // Phase 3: 模拟发送响应
            responded_count.fetch_add(1, std::memory_order_relaxed);
        });

        final_futures.push_back(std::move(f));
    }

    // 等待所有连接处理完毕
    for (auto& f : final_futures) {
        f.get();
    }

    TEST_ASSERT(received_count.load() == num_connections,
        "all connections received");
    TEST_ASSERT(processed_count.load() == num_connections,
        "all connections processed");
    TEST_ASSERT(responded_count.load() == num_connections,
        "all connections responded");

    return true;
}

// ---------------------------------------------------------------------------
// Test 4: 关闭安全性 — 析构时排空所有已提交任务
// ---------------------------------------------------------------------------
static bool test_shutdown_drain() {
    constexpr int num_tasks = 500;

    std::atomic<int> completed{0};

    {
        ThreadPool pool(4);

        // 提交大量任务（不获取 future，依赖析构排空）
        for (int i = 0; i < num_tasks; ++i) {
            // 不使用返回值，仅依赖析构保证执行
            auto f = pool.submit([&completed] {
                completed.fetch_add(1, std::memory_order_relaxed);
            });
            (void)f;  // 忽略 future，依赖析构
        }

        // pool 在此作用域结束析构，排空所有任务
    }

    // 析构后全部任务应已完成
    TEST_ASSERT(completed.load() == num_tasks,
        "all tasks drained on shutdown: " << completed.load()
        << " vs " << num_tasks);

    return true;
}

// ---------------------------------------------------------------------------
// Test 5: LockFreeQueue 基本功能测试
// ---------------------------------------------------------------------------
static bool test_lockfree_queue() {
    LockFreeQueue<int, 8> q;

    // 空队列状态
    TEST_ASSERT(q.empty(), "queue should be empty initially");
    TEST_ASSERT(!q.full(), "queue should not be full initially");
    TEST_ASSERT(q.size() == 0, "size should be 0 initially");

    // 入队直到满
    for (int i = 0; i < 8; ++i) {
        TEST_ASSERT(q.enqueue(std::move(i)), "enqueue should succeed");
    }
    TEST_ASSERT(q.full(), "queue should be full after 8 enqueues");

    // 入队失败（队列满）
    int extra = 99;
    TEST_ASSERT(!q.enqueue(std::move(extra)), "enqueue should fail when full");

    // 出队所有元素
    for (int i = 0; i < 8; ++i) {
        int val = -1;
        TEST_ASSERT(q.dequeue(val), "dequeue should succeed");
        TEST_ASSERT(val == i, "dequeued value should match: " << val << " vs " << i);
    }
    TEST_ASSERT(q.empty(), "queue should be empty after 8 dequeues");

    // 出队失败（队列空）
    int val = -1;
    TEST_ASSERT(!q.dequeue(val), "dequeue should fail when empty");

    return true;
}

// ---------------------------------------------------------------------------
// Test 6: LockFreeQueue 多线程压力测试
// ---------------------------------------------------------------------------
static bool test_lockfree_queue_threaded() {
    constexpr size_t Q_CAP = 256;
    LockFreeQueue<int, Q_CAP> q;

    constexpr int num_producers = 4;
    constexpr int num_consumers = 4;
    constexpr int items_per_producer = 10000;
    constexpr long long total_items = static_cast<long long>(num_producers) * items_per_producer;

    // 每个 producer 入队的起始值（按 producer 编号偏移，避免重复）
    std::atomic<long long> enqueued_sum{0};
    std::atomic<long long> dequeued_sum{0};
    std::atomic<long long> enqueued_count{0};
    std::atomic<long long> dequeued_count{0};
    std::atomic<bool> producers_done{false};

    // 启动 producers
    std::vector<std::thread> producers;
    for (int p = 0; p < num_producers; ++p) {
        producers.emplace_back([&q, &enqueued_sum, &enqueued_count, items_per_producer, p] {
            for (int i = 0; i < items_per_producer; ++i) {
                int val = p * items_per_producer + i;
                while (!q.enqueue(std::move(val))) {
                    std::this_thread::yield();
                }
                enqueued_sum.fetch_add(val, std::memory_order_relaxed);
                enqueued_count.fetch_add(1, std::memory_order_relaxed);
            }
        });
    }

    // 启动 consumers
    std::vector<std::thread> consumers;
    for (int c = 0; c < num_consumers; ++c) {
        consumers.emplace_back([&q, &dequeued_sum, &dequeued_count, &producers_done] {
            while (true) {
                int val = -1;
                if (q.dequeue(val)) {
                    dequeued_sum.fetch_add(val, std::memory_order_relaxed);
                    dequeued_count.fetch_add(1, std::memory_order_relaxed);
                } else if (producers_done.load(std::memory_order_acquire)) {
                    // 生产者完成且队列空 → 退出
                    // 再尝试一次（可能在 producers_done 设置前瞬间入队）
                    if (q.dequeue(val)) {
                        dequeued_sum.fetch_add(val, std::memory_order_relaxed);
                        dequeued_count.fetch_add(1, std::memory_order_relaxed);
                        continue;
                    }
                    break;
                }
            }
        });
    }

    // 等待所有 producer 完成
    for (auto& t : producers) {
        t.join();
    }
    producers_done.store(true, std::memory_order_release);

    // 等待所有 consumer 完成
    for (auto& t : consumers) {
        t.join();
    }

    // 验证
    TEST_ASSERT(enqueued_count.load() == total_items,
        "all items enqueued");
    TEST_ASSERT(dequeued_count.load() == total_items,
        "all items dequeued: " << dequeued_count.load() << " vs " << total_items);
    TEST_ASSERT(enqueued_sum.load() == dequeued_sum.load(),
        "sums match: " << enqueued_sum.load() << " vs " << dequeued_sum.load());
    TEST_ASSERT(q.empty(), "queue should be empty at end");

    return true;
}

// ============================================================================
// main
// ============================================================================
int main() {
    std::cout << "=== ThreadPool & LockFreeQueue Tests ===" << std::endl;
    std::cout << "Hardware concurrency: "
              << std::thread::hardware_concurrency() << std::endl;
    std::cout << std::endl;

    std::cout << "[LockFreeQueue]" << std::endl;
    RUN_TEST(lockfree_queue);
    RUN_TEST(lockfree_queue_threaded);

    std::cout << std::endl;
    std::cout << "[ThreadPool]" << std::endl;
    RUN_TEST(basic_submit);
    RUN_TEST(multi_producer);
    RUN_TEST(network_scenario);
    RUN_TEST(shutdown_drain);

    std::cout << std::endl;
    std::cout << "=== Results: " << g_passed << " passed, "
              << g_failed << " failed ===" << std::endl;

    return (g_failed == 0) ? 0 : 1;
}