#!/usr/bin/python3
# Copyright 2024 Helmut Grohne <helmut@subdivi.de>
# SPDX-License-Identifier: GPL-3

"""Estimate a suitable concurrency for a workload requiring significant amounts
of RAM from the environment.
"""

import argparse
import itertools
import logging
import os
import pathlib
import subprocess
import sys
import typing


_MANUAL_PAGE = """
=head1 NAME

guess_concurrency - Estimate suitable workload concurrency

=head1 SYNOPSIS

B<guess_concurrency> [B<-h>] [B<--detect> I<METHOD>] [B<--max> I<N>] [B<--min> I<N>] [B<--require-mem> I<MEMLIST>]

=head1 DESCRIPTION

Guess a suitable concurrency level given constraints from the runtime
environment. Environment variables such as B<CMAKE_BUILD_PARALLEL_LEVEL>,
B<DEB_BUILD_OPTIONS> and B<RPM_BUILD_NCPUS> are consulted for initial
guesses falling back to querying the operating system in order to guess a
suitable concurrency level. A user may further restrict the emitted concurrency
by specifying a minimum or maximum and by requiring sufficient memory to
support a level of concurrency to be available.

=head1 OPTIONS

=over 8

=item B<--debug>

additionally print diagnostics to stderr about what was detected and how.

=item B<--detect> I<METHODS>

supply a comma-separated sequence of detection methods (nproc, python or cores)
to be tried in order or a CPU count. Note that environment variables take
precedence over these methods and a fallback count.

=item B<--max> I<N>

limit the number of detected cores to a given maximum

=item B<--min> I<N>

limit the number of detected cores to a given minimum

=item B<--require-mem> I<MEMLIST>

specify per-core required memory as a comma separated list. Each entry
indicates how much additional RAM is needed for the next process. Numbers can
be given as floating point number optionally suffixed with a case-insensitive
binary suffix (k/m/g/t/e). A completely empty spot copies the previous number.
A concurrency I<N> is considered acceptable if the sum of the first I<N>
entries of this list does not exceed the available RAM.

=back

=head1 EXAMPLES

Ensure that there is at least 4GB per used core.

    guess_concurrency --require-mem 4G

Two linker processes may consume 6GB each and they may run concurrently. All
other concurrent processes are expected to not use more than 2GB each.

    guess_concurrency --require-mem 6G,,2G

=head1 CAVEATS

In Debian, this tool can be used by source packages to prevent build failures
arising from exhausting memory of the build machine. It should not be used for
packages that generally succeed building without it. Pessimistic use of the
B<--require-mem> option can significantly lengthen the time to build a package
and thus may become a problem to build daemon capacity especially for slower
architectures. If linking is a problem, consider alternative mechanisms. The
ninja build system provides a pool mechanism. When using it with CMake, you can
configure it via B<CMAKE_JOB_POOL_LINK>. When using it with Meson, you can set
the B<backend_max_links> option.

=head1 SEE ALSO

    dpkg-buildpackage(1)

=cut
"""


logger = logging.getLogger("guess_concurrency")


def positive_integer(decstr: str) -> int:
    """Parse a positive, integral number from a string and raise a ValueError
    otherwise.

    >>> positive_integer(-1)
    Traceback (most recent call last):
        ...
    ValueError: integer must be positive
    >>> positive_integer(0)
    Traceback (most recent call last):
        ...
    ValueError: integer must be positive
    >>> positive_integer(1)
    1
    """
    value = int(decstr)  # raises ValueError
    if value < 1:
        raise ValueError("integer must be positive")
    return value


def parse_size(expression: str) -> int:
    """Parse an expression representing a data size with an optional unit
    suffix into an integer. Raises a ValueError on failure.

    >>> parse_size("5")
    5
    >>> parse_size("4KB")
    4096
    >>> parse_size("0.9g")
    966367641
    >>> parse_size("-1")
    Traceback (most recent call last):
        ...
    ValueError: number must be positive
    """
    expression = expression.lower()
    if expression.endswith("b"):
        expression = expression[:-1]
    suffixes = {
        "k": 2**10,
        "m": 2**20,
        "g": 2**30,
        "t": 2**40,
        "e": 2**50,
    }
    factor = 1
    if expression[-1:] in suffixes:
        factor = suffixes[expression[-1]]
        expression = expression[:-1]
    fval = float(expression)  # propagate ValueError
    value = int(fval * factor)  # propagate ValueError
    if value < 1:
        raise ValueError("number must be positive")
    return value


