// Copyright (c) 2023, 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.
//
/** @file
 *  @brief Helper functions for dealing with docker containers.  These functions will
 *         fail gracefully on Windows or Mac.
 */
#pragma once

#include "../../carb/Defines.h"

#if CARB_PLATFORM_LINUX
#    include <fcntl.h>
#    include <unistd.h>
#endif


namespace omni
{
namespace extras
{

#if CARB_PLATFORM_LINUX || defined(DOXYGEN_BUILD)

#    ifndef DOXYGEN_SHOULD_SKIP_THIS
namespace detail
{

inline int readIntFromFile(const char* file) noexcept
{
    auto fd = ::open(file, O_RDONLY, 0);
    if (fd == -1)
    {
        return -1;
    }

    char buffer[64];
    auto size = CARB_RETRY_EINTR(::read(fd, buffer, sizeof(buffer) - 1));
    ::close(fd);

    if (size <= 0)
    {
        return -1;
    }

    buffer[size] = '\0';
    return std::atoi(buffer);
}

inline bool isRunningInContainer() noexcept
{
    FILE* fp;

    // first (and easiest) check is to check whether the `/.dockerenv` file exists.  This file
    // is not necessarily always present though.
    if (access("/.dockerenv", F_OK) == 0)
    {
        return true;
    }

    // a more reliable but more expensive check is to verify the control group of `init`.  If
    // running under docker, all of the entries will have a path that starts with `/docker` or
    // `/lxc` instead of just `/`.
    fp = fopen("/proc/1/cgroup", "r");

    if (fp != nullptr)
    {
        char line[256];

        while (fgets(line, CARB_COUNTOF(line), fp) != nullptr)
        {
            if (feof(fp) || ferror(fp))
                break;

            if (strstr(line, ":/docker") != nullptr || strstr(line, ":/lxc") != nullptr)
            {
                return true;
            }
        }

        fclose(fp);
    }

    return false;
}

} // namespace detail
#    endif


/** Attempts to read the current effective CPU usage quota for processes in a container.
 *
 *  @returns The effective number of logical cores that processes in the container will have
 *           access to if a limit has been imposed.  If no limit has been imposed when running
 *           the container, -1 will be returned.
 *
 *  @remarks This reads and calculates the effective CPU usage quota for processes in a container.
 *           If a limit has been imposed, the result is the number of logical cores that can be
 *           used.  Note that this does not actually guarantee that the processes will not be
 *           able to run on some cores, but rather that the CPU scheduler will only give a
 *           certain percentage of its time to processes in the container thereby effectively
 *           making their performance similar to running on a system with the returned number
 *           of cores.
 *
 *  @note This will still return a valid result outside of a container if a CPU usage limit has
 *        been imposed on the system.  By default, Linux systems allow unlimited CPU usage.
 */
inline int getDockerCpuLimit() noexcept
{
    // See:
    // https://docs.docker.com/config/containers/resource_constraints/#cpu
    // https://engineering.squarespace.com/blog/2017/understanding-linux-container-scheduling
    const auto cfsQuota = detail::readIntFromFile("/sys/fs/cgroup/cpu/cpu.cfs_quota_us");
    const auto cfsPeriod = detail::readIntFromFile("/sys/fs/cgroup/cpu/cpu.cfs_period_us");
    if (cfsQuota > 0 && cfsPeriod > 0)
    {
        // Since we can have fractional CPUs, round up half a CPU to a whole CPU, but make sure we have an even number.
        return ::carb_max(1, (cfsQuota + (cfsPeriod / 2)) / cfsPeriod);
    }
    return -1;
}

/** Attempts to detect whether this process is running inside a container.
 *
 *  @returns `true` if this process is running inside a container.  Returns `false` otherwise.
 *
 *  @remarks This detects if the calling process is running inside a container.  This is done
 *           by checking for certain files or their contents that are known (and documented)
 *           to be modified by Docker.  Note that this currently only supports detecting Docker
 *           and LXC containers.
 */
inline bool isRunningInContainer()
{
    static bool s_inContainer = detail::isRunningInContainer();
    return s_inContainer;
}

#else

inline int getDockerCpuLimit() noexcept
{
    return -1;
}

inline bool isRunningInContainer() noexcept
{
    return false;
}

#endif

} // namespace extras
} // namespace omni
