arch-image-builder/builder/lib/cgroup.py

97 lines
1.8 KiB
Python
Raw Normal View History

import os
import time
import signal
from logging import getLogger
log = getLogger(__name__)
class CGroup:
fs: str = "/sys/fs/cgroup"
name: str
@property
def path(self) -> str:
"""
Get full path of this cgroup
"""
return os.path.join(self.fs, self.name)
@property
def valid(self) -> bool:
"""
Can read or write to this cgroup
"""
return os.path.exists(self.path)
def create(self):
"""
Create this cgroup now
"""
if self.valid: return
os.mkdir(self.path)
def destroy(self):
"""
Destroy the cgroup
"""
if not self.valid: return
os.rmdir(self.path)
def add_pid(self, pid: int):
"""
Add a pid to track
"""
if not self.valid: return
procs = os.path.join(self.path, "cgroup.procs")
with open(procs, "w") as f:
f.write(f"{pid}\n")
def list_pid(self) -> list[int]:
"""
List all tracked children progress id
"""
ret: list[int] = []
if not self.valid: return ret
procs = os.path.join(self.path, "cgroup.procs")
with open(procs, "r") as f:
for line in f:
ret.append(int(line))
return ret
def kill_all(self, sig: int = signal.SIGTERM, timeout: int = 10, kill: int = 8):
"""
Kill all children process and wait them exit
"""
if not self.valid: return
pids = self.list_pid()
remain = 0
while True:
# send a signal
for pid in pids:
log.debug(f"killing {pid}")
try: os.kill(pid, sig)
except: pass
# waitpid to clean zombie
try: os.waitpid(-1, os.WNOHANG)
except: pass
# check all children was exited
pids = self.list_pid()
if len(pids) <= 0: break
# set to SIGKILL when reached kill time
if 0 < kill <= remain:
sig = signal.SIGKILL
# timeoutd, throw out
if remain >= timeout:
raise TimeoutError("killing pids timedout")
# wait...
time.sleep(1)
def __init__(self, name: str, fs: str = None):
if fs: self.fs = fs
self.name = name