avoid executing external tools by using Python's built-in os module to interact with the filesystem directly

This commit is contained in:
Sergey Krashevich 2023-04-28 01:33:37 +03:00
parent 83006eeb65
commit 2e549514a5
No known key found for this signature in database
GPG Key ID: 625171324E7D3856

View File

@ -9,6 +9,7 @@ import signal
import traceback import traceback
import urllib.parse import urllib.parse
import yaml import yaml
import os
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from collections import Counter from collections import Counter
@ -740,92 +741,71 @@ def escape_special_characters(path: str) -> str:
def get_cgroups_version() -> str: def get_cgroups_version() -> str:
"""Determine what version of cgroups is enabled""" """Determine what version of cgroups is enabled."""
stat_command = ["stat", "-fc", "%T", "/sys/fs/cgroup"] cgroup_path = "/sys/fs/cgroup"
p = sp.run( try:
stat_command, stat_info = os.stat(cgroup_path)
encoding="ascii", value = os.statvfs(cgroup_path).f_type
capture_output=True,
)
if p.returncode == 0: if value == os.fsencode('cgroup2fs'):
value: str = p.stdout.strip().lower()
if value == "cgroup2fs":
return "cgroup2" return "cgroup2"
elif value == "tmpfs": elif value == os.fsencode('tmpfs'):
return "cgroup" return "cgroup"
else: else:
logger.debug( logger.debug(f"Could not determine cgroups version: unhandled filesystem {value}")
f"Could not determine cgroups version: unhandled filesystem {value}" except Exception as e:
) logger.debug(f"Could not determine cgroups version: {e}")
else:
logger.debug(f"Could not determine cgroups version: {p.stderr}")
return "unknown" return "unknown"
def get_docker_memlimit_bytes() -> int: def get_docker_memlimit_bytes() -> int:
"""Get mem limit in bytes set in docker if present. Returns -1 if no limit detected""" """Get mem limit in bytes set in docker if present. Returns -1 if no limit detected."""
# check running a supported cgroups version # check running a supported cgroups version
if get_cgroups_version() == "cgroup2": if get_cgroups_version() == "cgroup2":
memlimit_command = ["cat", "/sys/fs/cgroup/memory.max"] memlimit_path = "/sys/fs/cgroup/memory.max"
p = sp.run( try:
memlimit_command, with open(memlimit_path, 'r') as f:
encoding="ascii", value = f.read().strip()
capture_output=True,
)
if p.returncode == 0:
value: str = p.stdout.strip()
if value.isnumeric(): if value.isnumeric():
return int(value) return int(value)
elif value.lower() == "max": elif value.lower() == "max":
return -1 return -1
else: except Exception as e:
logger.debug(f"Unable to get docker memlimit: {p.stderr}") logger.debug(f"Unable to get docker memlimit: {e}")
return -1 return -1
def get_cpu_stats() -> dict[str, dict]: def get_cpu_stats() -> dict[str, dict]:
"""Get cpu usages for each process id""" """Get cpu usages for each process id."""
usages = {} usages = {}
# -n=2 runs to ensure extraneous values are not included
top_command = ["top", "-b", "-n", "2"]
docker_memlimit = get_docker_memlimit_bytes() / 1024 docker_memlimit = get_docker_memlimit_bytes() / 1024
p = sp.run( for pid in os.listdir('/proc'):
top_command, if pid.isdigit():
encoding="ascii",
capture_output=True,
)
if p.returncode != 0:
logger.error(p.stderr)
return usages
else:
lines = p.stdout.split("\n")
for line in lines:
stats = list(filter(lambda a: a != "", line.strip().split(" ")))
try: try:
if docker_memlimit > 0: with open(f'/proc/{pid}/stat', 'r') as f:
mem_res = int(stats[5]) stat_info = f.read().split()
mem_pct = str(
round((float(mem_res) / float(docker_memlimit)) * 100, 1)
)
else:
mem_pct = stats[9]
usages[stats[0]] = { with open('/proc/meminfo', 'r') as f:
"cpu": stats[8], mem_info = f.readlines()
total_mem = int(mem_info[0].split()[1])
mem_res = int(stat_info[23]) / 1024
if docker_memlimit > 0:
mem_pct = round((float(mem_res) / float(docker_memlimit)) * 100, 1)
else:
mem_pct = round((float(mem_res) / float(total_mem)) * 100, 1)
cpu_total_time = sum(int(x) for x in stat_info[13:17])
cpu_usage = round(cpu_total_time / os.sysconf("SC_CLK_TCK"), 1)
usages[pid] = {
"cpu": cpu_usage,
"mem": mem_pct, "mem": mem_pct,
} }
except: except: