carb/extras/MemoryUsage.h

File members: carb/extras/MemoryUsage.h

// Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto. Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
//
#pragma once

#include "../Defines.h"

#include "../logging/Log.h"

#if CARB_PLATFORM_LINUX
#    include <sys/sysinfo.h>

#    include <stdio.h>
#    include <unistd.h>
#    include <sys/time.h>
#    include <sys/resource.h>
#elif CARB_PLATFORM_MACOS
#    include <mach/task.h>
#    include <mach/mach_init.h>
#    include <mach/mach_host.h>
#    include <os/proc.h>
#    include <sys/resource.h>
#elif CARB_PLATFORM_WINDOWS
#    include "../CarbWindows.h"
#endif

namespace carb
{
namespace extras
{

enum class MemoryQueryType
{
    eAvailable,
    eTotal,
};

size_t getPhysicalMemory(MemoryQueryType type); // forward declare

inline size_t getCurrentProcessMemoryUsage()
{
#if CARB_PLATFORM_LINUX
    unsigned long rss = 0;
    long pageSize;

    pageSize = sysconf(_SC_PAGESIZE);
    if (pageSize < 0)
    {
        CARB_LOG_ERROR("failed to retrieve the page size");
        return 0;
    }

    // fopen() and fgets() can use the heap, and we may be called from the crashreporter, so avoid using those
    // functions to avoid heap usage.
    auto fd = open("/proc/self/statm", 0, O_RDONLY);
    if (fd < 0)
    {
        CARB_LOG_ERROR("failed to open /proc/self/statm");
        return 0;
    }

    char line[256];
    auto readBytes = CARB_RETRY_EINTR(read(fd, line, CARB_COUNTOF(line) - 1));
    if (readBytes > 0)
    {
        char* endp;
        // Skip the first, read the second
        strtoul(line, &endp, 10);
        rss = strtoul(endp, nullptr, 10);
    }

    close(fd);

    return rss * pageSize;
#elif CARB_PLATFORM_WINDOWS
    CARBWIN_PROCESS_MEMORY_COUNTERS mem;
    mem.cb = sizeof(mem);
    if (!K32GetProcessMemoryInfo(GetCurrentProcess(), (PPROCESS_MEMORY_COUNTERS)&mem, sizeof(mem)))
    {
        CARB_LOG_ERROR("GetProcessMemoryInfo failed");
        return 0;
    }

    return mem.WorkingSetSize;
#elif CARB_PLATFORM_MACOS
    mach_msg_type_number_t count = TASK_BASIC_INFO_COUNT;
    task_basic_info info = {};
    kern_return_t r = task_info(mach_task_self(), TASK_BASIC_INFO, reinterpret_cast<task_info_t>(&info), &count);
    if (r != KERN_SUCCESS)
    {
        CARB_LOG_ERROR("task_info() failed (%d)", int(r));
        return 0;
    }
    return info.resident_size;
#else
#    warning "getMemoryUsage() has no implementation"
    return 0;
#endif
}

inline size_t getPeakProcessMemoryUsage()
{
#if CARB_PLATFORM_WINDOWS
    CARBWIN_PROCESS_MEMORY_COUNTERS mem;
    mem.cb = sizeof(mem);
    if (!K32GetProcessMemoryInfo(GetCurrentProcess(), (PPROCESS_MEMORY_COUNTERS)&mem, sizeof(mem)))
    {
        CARB_LOG_ERROR("GetProcessMemoryInfo failed");
        return 0;
    }

    return mem.PeakWorkingSetSize;
#elif CARB_POSIX
    rusage usage_data;
    size_t scale;
    getrusage(RUSAGE_SELF, &usage_data);

#    if CARB_PLATFORM_LINUX
    scale = 1024; // Linux provides this value in kilobytes.
#    elif CARB_PLATFORM_MACOS
    scale = 1; // MacOS provides this value in bytes.
#    else
    CARB_UNSUPPORTED_PLATFORM();
#    endif

    // convert to bytes.
    return usage_data.ru_maxrss * scale;
#else
    CARB_UNSUPPORTED_PLATFORM();
#endif
}

struct SystemMemoryInfo
{
    uint64_t totalPhysical;

    uint64_t availablePhysical;

    uint64_t totalPageFile;

