Jungfraujoch/tests/HostMemoryTest.cpp

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <catch2/catch_all.hpp>

#include <algorithm>
#include <atomic>
#include <chrono>
#include <thread>
#include <tuple>
#include <vector>

#include "../common/HostMemory.h"

TEST_CASE("HostMemory_BudgetFromTheTotal", "[HostMemory]") {
    const size_t total = get_host_total_memory();
    REQUIRE(total > 0);
    CHECK(DefaultHostMemoryBudget() == static_cast<size_t>(HOST_MEMORY_BUDGET_FRACTION * static_cast<double>(total)));
}

// As many reservations at once as the pool holds, and one at a time where a single one is larger than it.
TEST_CASE("HostWorkerPool_BoundsWorkersByBytes", "[HostMemory]") {
    HostWorkerPool &pool = HostWorkerPool::Instance();
    for (const auto &[size, bytes, at_most] : std::vector<std::tuple<size_t, size_t, int>>{
             {100, 40, 2}, {100, 100, 1}, {100, 1000, 1}, {1000, 10, 8}}) {
        CAPTURE(size, bytes);
        pool.SetSize(size);
        std::atomic<int> running{0}, most{0}, done{0};
        std::vector<std::thread> workers;
        for (int i = 0; i < 8; i++)
            workers.emplace_back([&] {
                const HostWorkerPool::Reservation reserved(pool, bytes);
                const int now = ++running;
                int seen = most.load();
                while (now > seen && !most.compare_exchange_weak(seen, now)) {}
                std::this_thread::sleep_for(std::chrono::milliseconds(20));
                --running;
                ++done;
            });
        for (auto &w : workers)
            w.join();
        CHECK(done == 8);
        CHECK(most >= 1);
        CHECK(most <= at_most);
    }
    pool.SetSize(0);
    CHECK(pool.Size() == static_cast<size_t>(HOST_WORKER_POOL_FRACTION * static_cast<double>(DefaultHostMemoryBudget())));
}

// One ledger: what the run holds (Held) and what the workers reserve together stay within the budget. Workers
// take what the Held bytes leave of it, up to the pool's size, and a worker larger than that runs alone rather
// than never.
TEST_CASE("HostWorkerPool_HeldBytesShareTheBudget", "[HostMemory]") {
    HostWorkerPool &pool = HostWorkerPool::Instance();
    pool.SetBudget(100);
    pool.SetSize(50);
    CHECK(pool.Size() == 50);
    const auto most_at_once = [&](size_t bytes) {
        std::atomic<int> running{0}, most{0}, done{0};
        std::vector<std::thread> workers;
        for (int i = 0; i < 8; i++)
            workers.emplace_back([&] {
                const HostWorkerPool::Reservation reserved(pool, bytes);
                const int now = ++running;
                int seen = most.load();
                while (now > seen && !most.compare_exchange_weak(seen, now)) {}
                std::this_thread::sleep_for(std::chrono::milliseconds(20));
                --running;
                ++done;
            });
        for (auto &w : workers)
            w.join();
        CHECK(done == 8);
        return most.load();
    };
    {
        HostWorkerPool::Held run(pool, 70);   // leaves 30 of the budget: the pool's 50 does not fit beside it
        CHECK(pool.Size() == 30);
        CHECK(most_at_once(10) <= 3);
        run.Set(20);                          // leaves 80: the pool's own size bounds again
        CHECK(pool.Size() == 50);
        CHECK(most_at_once(10) <= 5);
        run.Set(100);                         // nothing left: one at a time, each still runs
        CHECK(pool.Size() == 0);
        CHECK(most_at_once(10) == 1);
    }
    CHECK(pool.Size() == 50);
    pool.SetBudget(0);
    pool.SetSize(0);
}