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()
|