task_group_storage_test.cpp

#include <gtest/gtest.h>

import std;
import mcpplibs.cmp;

namespace {

struct alignas(std::max_align_t) AllocationHeader final {
    std::size_t size_ {};
};

std::atomic<std::size_t> gLiveBytes {};

}  // namespace

void* operator new(std::size_t size) {
    const auto actual = size == 0 ? 1 : size;
    auto* memory = static_cast<AllocationHeader*>(
        std::malloc(sizeof(AllocationHeader) + actual));

    if (!memory) {
        throw std::bad_alloc {};
    }

    memory->size_ = actual;
    gLiveBytes.fetch_add(actual, std::memory_order_relaxed);
    return memory + 1;
}

void operator delete(void* pointer) noexcept {
    if (pointer) {
        auto* memory = static_cast<AllocationHeader*>(pointer) - 1;
        gLiveBytes.fetch_sub(memory->size_, std::memory_order_relaxed);
        std::free(memory);
    }
}

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

namespace {

using namespace mcpplibs::cmp;

struct Retention final {
    std::size_t before_ {};
    std::size_t peak_ {};
    std::size_t after_ {};
    int completed_ {};
    bool caught_ {};
};

Task<void> immediate_child(int& completed, bool fail) {
    ++completed;

    if (fail) {
        throw std::runtime_error { "first failure" };
    }

    co_return;
}

Task<Retention> measure_retention(int count, bool fail) {
    TaskGroup group {};
    Retention retention {};
    auto joining = group.join();
    retention.before_ = gLiveBytes.load();

    for (int index {}; index < count; ++index) {
        group.spawn(immediate_child(retention.completed_, fail && index == 0));
    }

    retention.peak_ = gLiveBytes.load();

    try {
        co_await std::move(joining);
    } catch (const std::runtime_error& error) {
        retention.caught_ = std::string_view { error.what() } == "first failure";
    }

    retention.after_ = gLiveBytes.load();
    co_return retention;
}

TEST(CmpTaskGroupStorageTest, JoinReleasesTerminalFramesBeforeGroupDestruction) {
    for (const int count : { 1'000, 10'000, 50'000 }) {
        for (const bool fail : { false, true }) {
            RunLoop loop {};
            const auto retention = loop.run(measure_retention(count, fail));
            EXPECT_EQ(retention.completed_, count);
            EXPECT_EQ(retention.caught_, fail);
            EXPECT_GT(retention.peak_, retention.before_);
            EXPECT_LE(retention.after_, retention.before_);
            std::println("storage count={} failure={} before={} peak={} after={}",
                count, fail, retention.before_, retention.peak_, retention.after_);
        }
    }
}

}  // namespace