#!/usr/bin/env python3
"""
Scan the cgroup tree for CPU (CFS bandwidth) throttling.

Pass 1: every cgroup that has *ever* been throttled (throttled_time > 0).
Pass 2: after an interval, every cgroup that throttled *during* the window
        (delta throttled_time > 0).

Targets cgroups v1 (throttled_time, nanoseconds) but also understands the
v2 field name (throttled_usec) so it degrades gracefully on unified hierarchies.
"""

import argparse
import os
import sys
import time

DEFAULT_ROOT = "/sys/fs/cgroup"


def parse_cpu_stat(path):
    """Parse a cpu.stat file into a dict of {key: int}, or None if unreadable."""
    stats = {}
    try:
        with open(path) as f:
            for line in f:
                parts = line.split()
                if len(parts) == 2:
                    try:
                        stats[parts[0]] = int(parts[1])
                    except ValueError:
                        pass
    except OSError:
        return None
    return stats or None


def throttled_ns(stats):
    """Return throttled time in nanoseconds regardless of v1/v2 field name."""
    if "throttled_time" in stats:            # cgroup v1: nanoseconds
        return stats["throttled_time"]
    if "throttled_usec" in stats:            # cgroup v2: microseconds
        return stats["throttled_usec"] * 1000
    return 0


def scan(root):
    """
    Walk the cgroup tree, returning {cgroup_path: stats_dict} for every
    directory that exposes a cpu.stat file.
    """
    results = {}
    for dirpath, _dirnames, filenames in os.walk(root):
        if "cpu.stat" in filenames:
            stats = parse_cpu_stat(os.path.join(dirpath, "cpu.stat"))
            if stats is not None:
                results[dirpath] = stats
    return results


def human_ns(ns):
    if ns >= 1_000_000_000:
        return f"{ns / 1_000_000_000:.2f}s"
    if ns >= 1_000_000:
        return f"{ns / 1_000_000:.1f}ms"
    if ns >= 1_000:
        return f"{ns / 1_000:.1f}us"
    return f"{ns}ns"


def display_name(path, root):
    """Strip the mount prefix so output shows the logical cgroup path."""
    rel = os.path.relpath(path, root)
    return "/" if rel == "." else "/" + rel


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--root", default=DEFAULT_ROOT,
                    help=f"cgroup mount root to scan (default: {DEFAULT_ROOT})")
    ap.add_argument("--interval", type=float, default=5.0,
                    help="seconds to wait between the two passes (default: 5)")
    args = ap.parse_args()

    if not os.path.isdir(args.root):
        sys.exit(f"error: {args.root} is not a directory")

    # ---- Pass 1: lifetime throttling ------------------------------------
    first = scan(args.root)
    ever = sorted(
        ((p, s) for p, s in first.items() if throttled_ns(s) > 0),
        key=lambda kv: throttled_ns(kv[1]),
        reverse=True,
    )

    print(f"=== Cgroups that have EVER throttled ({len(ever)}) ===")
    if not ever:
        print("  (none)")
    else:
        for path, stats in ever:
            t = throttled_ns(stats)
            n = stats.get("nr_throttled", 0)
            print(f"  {human_ns(t):>9}  nr_throttled={n:<8} {display_name(path, args.root)}")

    # ---- Pass 2: throttling during the interval -------------------------
    print(f"\nSleeping {args.interval:g}s...\n")
    time.sleep(args.interval)
    second = scan(args.root)

    deltas = []
    for path, stats in second.items():
        base = throttled_ns(first[path]) if path in first else 0
        d_ns = throttled_ns(stats) - base
        if d_ns > 0:
            base_n = first[path].get("nr_throttled", 0) if path in first else 0
            d_n = stats.get("nr_throttled", 0) - base_n
            deltas.append((path, d_ns, d_n))

    deltas.sort(key=lambda x: x[1], reverse=True)

    print(f"=== Cgroups that throttled in the last {args.interval:g}s ({len(deltas)}) ===")
    if not deltas:
        print("  (none)")
    else:
        for path, d_ns, d_n in deltas:
            print(f"  +{human_ns(d_ns):>9}  +{d_n:<6} throttled  {display_name(path, args.root)}")


if __name__ == "__main__":
    main()