def guess_python() -> int | None:
    """Estimate the number of processors using Python's standard library
    functions.

    >>> guess_python() > 0
    True
    """
    process_cpu_count = getattr(os, "process_cpu_count", None)
    if process_cpu_count is not None:
        count = process_cpu_count()
        assert isinstance(count, int) or count is None  # make mypy happy
        return count
    # Fallback for Python < 3.13.
    return len(os.sched_getaffinity(0))


def guess_nproc() -> int:
    """Estimate number of processors using coreutils' nproc.

    >>> guess_nproc() > 0
    True
    """
    return positive_integer(
        subprocess.check_output(["nproc"], encoding="ascii")
    )


def guess_cores() -> int | None:
    """Estimate the number of cores (not SMT threads).  This is done by
    counting the number of distinct "core id" values found in sysfs.
    """
    cpus = set(
        int(fn.read_text())
        for fn in pathlib.Path("/sys/devices/system/cpu").glob(
            "cpu*/topology/core_id"
        )
    )
    if not cpus:
        return None
    return len(cpus)


def guess_deb_build_parallel(
    environ: typing.Mapping[str, str] = os.environ
) -> int | None:
    """Parse a possible parallel= assignment in a DEB_BUILD_OPTIONS environment
    variable.

    >>> guess_deb_build_parallel({})
    >>> guess_deb_build_parallel({"DEB_BUILD_OPTIONS": "nocheck parallel=3"})
    3
    """
    try:
        options = environ["DEB_BUILD_OPTIONS"]
    except KeyError:
        return None
    for option in options.split():
        if option.startswith("parallel="):
            option = option.removeprefix("parallel=")
            try:
                return positive_integer(option)
            except ValueError:
                pass
    return None


def guess_from_environment(
    variable: str, environ: typing.Mapping[str, str] = os.environ
) -> int | None:
    """Read a number from an environment variable.

    >>> guess_from_environment("CPUS", {"CPUS": 4})
    4
    >>> guess_from_environment("CPUS", {"other": 3})
    """
    try:
        return positive_integer(environ[variable])
    except (KeyError, ValueError):
        return None


def guess_memavailable() -> int:
    """Estimate the available memory from /proc/meminfo in bytes."""
    fields: dict[str, int | None] = {
        "MemAvailable": None,
        "Shmem": None,
        "SwapFree": None,
    }
    with open("/proc/meminfo", encoding="ascii") as fh:
        for line in fh:
            try:
                key, value = line.split(":", 1)
            except ValueError:
                continue
            if key not in fields:
                continue
            mult = 1
            value = value.strip()
            if value.endswith("kB"):
                mult = 1024
                value = value.removesuffix("kB")
            try:
                amount = int(value.strip())
            except ValueError:
                continue
            if amount < 0:
                continue
            fields[key] = amount * mult
    # This is the amount of free and reclaimable memory. The reclaiming here
    # covers evicting page cache with unmodified disk copies.
    result = fields["MemAvailable"]
    if result is None:
        raise RuntimeError("no MemAvailable line found in /proc/meminfo")
    if fields["Shmem"] is None or fields["SwapFree"] is None:
        return result
    # While tmpfs usage is also considered page cache, it is not considered
    # reclaimable by linux, but it kinda is as it can be swapped out. This
    # is the typical setup on Debian buildds.
    return result + min(fields["Shmem"], fields["SwapFree"])


def guess_cgroup_memory() -> int | None:
    """Return the smallest "memory.high" or "memory.max" limit of the current
    cgroup or any parent of it if any.
    """
    guess: int | None = None
    mygroup = pathlib.PurePath(
        pathlib.Path("/proc/self/cgroup")
        .read_text(encoding=sys.getfilesystemencoding())
        .strip()
        .split(":", 2)[2]
    ).relative_to("/")
    sfc = pathlib.Path("/sys/fs/cgroup")
    for group in itertools.chain((mygroup,), mygroup.parents):
        for entry in ("memory.max", "memory.high"):
            try:
                value = positive_integer(
                    (sfc / group / entry).read_text(encoding="ascii")
                )
            except (FileNotFoundError, ValueError):
                pass
            else:
                if guess is None:
                    guess = value
                else:
                    guess = min(guess, value)
    return guess


def parse_required_memory(expression: str) -> list[int]:
    """Parse comma-separated list of memory expressions. Empty expressions copy
    the previous entry.

    >>> parse_required_memory("1k,9,,1")
    [1024, 9, 9, 1]
    """
    values: list[int] = []
    for memexpr in expression.split(","):
        if not memexpr:
            if values:
                values.append(values[-1])
            else:
                raise ValueError("initial memory expression cannot be empty")
        else:
            values.append(parse_size(memexpr))
    return values


