Viewing File: /opt/imunify360/venv/versions/imunify-core-8.12.1-2/defence360agent/utils/resource_limits.py

import asyncio
import logging

from contextlib import suppress
from enum import Enum
from os import fsdecode
from pathlib import Path
from typing import List, Optional, Tuple

from defence360agent.utils import OsReleaseInfo

logger = logging.getLogger(__name__)


RUN_WITH_INTENSITY = "/usr/libexec/run-with-intensity"

LVECTL_BIN_PATH = Path("/usr/sbin/lvectl")
PROC_LVE_LIST_PATH = Path("/proc/lve/list")
PROC_PATH = Path("/proc")
CGROUP_PATH = Path("/sys/fs/cgroup")

# cgroup v1 reports "no limit" as a page counter near its maximum
_CGROUP_V1_NO_LIMIT = 1 << 50


class LimitsMethod(Enum):
    NICE = "nice"
    LVE = "lve"
    CGROUPS = "cgroups"


async def get_current_method() -> LimitsMethod:
    """Returns limit method, used in run-with-intensity tool."""
    proc = await asyncio.create_subprocess_exec(
        RUN_WITH_INTENSITY,
        "show",
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )

    stdout, stderr = await proc.communicate()
    stdout = fsdecode(stdout).strip()

    if stdout == "nice":
        return LimitsMethod.NICE
    if stdout == "lve":
        return LimitsMethod.LVE
    if stdout == "cgroups":
        return LimitsMethod.CGROUPS
    raise LookupError(
        "Parsing of used limitation method failed\nstdout: {}\nstderr: {}"
        .format(stdout, fsdecode(stderr).strip())
    )


async def create_subprocess(
    cmd: List[str],
    key: str,
    intensity_cpu: int,
    intensity_io: int,
    **subprocess_kwargs
) -> asyncio.subprocess.Process:
    """
    Creates asyncio.Process with limited resources (cpu & io),
    using run-with-intensity tool.

    :param cmd: command to execute
    :param intensity_cpu: cpu intensity limit
    :param intensity_io: io intensity limit
    :param subprocess_kwargs: keyword arguments for create_subprocess_exec func
    :return: executed Process
    """
    limits_cmd = [
        RUN_WITH_INTENSITY,
        "run",
        "--intensity-cpu",
        str(intensity_cpu),
        "--intensity-io",
        str(intensity_io),
    ]
    limits_cmd.extend(["--key", key])

    return await asyncio.create_subprocess_exec(
        *(limits_cmd + cmd), **subprocess_kwargs
    )


def _status_field_kb(pid: int, field: str) -> Optional[int]:
    try:
        status = (PROC_PATH / str(pid) / "status").read_text()
    except OSError:
        return None
    for line in status.splitlines():
        name, _, value = line.partition(":")
        if name == field:
            with suppress(IndexError, ValueError):
                return int(value.split()[0])
    return None


def _child_pids(pid: int) -> List[int]:
    children: List[int] = []
    try:
        threads = sorted((PROC_PATH / str(pid) / "task").iterdir())
    except OSError:
        return children
    for thread in threads:
        try:
            listed = (thread / "children").read_text()
        except OSError:
            continue
        with suppress(ValueError):
            children.extend(int(child) for child in listed.split())
    return children


def _process_tree(pid: int) -> List[int]:
    tree, pending = [], [pid]
    while pending:
        current = pending.pop()
        tree.append(current)
        pending.extend(_child_pids(current))
    return tree


def peak_rss_kb(pid: int) -> Optional[int]:
    """Peak RSS of the process and its descendants, summed, in kB."""
    total = None
    for process in _process_tree(pid):
        peak = _status_field_kb(process, "VmHWM")
        if peak is not None:
            total = (total or 0) + peak
    return total


def _cgroup_limit_bytes(path: Path) -> Optional[int]:
    try:
        value = path.read_text().strip()
    except OSError:
        return None
    if value == "max":
        return None
    with suppress(ValueError):
        limit = int(value)
        if limit < _CGROUP_V1_NO_LIMIT:
            return limit
    return None


def memory_bound_kb(pid: int) -> Optional[Tuple[int, str]]:
    """The cgroup memory limit the process runs under, in kB, and its source."""
    try:
        cgroups = (PROC_PATH / str(pid) / "cgroup").read_text()
    except OSError:
        return None
    for line in cgroups.splitlines():
        _, _, rest = line.partition(":")
        controllers, _, cgroup = rest.partition(":")
        relative = cgroup.lstrip("/")
        if not controllers:
            limit = _cgroup_limit_bytes(CGROUP_PATH / relative / "memory.max")
            source = "cgroup v2"
        elif "memory" in controllers.split(","):
            limit = _cgroup_limit_bytes(
                CGROUP_PATH / "memory" / relative / "memory.limit_in_bytes"
            )
            source = "cgroup v1"
        else:
            continue
        if limit is not None:
            return limit // 1024, source
    return None


def is_lve_active() -> bool:
    """Checks that LVE-utils is active resource limiter."""
    # to avoid possible errors such as DEF-11941
    # make sure that OS is CL
    return PROC_LVE_LIST_PATH.exists() and OsReleaseInfo.is_cloudlinux()


def has_lvectl() -> bool:
    """Checks that LVE-utils is installed."""
    return LVECTL_BIN_PATH.exists()