Jungfraujoch/common/HostMemory.cpp

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

#include "HostMemory.h"

#include <algorithm>

#if defined(_WIN32)
#ifndef NOMINMAX
#define NOMINMAX
#endif
#include <windows.h>
#elif defined(__APPLE__)
#include <cstdint>
#include <sys/types.h>
#include <sys/sysctl.h>
#else
#include <unistd.h>
#endif

size_t get_host_total_memory() {
#if defined(_WIN32)
    MEMORYSTATUSEX status{};
    status.dwLength = sizeof(status);
    if (!GlobalMemoryStatusEx(&status))
        return 0;
    return static_cast<size_t>(status.ullTotalPhys);
#elif defined(__APPLE__)
    uint64_t bytes = 0;
    size_t len = sizeof(bytes);
    if (sysctlbyname("hw.memsize", &bytes, &len, nullptr, 0) != 0)
        return 0;
    return static_cast<size_t>(bytes);
#else
    const long pages = sysconf(_SC_PHYS_PAGES), page = sysconf(_SC_PAGE_SIZE);
    if (pages <= 0 || page <= 0)
        return 0;
    return static_cast<size_t>(pages) * static_cast<size_t>(page);
#endif
}

size_t DefaultHostMemoryBudget() {
    return static_cast<size_t>(HOST_MEMORY_BUDGET_FRACTION * static_cast<double>(get_host_total_memory()));
}

HostWorkerPool &HostWorkerPool::Instance() {
    static HostWorkerPool pool;
    return pool;
}

HostWorkerPool::HostWorkerPool() {
    SetBudget(0);
    SetSize(0);
}

void HostWorkerPool::SetSize(size_t bytes) {
    std::lock_guard lock(m);
    size = bytes > 0 ? bytes
                     : static_cast<size_t>(HOST_WORKER_POOL_FRACTION * static_cast<double>(DefaultHostMemoryBudget()));
    cv.notify_all();
}

void HostWorkerPool::SetBudget(size_t bytes) {
    std::lock_guard lock(m);
    budget = bytes > 0 ? bytes : DefaultHostMemoryBudget();
    cv.notify_all();
}

size_t HostWorkerPool::available() const {
    return std::min(size, budget > held ? budget - held : 0);
}

size_t HostWorkerPool::Size() {
    std::lock_guard lock(m);
    return available();
}

HostWorkerPool::Reservation::Reservation(HostWorkerPool &pool, size_t bytes) : pool(pool), bytes(bytes) {
    std::unique_lock lock(pool.m);
    pool.cv.wait(lock, [&] { return pool.reserved == 0 || pool.reserved + bytes <= pool.available(); });
    pool.reserved += bytes;
}

HostWorkerPool::Reservation::~Reservation() {
    {
        std::lock_guard lock(pool.m);
        pool.reserved -= bytes;
    }
    pool.cv.notify_all();
}

HostWorkerPool::Held::Held(HostWorkerPool &pool, size_t bytes) : pool(pool), bytes(0) {
    Set(bytes);
}

HostWorkerPool::Held::~Held() {
    Set(0);
}

void HostWorkerPool::Held::Set(size_t new_bytes) {
    {
        std::lock_guard lock(pool.m);
        pool.held = pool.held - bytes + new_bytes;
        bytes = new_bytes;
    }
    pool.cv.notify_all();
}