    uint64_t availablePageFile;

    uint64_t totalVirtual;

    uint64_t availableVirtual;
};

#if CARB_PLATFORM_LINUX
inline size_t getMemorySizeMultiplier(const char* str)
{
    size_t multiplier = 1;

    // strip leading whitespace.
    while (*str == ' ' || *str == '\t')
        str++;

    // check the prefix of the multiplier string (ie: "kB", "gB", etc).
    switch (tolower(*str))
    {
        case 'e':
            multiplier *= 1024ull; // fall through...
            CARB_FALLTHROUGH;

        case 'p':
            multiplier *= 1024ull; // fall through...
            CARB_FALLTHROUGH;

        case 't':
            multiplier *= 1024ull; // fall through...
            CARB_FALLTHROUGH;

        case 'g':
            multiplier *= 1024ull; // fall through...
            CARB_FALLTHROUGH;

        case 'm':
            multiplier *= 1024ull; // fall through...
            CARB_FALLTHROUGH;

        case 'k':
            multiplier *= 1024ull; // fall through...
            CARB_FALLTHROUGH;

        default:
            break;
    }

    return multiplier;
}

inline size_t getMemoryValueByName(const char* filename, const char* name, size_t nameLen = 0)
{
    constexpr static char kProcMemInfo[] = "/proc/meminfo";
    size_t bytes = 0;
    char line[256];
    size_t nameLength = nameLen;

    if (filename == nullptr)
        filename = kProcMemInfo;

    if (nameLength == 0)
        nameLength = strlen(name);

    // fopen() and fgets() can use the heap, and we may be called from the crashreporter, so avoid using those
    // functions to avoid heap usage.
    auto fd = open(filename, 0, O_RDONLY);
    if (fd < 0)
        return 0;

    ssize_t readBytes;
    while ((readBytes = CARB_RETRY_EINTR(read(fd, line, CARB_COUNTOF(line) - 1))) > 0)
    {
        line[readBytes] = '\0';

        auto lf = strchr(line, '\n');
        if (lf)
        {
            *lf = '\0';
            // Seek back to the start of the next line for the next read
            lseek(fd, -off_t(line + readBytes - lf) + 1, SEEK_CUR);
        }

        if (strncmp(line, name, nameLength) == 0)
        {
            // Found the key that we're looking for => parse its value in Kibibytes and succeed.
            char* endp;
            bytes = strtoull(line + nameLength, &endp, 10);
            bytes *= getMemorySizeMultiplier(endp);
            break;
        }
    }

    close(fd);
    return bytes;
}
#endif

inline bool getSystemMemoryInfo(SystemMemoryInfo& out)
{
#if CARB_PLATFORM_LINUX
    struct sysinfo info = {};
    struct rlimit limit = {};
    int result;
    size_t bytes;

    // collect the total memory counts.
    result = sysinfo(&info);

    if (result != 0)
    {
        CARB_LOG_WARN("sysinfo() returned %d", result);

        // retrieve the values from '/proc/meminfo' instead.
        out.totalPhysical = getMemoryValueByName(nullptr, "MemTotal:", sizeof("MemTotal:") - 1);
        out.totalPageFile = getMemoryValueByName(nullptr, "SwapTotal:", sizeof("SwapTotal:") - 1);
    }

    else
    {
        out.totalPhysical = (uint64_t)info.totalram * info.mem_unit;
        out.totalPageFile = (uint64_t)info.totalswap * info.mem_unit;
    }

    // get the virtual memory info.
    if (getrlimit(RLIMIT_AS, &limit) == 0)
    {
        out.totalVirtual = limit.rlim_cur;
        out.availableVirtual = 0;

        // retrieve the total VM usage for the calling process.
        bytes = getMemoryValueByName("/proc/self/status", "VmSize:", sizeof("VmSize:") - 1);

        if (bytes != 0)
        {
            if (bytes > out.totalVirtual)
            {
                CARB_LOG_WARN(
                    "retrieved a larger VM size than total VM space (!?) {bytes = %zu, "
                    "totalVirtual = %" PRIu64 "}",
                    bytes, out.totalVirtual);
            }

            else
                out.availableVirtual = out.totalVirtual - bytes;
        }
    }

    else
    {
        CARB_LOG_WARN("failed to retrieve the total address space {errno = %d/%s}", errno, strerror(errno));
        out.totalVirtual = 0;
        out.availableVirtual = 0;
    }

    // collect the available RAM as best we can.  The values in '/proc/meminfo' are typically
    // more accurate than what sysinfo() gives us due to the 'mem_unit' value.
    bytes = getMemoryValueByName(nullptr, "MemAvailable:", sizeof("MemAvailable:") - 1);

    if (bytes != 0)
        out.availablePhysical = bytes;

    else
        out.availablePhysical = (uint64_t)info.freeram * info.mem_unit;

    // collect the available swap space as best we can.
    bytes = getMemoryValueByName(nullptr, "SwapFree:", sizeof("SwapFree:") - 1);

    if (bytes != 0)
        out.availablePageFile = bytes;

    else
        out.availablePageFile = (uint64_t)info.freeswap * info.mem_unit;

    return true;
#elif CARB_PLATFORM_WINDOWS
    CARBWIN_MEMORYSTATUSEX status;
    status.dwLength = sizeof(status);

    if (!GlobalMemoryStatusEx((LPMEMORYSTATUSEX)&status))
    {
        CARB_LOG_ERROR("GlobalMemoryStatusEx() failed {error = %d}", GetLastError());
        return false;
    }

    out.totalPhysical = (uint64_t)status.ullTotalPhys;
    out.totalPageFile = (uint64_t)status.ullTotalPageFile;
    out.totalVirtual = (uint64_t)status.ullTotalVirtual;
    out.availablePhysical = (uint64_t)status.ullAvailPhys;
    out.availablePageFile = (uint64_t)status.ullAvailPageFile;
    out.availableVirtual = (uint64_t)status.ullAvailVirtual;

    return true;
#elif CARB_PLATFORM_MACOS
    int mib[2];
    size_t length;
    mach_msg_type_number_t count;
    kern_return_t r;
    xsw_usage swap = {};
    task_basic_info info = {};
    struct rlimit limit = {};

    // get the system's swap usage
    mib[0] = CTL_HW, mib[1] = VM_SWAPUSAGE;
    length = sizeof(swap);
    if (sysctl(mib, CARB_COUNTOF(mib), &swap, &length, nullptr, 0) != 0)
    {
        CARB_LOG_ERROR("sysctl() for VM_SWAPUSAGE failed (errno = %d)", errno);
        return false;
    }

    count = TASK_BASIC_INFO_COUNT;
    r = task_info(mach_task_self(), TASK_BASIC_INFO, reinterpret_cast<task_info_t>(&info), &count);
    if (r != KERN_SUCCESS)
    {
        CARB_LOG_ERROR("task_info() failed (result = %d, errno = %d)", int(r), errno);
        return false;
    }

    // it's undocumented but RLIMIT_AS is supported
    if (getrlimit(RLIMIT_AS, &limit) != 0)
    {
        CARB_LOG_ERROR("getrlimit(RLIMIT_AS) failed (errno = %d)", errno);
        return false;
    }

    out.totalVirtual = limit.rlim_cur;
    out.availableVirtual = out.totalVirtual - info.virtual_size;
    out.totalPhysical = getPhysicalMemory(MemoryQueryType::eTotal);
    out.availablePhysical = getPhysicalMemory(MemoryQueryType::eAvailable);
    out.totalPageFile = swap.xsu_total;
    out.availablePageFile = swap.xsu_avail;

    return true;
#else
#    warning "getSystemMemoryInfo() has no implementation"
    return 0;
#endif
}

inline size_t getPhysicalMemory(MemoryQueryType type)
{
#if CARB_PLATFORM_LINUX
    // this is a linux-specific system call
    struct sysinfo info;
    int result;
    const char* search;
    size_t searchLength;
    size_t bytes;

    // attempt to read the available memory from '/proc/meminfo' first.
    if (type == MemoryQueryType::eTotal)
    {
        search = "MemTotal:";
        searchLength = sizeof("MemTotal:") - 1;
    }

    else
    {
        search = "MemAvailable:";
        searchLength = sizeof("MemAvailable:") - 1;
    }

    bytes = getMemoryValueByName(nullptr, search, searchLength);

    if (bytes != 0)
        return bytes;

    // fall back to sysinfo() to get the amount of free RAM if it couldn't be found in
    // '/proc/meminfo'.
    result = sysinfo(&info);
    if (result != 0)
    {
        CARB_LOG_ERROR("sysinfo() returned %d", result);
        return 0;
    }

    if (type == MemoryQueryType::eTotal)
        return info.totalram * info.mem_unit;
    else
        return info.freeram * info.mem_unit;
#elif CARB_PLATFORM_WINDOWS
    CARBWIN_MEMORYSTATUSEX status;
    status.dwLength = sizeof(status);

    if (!GlobalMemoryStatusEx((LPMEMORYSTATUSEX)&status))
    {
        CARB_LOG_ERROR("GlobalMemoryStatusEx failed");
        return 0;
    }

    if (type == MemoryQueryType::eTotal)
        return status.ullTotalPhys;
    else
        return status.ullAvailPhys;
#elif CARB_PLATFORM_MACOS
    int mib[2];
    size_t memSize = 0;
    size_t length;

    if (type == MemoryQueryType::eTotal)
    {
        mib[0] = CTL_HW, mib[1] = HW_MEMSIZE;
        length = sizeof(memSize);
        if (sysctl(mib, CARB_COUNTOF(mib), &memSize, &length, nullptr, 0) != 0)
        {
            CARB_LOG_ERROR("sysctl() for HW_MEMSIZE failed (errno = %d)", errno);
            return false;
        }
    }
    else
    {
        mach_msg_type_number_t count;
        kern_return_t r;
        vm_statistics_data_t vm = {};
        size_t pageSize = getpagesize();

        count = HOST_VM_INFO_COUNT;
        r = host_statistics(mach_host_self(), HOST_VM_INFO, reinterpret_cast<host_info_t>(&vm), &count);
        if (r != KERN_SUCCESS)
        {
            CARB_LOG_ERROR("host_statistics() failed (%d)", int(r));
            return false;
        }

        memSize = (vm.free_count + vm.inactive_count) * pageSize;
    }

    return memSize;
#else
#    warning "getPhysicalMemoryAvailable() has no implementation"
    return 0;
#endif
}

enum class MemoryScaleType
{
    eBinaryScale,