def guess_memory_concurrency(memory: int, usage: list[int]) -> int:
    """Estimate the maximum number of cores that can be used given the
    available memory and a sequence of per-core memory consumption.

    >>> guess_memory_concurrency(4, [1])
    4
    >>> guess_memory_concurrency(10, [5, 4, 3])
    2
    >>> guess_memory_concurrency(2, [3])
    1
    """
    concurrency = 0
    for use in usage[:-1]:
        if use > memory:
            break
        memory -= use
        concurrency += 1
    else:
        concurrency += memory // usage[-1]
    return max(1, concurrency)


def clamp(
    value: int, lower: int | None = None, upper: int | None = None
) -> int:
    """Return an adjusted value that does not exceed the lower or upper limits
    if any.

    >>> clamp(5, upper=4)
    4
    >>> clamp(9, 2)
    9
    """
    if upper is not None and upper < value:
        value = upper
    if lower is not None and lower > value:
        value = lower
    return value


def get_argparser() -> argparse.ArgumentParser:
    """Construct the ArgumentParser for the CLI."""
    parser = argparse.ArgumentParser(
        description="""

Guess a suitable concurrency level given constraints from the runtime
environment. Environment variables such as DEB_BUILD_OPTIONS, RPM_BUILD_NCPUS
and CMAKE_BUILD_PARALLEL_LEVEL are consulted for initial guesses falling back
to querying the operating system in order to guess a suitable concurrency
level. A user may further restrict the emitted concurrency by specifying a
minimum or maximum and by requiring sufficient memory to support a level of
concurrency to be available.

""".strip(),
    )
    parser.add_argument(
        "--detect",
        action="store",
        default="nproc",
        metavar="METHOD",
        help="supply a processor count or select a detection method "
        "(nproc, python or cores)",
    )
    parser.add_argument(
        "--max",
        action="store",
        type=positive_integer,
        default=None,
        metavar="N",
        help="limit the number of detected cores to a given maximum",
    )
    parser.add_argument(
        "--min",
        action="store",
        type=positive_integer,
        default=None,
        metavar="N",
        help="limit the number of detected cores to a given minimum",
    )
    parser.add_argument(
        "--require-mem",
        action="store",
        type=parse_required_memory,
        default=[],
        metavar="MEMLIST",
        help="specify per-core required memory as a comma separated list",
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        default=False,
        help="output detected values to stderr for understanding the result",
    )
    return parser


def main() -> None:
    """Command line interface entry point."""
    parser = get_argparser()
    args = parser.parse_args()
    if args.min is not None and args.max is not None and args.min > args.max:
        parser.error("--min value larger than --max value")
    if args.debug:
        logging.basicConfig(
            stream=sys.stderr, level=logging.DEBUG, format="%(message)s"
        )

    guess = None
    guessers: list[tuple[typing.Callable[[], int | None], str]] = [
        (
            lambda: guess_from_environment("CMAKE_BUILD_PARALLEL_LEVEL"),
            "CMAKE_BUILD_PARALLEL_LEVEL",
        ),
        (guess_deb_build_parallel, "DEB_BUILD_OPTIONS"),
        (lambda: guess_from_environment("RPM_BUILD_NCPUS"), "RPM_BUILD_NCPUS"),
    ]
    for guess_function, description in guessers:
        guess = guess_function()
        if guess is not None:
            logger.debug("discovered cpu count %d from %s", guess, description)
            break
    else:
        for detector in args.detect.split(","):
            try:
                guess = positive_integer(detector)
            except ValueError:
                try:
                    guess = {
                        "nproc": guess_nproc,
                        "python": guess_python,
                        "cores": guess_cores,
                    }[detector]()
                except KeyError:
                    parser.error("invalid argument to --detect")
                else:
                    if guess is not None:
                        logger.debug(
                            "discovered cpu count %d from %s", guess, detector
                        )
                        break
            else:
                logger.debug("user provided cpu count %d", guess)
                break
        if guess is None:
            print("failed to guess processor count", file=sys.stderr)
            sys.exit(1)
    if args.require_mem and guess > 1:
        memory = clamp(guess_memavailable(), upper=guess_cgroup_memory())
        logger.debug("discovered %d MB available memory", memory >> 20)
        guess = clamp(
            guess, upper=guess_memory_concurrency(memory, args.require_mem)
        )
    guess = clamp(guess, args.min, args.max)
    print(guess)


if __name__ == "__main__":
    main()
