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

Code before:

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)


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


class ProcessExt(psutil.Process):
    """Extends psutil.Process with better thread support."""

    def threads_cpu_times(self) -> Dict[str, CpuTimes]:
        """Return a dict mapping thread name to CPU usage (in seconds) since start."""
        cpu_times = {}
        for pthread in super().threads():
            try:
                name = psutil.Process(pthread.id).name()
                cpu_times[name] = CpuTimes(pthread.user_time, pthread.system_time)
            except psutil.NoSuchProcess as exc:
                logging.warning("Thread %d no longer exists", pthread.id, exc_info=exc)
                continue
        return cpu_times

    def threads_cpu_percent(
        self,
        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 = {}

        time_before = time.time()
        cpu_times_before = self.threads_cpu_times()

        time.sleep(interval)

        cpu_times_after = self.threads_cpu_times()
        time_after = time.time()

        time_diff = time_after - time_before
        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]

Review comment (suggestion):
maybe for these "nice if it works, don't care if it doesn't" kinda things, we can just use bare excepts? And maybe even just put the entire body of the function into one. That'd also save us the "except-pass" blocks at the call sites.

Code after:

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)


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


def get_cpu_times(process: psutil.Process) -> Dict[str, CpuTimes]:
    """Return a dict mapping thread name to CPU usage (in seconds) since start."""
    # We're consciously ignoring whatever erorr is returned by psutil and returning
    # empty {} as result in case of any error retrieving the process threads
    # information
    # pylint: disable=locally-disabled, broad-exception-caught

    threads = []
    try:
        threads = process.threads()
    except Exception as exc:
        logging.warning("Process %d does not exist", process.pid, exc_info=exc)
        return {}

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

    return cpu_times


def get_cpu_utilization(
    process: psutil.Process,
    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(process)
    time.sleep(interval)
    cpu_times_after = get_cpu_times(process)