    eDecimalScale,
};

inline double getFriendlyMemorySize(size_t bytes, const char** suffix, MemoryScaleType scale = MemoryScaleType::eBinaryScale)
{
    constexpr size_t kEib = 1024ull * 1024 * 1024 * 1024 * 1024 * 1024;
    constexpr size_t kPib = 1024ull * 1024 * 1024 * 1024 * 1024;
    constexpr size_t kTib = 1024ull * 1024 * 1024 * 1024;
    constexpr size_t kGib = 1024ull * 1024 * 1024;
    constexpr size_t kMib = 1024ull * 1024;
    constexpr size_t kKib = 1024ull;
    constexpr size_t kEb = 1000ull * 1000 * 1000 * 1000 * 1000 * 1000;
    constexpr size_t kPb = 1000ull * 1000 * 1000 * 1000 * 1000;
    constexpr size_t kTb = 1000ull * 1000 * 1000 * 1000;
    constexpr size_t kGb = 1000ull * 1000 * 1000;
    constexpr size_t kMb = 1000ull * 1000;
    constexpr size_t kKb = 1000ull;
    constexpr size_t limits[2][6] = { { kEib, kPib, kTib, kGib, kMib, kKib }, { kEb, kPb, kTb, kGb, kMb, kKb } };
    constexpr const char* suffixes[2][6] = { { "EiB", "PiB", "TiB", "GiB", "MiB", "KiB" },
                                             { "EB", "PB", "TB", "GB", "MB", "KB" } };

    if (scale != MemoryScaleType::eBinaryScale && scale != MemoryScaleType::eDecimalScale)
    {
        *suffix = "bytes";
        return (double)bytes;
    }

    for (size_t i = 0; i < CARB_COUNTOF(limits[(size_t)scale]); i++)
    {
        size_t limit = limits[(size_t)scale][i];

        if (bytes >= limit)
        {
            *suffix = suffixes[(size_t)scale][i];
            return (double)bytes / (double)limit;
        }
    }

    *suffix = "bytes";
    return (double)bytes;
}

} // namespace extras
} // namespace carb