Repository: firecracker-microvm/firecracker
File: tests/framework/utils.py (Python)

Code before:
        for thread in proc.threads():
            threads_map[psutil.Process(thread.id).name()].append(thread.id)
        return threads_map
    except psutil.NoSuchProcess:
        return {}


def get_cpu_affinity(pid: int) -> list:
    """Get CPU affinity for a thread."""
    return psutil.Process(pid).cpu_affinity()


def set_cpu_affinity(pid: int, cpulist: list) -> list:
    """Set CPU affinity for a thread."""
    real_cpulist = list(map(CpuMap, cpulist))
    return psutil.Process(pid).cpu_affinity(real_cpulist)


def get_thread_name(pid: int, tid: int) -> str:
    """Return thread name from pid and tid pair."""
    return Path("/proc", str(pid), "task", str(tid), "comm").read_text("utf-8").strip()


CpuTimes = namedtuple("CpuTimes", ["user", "system"])


def get_cpu_times(pid: int) -> Dict[str, CpuTimes]:
    """Return a dict mapping thread name to CPU usage (in seconds) since start."""
    cpu_times = {}
    for thread in psutil.Process(pid).threads():
        thread_name = get_thread_name(pid, thread.id)
        cpu_times[thread_name] = CpuTimes(thread.user_time, thread.system_time)
    return cpu_times


def get_cpu_utilization(
    pid: int,
    interval: int = 1,
    split_user_system: bool = False,
) -> Dict[str, float | CpuTimes]:
    """Return current process per thread CPU utilization over the interval (seconds)."""
    cpu_utilization = {}
    cpu_times_before = get_cpu_times(pid)
    time.sleep(interval)
    cpu_times_after = get_cpu_times(pid)
    threads = set(cpu_times_before.keys()) & set(cpu_times_after.keys())
    for thread_name in threads:
        before = cpu_times_before[thread_name]
        after = cpu_times_after[thread_name]
        user = (after.user - before.user) / interval * 100
        system = (after.system - before.system) / interval * 100

Review comment (suggestion):
I could have just returned the tid as a string, but I wanted to avoid having spurious metrics emitted from the cpu monitor thread in the block test.

Code after:
        for thread in proc.threads():
            threads_map[psutil.Process(thread.id).name()].append(thread.id)
        return threads_map
    except psutil.NoSuchProcess:
        return {}


def get_cpu_affinity(pid: int) -> list:
    """Get CPU affinity for a thread."""
    return psutil.Process(pid).cpu_affinity()


def set_cpu_affinity(pid: int, cpulist: list) -> list:
    """Set CPU affinity for a thread."""
    real_cpulist = list(map(CpuMap, cpulist))
    return psutil.Process(pid).cpu_affinity(real_cpulist)


def get_thread_name(pid: int, tid: int) -> str:
    """Return thread name from pid and tid pair."""
    try:
        return (
            Path("/proc", str(pid), "task", str(tid), "comm").read_text("utf-8").strip()
        )
    except FileNotFoundError as exc:
        raise psutil.NoSuchProcess(tid) from exc


CpuTimes = namedtuple("CpuTimes", ["user", "system"])


def get_cpu_times(pid: int) -> Dict[str, CpuTimes]:
    """Return a dict mapping thread name to CPU usage (in seconds) since start."""
    threads = []
    try:
        threads = psutil.Process(pid).threads()
    except psutil.NoSuchProcess as exc:
        logging.warning("Process %d does not exist", pid, exc_info=exc)
        return {}

    cpu_times = {}
    for thread in threads:
        try:
            thread_name = get_thread_name(pid, thread.id)
            cpu_times[thread_name] = CpuTimes(thread.user_time, thread.system_time)
        except psutil.NoSuchProcess as exc:
            logging.warning("Thread %d no longer exists", thread.id, exc_info=exc)
            continue

    return cpu_times