run_loop_failure_test.cpp

#include <gtest/gtest.h>

import std;
import mcpplibs.cmp;

namespace {

std::atomic<bool> gFailAllocation { false };
std::atomic<int> gRejectedAllocation {};

}  // namespace

void* operator new(std::size_t size) {
    if (gFailAllocation.load(std::memory_order_relaxed)) {
        gRejectedAllocation.fetch_add(1, std::memory_order_relaxed);
        throw std::bad_alloc {};
    }

    if (auto* memory = std::malloc(size == 0 ? 1 : size)) {
        return memory;
    }

    throw std::bad_alloc {};
}

void operator delete(void* memory) noexcept {
    std::free(memory);
}

void operator delete(void* memory, std::size_t) noexcept {
    std::free(memory);
}

namespace {

using namespace mcpplibs::cmp;
using Scheduler = RunLoop::Scheduler;

class AllocationFailure final {
public:
    AllocationFailure() noexcept {
        gFailAllocation.store(true, std::memory_order_relaxed);
    }

    ~AllocationFailure() {
        gFailAllocation.store(false, std::memory_order_relaxed);
    }
};

Task<void> count_timer(
    Scheduler scheduler,
    Scheduler::TimePoint deadline,
    int& completed) {
    co_await scheduler.schedule_at(deadline);
    ++completed;
}

Task<int> drain_admitted_timers(Scheduler scheduler) {
    TaskGroup group {};
    int completed {};
    const auto deadline = Scheduler::Clock::now() + std::chrono::milliseconds(100);

    for (int index {}; index < 4'096; ++index) {
        group.spawn(count_timer(scheduler, deadline, completed));
    }

    // join 帧在故障开启前创建，故障仅作用于已接纳工作的推进
    auto joining = group.join();
    AllocationFailure failure {};
    co_await std::move(joining);
    co_return completed;
}

Task<void> fail_timer_admission(Scheduler scheduler, bool& caught) {
    AllocationFailure failure {};

    try {
        co_await scheduler.schedule_after(std::chrono::seconds(1));
    } catch (const std::bad_alloc&) {
        caught = true;
    }
}

Task<void> queued_without_allocation(Scheduler scheduler, int& completed) {
    AllocationFailure failure {};
    co_await scheduler.schedule();
    ++completed;
    co_await scheduler.schedule_after(Scheduler::Duration::zero());
    ++completed;
}

Task<void> return_from_worker(
    ThreadPool::Scheduler worker,
    Scheduler caller,
    std::latch& entered,
    std::binary_semaphore& release,
    int& completed) {
    co_await worker.schedule();
    entered.count_down();
    release.acquire();
    co_await caller.schedule();
    ++completed;
}

Task<int> drain_mixed_children(Scheduler scheduler, ThreadPool::Scheduler worker) {
    TaskGroup group {};
    int completed {};
    std::latch entered { 1 };
    std::binary_semaphore release { 0 };
    group.spawn(return_from_worker(worker, scheduler, entered, release, completed));
    const auto deadline = Scheduler::Clock::now() + std::chrono::milliseconds(100);

    for (int index {}; index < 4'096; ++index) {
        group.spawn(count_timer(scheduler, deadline, completed));
    }

    auto joining = group.join();
    entered.wait();
    AllocationFailure failure {};
    release.release();
    co_await std::move(joining);
    co_return completed;
}

Task<void> cancelled_timer(Scheduler scheduler, std::stop_token token, int& completed) {
    try {
        co_await scheduler.schedule_at(Scheduler::TimePoint::max(), token);
    } catch (const OperationCancelled&) {
        ++completed;
    }
}

Task<int> cancel_admitted_timers(Scheduler scheduler) {
    TaskGroup group {};
    std::stop_source stop {};
    int completed {};

    for (int index {}; index < 4'096; ++index) {
        group.spawn(cancelled_timer(scheduler, stop.get_token(), completed));
    }

    auto joining = group.join();
    AllocationFailure failure {};
    stop.request_stop();
    co_await std::move(joining);
    co_return completed;
}

Task<int> value();

TEST(CmpRunLoopFailureTest, TimersAndExternalWorkersDrainWithoutAllocating) {
    ThreadPool pool { 1 };
    RunLoop loop {};
    gRejectedAllocation.store(0);
    EXPECT_EQ(loop.run(drain_mixed_children(
        loop.get_scheduler(), pool.get_scheduler())), 4'097);
    EXPECT_EQ(gRejectedAllocation.load(), 0);
}

TEST(CmpRunLoopFailureTest, CancellationAndPromotionDoNotAllocate) {
    RunLoop loop {};
    gRejectedAllocation.store(0);
    EXPECT_EQ(loop.run(cancel_admitted_timers(loop.get_scheduler())), 4'096);
    EXPECT_EQ(gRejectedAllocation.load(), 0);
    EXPECT_EQ(loop.run(value()), 42);
}

Task<int> value() {
    co_return 42;
}

Task<void> fail_admission_after_children(
    Scheduler scheduler,
    std::stop_source& stop,
    std::latch& entered,
    std::binary_semaphore& release,
    bool& caught) {
    entered.wait();
    {
        AllocationFailure failure {};
        try {
            co_await scheduler.schedule_at(Scheduler::TimePoint::max());
        } catch (const std::bad_alloc&) {
            caught = true;
        }
    }
    stop.request_stop();
    release.release();
}

Task<int> admission_failure_with_children(
    Scheduler scheduler,
    ThreadPool::Scheduler worker,
    bool tuple,
    bool& caught) {
    int completed {};
    std::stop_source stop {};
    std::latch entered { 1 };
    std::binary_semaphore release { 0 };

    if (tuple) {
        co_await when_all(
            cancelled_timer(scheduler, stop.get_token(), completed),
            cancelled_timer(scheduler, stop.get_token(), completed),
            return_from_worker(worker, scheduler, entered, release, completed),
            fail_admission_after_children(scheduler, stop, entered, release, caught));
    } else {
        TaskGroup group {};
        auto joining = group.join();
        group.spawn(cancelled_timer(scheduler, stop.get_token(), completed));
        group.spawn(cancelled_timer(scheduler, stop.get_token(), completed));
        group.spawn(return_from_worker(worker, scheduler, entered, release, completed));
        group.spawn(fail_admission_after_children(scheduler, stop, entered, release, caught));
        co_await std::move(joining);
    }
    co_return completed;
}

TEST(CmpRunLoopFailureTest, GroupAdmissionFailureDrainsAlreadyAcceptedTimersAndWorker) {
    ThreadPool pool { 1 };
    RunLoop loop {};
    bool caught {};
    gRejectedAllocation.store(0);
    EXPECT_EQ(loop.run(admission_failure_with_children(
        loop.get_scheduler(), pool.get_scheduler(), false, caught)), 3);
    EXPECT_TRUE(caught);
    EXPECT_EQ(gRejectedAllocation.load(), 1);
    EXPECT_EQ(loop.run(value()), 42);
}

TEST(CmpRunLoopFailureTest, WhenAllAdmissionFailureDrainsAlreadyAcceptedTimersAndWorker) {
    ThreadPool pool { 1 };
    RunLoop loop {};
    bool caught {};
    gRejectedAllocation.store(0);
    EXPECT_EQ(loop.run(admission_failure_with_children(
        loop.get_scheduler(), pool.get_scheduler(), true, caught)), 3);
    EXPECT_TRUE(caught);
    EXPECT_EQ(gRejectedAllocation.load(), 1);
    EXPECT_EQ(loop.run(value()), 42);
}

TEST(CmpRunLoopFailureTest, AdmittedTimersDrainWithoutAllocating) {
    RunLoop loop {};
    gRejectedAllocation.store(0);
    EXPECT_EQ(loop.run(drain_admitted_timers(loop.get_scheduler())), 4'096);
    EXPECT_EQ(gRejectedAllocation.load(), 0);
    EXPECT_EQ(loop.run(value()), 42);
}

TEST(CmpRunLoopFailureTest, TimerAdmissionFailureIsCatchableAndReusable) {
    RunLoop loop {};
    bool caught {};
    gRejectedAllocation.store(0);
    loop.run(fail_timer_admission(loop.get_scheduler(), caught));
    EXPECT_TRUE(caught);
    EXPECT_EQ(gRejectedAllocation.load(), 1);
    EXPECT_EQ(loop.run(value()), 42);
}

TEST(CmpRunLoopFailureTest, ReadyAndImmediateTimerQueuesDoNotAllocate) {
    RunLoop loop {};
    int completed {};
    gRejectedAllocation.store(0);
    loop.run(queued_without_allocation(loop.get_scheduler(), completed));
    EXPECT_EQ(completed, 2);
    EXPECT_EQ(gRejectedAllocation.load(), 0);
}

}  // namespace