cgroups-top

#!/usr/bin/env python3

import argparse
import io
import json
import os
import re
import shutil
import subprocess
import sys
import time
from pathlib import Path


BOLD = "\033[1m"
BOLD_RESET = "\033[22m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
RED = "\033[31m"
CYAN = "\033[36m"
RESET = "\033[0m"


def human_size(bytes_val):
    """Convert bytes to human-readable format."""
    for unit in ['B', 'KiB', 'MiB', 'GiB', 'TiB']:
        if bytes_val < 1024.0:
            if bytes_val >= 100:
                return f"{bytes_val:.0f}{unit}"
            return f"{bytes_val:.1f}{unit}"
        bytes_val /= 1024.0
    if bytes_val >= 100:
        return f"{bytes_val:.0f}PiB"
    return f"{bytes_val:.1f}PiB"


def bold(text, enabled=False):
    if enabled:
        return f"{BOLD}{text}{BOLD_RESET}"
    return text


def color(text, code, enabled=False):
    if enabled:
        return f"{code}{text}{RESET}"
    return text


def build_watch_title(data):
    """Build watch title with total memory, shmem, and swap across real cgroups."""
    real_rows = [row for row in data if row.get('service') not in ('[free]', '[tmpfs]')]
    total_mem = sum(row.get('memory_bytes', 0) or 0 for row in real_rows)
    total_shmem = sum(row.get('shmem_bytes', 0) or 0 for row in real_rows)
    total_swap = sum(row.get('swap_bytes', 0) or 0 for row in real_rows)
    return f"cgroups-top {human_size(total_mem)} {human_size(total_shmem)} {human_size(total_swap)}"


def get_proc_comm(pid):
    """Get command name for a PID."""
    try:
        return Path(f"/proc/{pid}/comm").read_text().strip()
    except (FileNotFoundError, PermissionError, ProcessLookupError):
        return None


def get_proc_cmdline(pid):
    """Get full command line for a PID."""
    try:
        cmdline = Path(f"/proc/{pid}/cmdline").read_text()
        if cmdline:
            return ' '.join(cmdline.split('\x00')).strip()
        return get_proc_comm(pid)
    except (FileNotFoundError, PermissionError, ProcessLookupError):
        return None


def get_proc_memory(pid):
    """Get memory usage for a PID from /proc/pid/status."""
    try:
        status = Path(f"/proc/{pid}/status").read_text()
        rss_kb = 0
        file_kb = 0
        anon_kb = 0
        swap_kb = 0
        shmem_kb = 0

        for line in status.split('\n'):
            if line.startswith('VmRSS:'):
                rss_kb = int(line.split()[1])
            elif line.startswith('RssFile:'):
                file_kb = int(line.split()[1])
            elif line.startswith('RssAnon:'):
                anon_kb = int(line.split()[1])
            elif line.startswith('VmSwap:'):
                swap_kb = int(line.split()[1])
            elif line.startswith('RssShmem:'):
                shmem_kb = int(line.split()[1])

        return {
            'rss_bytes': rss_kb * 1024,
            'file_bytes': file_kb * 1024,
            'anon_bytes': anon_kb * 1024,
            'swap_bytes': swap_kb * 1024,
            'shmem_bytes': shmem_kb * 1024
        }
    except (FileNotFoundError, PermissionError, ProcessLookupError, ValueError):
        return {'rss_bytes': 0, 'file_bytes': 0, 'anon_bytes': 0, 'swap_bytes': 0, 'shmem_bytes': 0}


def get_meminfo():
    """Parse /proc/meminfo into a dict of name -> kB."""
    info = {}
    for line in Path("/proc/meminfo").read_text().split('\n'):
        parts = line.split()
        if len(parts) >= 2:
            key = parts[0].rstrip(':')
            try:
                info[key] = int(parts[1])
            except ValueError:
                pass
    return info


def get_tmpfs_usage():
    """Get total bytes used across all tmpfs mounts."""
    total_used = 0
    seen = set()
    try:
        for line in Path("/proc/mounts").read_text().split('\n'):
            parts = line.split()
            if len(parts) >= 3 and parts[2] == 'tmpfs':
                mountpoint = parts[1]
                if mountpoint in seen:
                    continue
                seen.add(mountpoint)
                try:
                    stat = os.statvfs(mountpoint)
                    total_used += (stat.f_blocks - stat.f_bfree) * stat.f_frsize
                except (PermissionError, OSError):
                    pass
    except (FileNotFoundError, PermissionError):
        pass
    return total_used


def get_synthetic_entries():
    """Return synthetic entries for free memory+swap and tmpfs."""
    meminfo = get_meminfo()
    total_mem = meminfo.get('MemTotal', 0) * 1024
    free_mem = meminfo.get('MemFree', 0) * 1024
    available_mem = meminfo.get('MemAvailable', 0) * 1024
    free_swap = meminfo.get('SwapFree', 0) * 1024
    reclaimable_cache = max(0, available_mem - free_mem)
    tmpfs_used = get_tmpfs_usage()

    return [
        {
            'memory_bytes': free_mem,
            'swap_bytes': free_swap,
            'shmem_bytes': reclaimable_cache,
            'available_bytes': available_mem,
            'total_bytes': total_mem,
            'root_pid': None,
            'service': '[free]',
            'executables': [],
            'processes': [],
            'cgroup_path': '',
            'pids': [],
        },
        {
            'memory_bytes': tmpfs_used,
            'swap_bytes': None,
            'shmem_bytes': None,
            'root_pid': None,
            'service': '[tmpfs]',
            'executables': [],
            'processes': [],
            'cgroup_path': '',
            'pids': [],
        },
    ]


def collect_cgroup_data():
    """Collect cgroup memory and process data."""
    results = []

    for cgroup_dir in Path("/sys/fs/cgroup").rglob("*"):
        if not cgroup_dir.is_dir():
            continue

        memory_file = cgroup_dir / "memory.current"
        swap_file = cgroup_dir / "memory.swap.current"
        stat_file = cgroup_dir / "memory.stat"
        procs_file = cgroup_dir / "cgroup.procs"

        if not (memory_file.exists() and procs_file.exists()):
            continue

        try:
            # Extract shmem and inactive_file from memory.stat
            shmem_bytes = 0
            inactive_file_bytes = 0
            if stat_file.exists():
                for line in stat_file.read_text().split('\n'):
                    if line.startswith('shmem '):
                        shmem_bytes = int(line.split()[1])
                    elif line.startswith('inactive_file '):
                        inactive_file_bytes = int(line.split()[1])

            memory_bytes = int(memory_file.read_text().strip())
            # Don't include inactive file cache in the cgroup total
            memory_bytes -= inactive_file_bytes
            # Note: shmem_bytes is not subtracted here; it's counted both in this cgroup's
            # total and in the synthetic [tmpfs] entry, which is intentional so per-cgroup
            # shmem usage remains visible.

            swap_bytes = int(swap_file.read_text().strip()) if swap_file.exists() else 0

            pids = [int(p) for p in procs_file.read_text().strip().split('\n') if p]
            if not pids:
                continue

            root_pid = pids[0]

            # Get distinct executable names
            exes = set()
            for pid in pids:
                comm = get_proc_comm(pid)
                if comm:
                    exes.add(comm)

            # Collect per-process memory totals for this cgroup
            processes = []
            for pid in pids:
                cmdline = get_proc_cmdline(pid) or get_proc_comm(pid) or ''
                mem = get_proc_memory(pid)
                # Include RssFile (file_bytes) as it represents the "active" file cache for the process
                # rss_bytes already includes anon_bytes + file_bytes + shmem_bytes
                total = mem['rss_bytes'] + mem['swap_bytes']
                processes.append({
                    'pid': pid,
                    'cmdline': cmdline,
                    **mem,
                    'total_bytes': total
                })

            # Sort processes by total memory usage descending
            processes.sort(key=lambda x: x['total_bytes'], reverse=True)

            service_name = cgroup_dir.name[:40]

            results.append({
                'memory_bytes': memory_bytes,
                'swap_bytes': swap_bytes,
                'shmem_bytes': shmem_bytes,
                'root_pid': root_pid,
                'service': service_name,
                'executables': sorted(exes),
                'processes': processes,
                'cgroup_path': str(cgroup_dir),
                'pids': pids
            })

        except (FileNotFoundError, ValueError, PermissionError):
            continue

    # Sort by total usage (memory + swap).
    # memory_bytes already includes shmem.
    results.sort(key=lambda x: x['memory_bytes'] + (x['swap_bytes'] or 0), reverse=True)
    return results


def get_terminal_width():
    """Get terminal width if output is interactive."""
    is_tty = sys.stdout.isatty()
    if is_tty:
        terminal_width = shutil.get_terminal_size().columns
    else:
        terminal_width = None
    return is_tty, terminal_width


def print_cgroups_table(data, watch_mode=False, file=sys.stdout):
    """Print cgroup data in aligned table format."""
    if not data:
        return

    is_tty, terminal_width = get_terminal_width()

    # Calculate max widths
    max_pid_width = max(len(str(row['root_pid'])) if row['root_pid'] is not None else 0 for row in data)

    # Print header
    header = f"{'Total':>9} {'Memory':>9} {'Shmem':>9} {'Swap':>9} {'PID':>{max_pid_width}} {'Service':40} {'Executables (Mem)'}"
    print(bold(header, watch_mode), file=file)

    # Calculate fixed width for clipping. Executables column will include per-executable memory totals.
    fixed_width = 9 + 1 + 9 + 1 + 9 + 1 + 9 + 1 + max_pid_width + 1 + 40 + 1
    if is_tty:
        exes_width = max(20, terminal_width - fixed_width)
    else:
        exes_width = None

    for row in data:
        # memory_bytes already includes shmem
        total_bytes = row['memory_bytes'] + (row['swap_bytes'] or 0)
        total_str = human_size(total_bytes)
        memory_str = human_size(row['memory_bytes']) if row['shmem_bytes'] is not None else ''
        shmem_str = human_size(row['shmem_bytes']) if row['shmem_bytes'] is not None else ''
        swap_str = human_size(row['swap_bytes']) if row['swap_bytes'] is not None else ''
        pid_str = str(row['root_pid']) if row['root_pid'] is not None else ''
        service = row['service']

        # Aggregate memory per executable name (prefer /proc/<pid>/comm).
        # Count anon+swap for each process, but include shmem and mapped files only once per executable
        exe_anon_swap = {}
        exe_shmem_once = {}
        exe_file_once = {}
        for p in (row.get('processes') or []):
            try:
                comm = get_proc_comm(p['pid']) or ''
            except Exception:
                comm = ''
            if not comm:
                # fallback to first token of cmdline
                cmd = p.get('cmdline') or ''
                comm = cmd.split(' ')[0] if cmd else ''

            anon = p.get('anon_bytes', 0)
            swap = p.get('swap_bytes', 0)
            shmem = p.get('shmem_bytes', 0)
            file_mem = p.get('file_bytes', 0)

            exe_anon_swap[comm] = exe_anon_swap.get(comm, 0) + anon + swap
            # include shared memory and mapped files only once per executable: track the max seen
            exe_shmem_once[comm] = max(exe_shmem_once.get(comm, 0), shmem)
            exe_file_once[comm] = max(exe_file_once.get(comm, 0), file_mem)

        # Build sorted list by memory desc (anon+swap summed per-exe + shmem/file counted once per-exe)
        exe_items = []
        for name, anon_swap_sum in exe_anon_swap.items():
            shmem_once = exe_shmem_once.get(name, 0)
            file_once = exe_file_once.get(name, 0)
            total = anon_swap_sum + shmem_once + file_once
            exe_items.append((name, total))

        exe_items.sort(key=lambda x: x[1], reverse=True)
        exes_with_mem = [f"{human_size(size)} {name}" for name, size in exe_items if name]
        exes_str = ', '.join(exes_with_mem)

        # Clip executables string to available width
        if exes_width and len(exes_str) > exes_width:
            exes_str = exes_str[:exes_width-3] + '...'

        memory_col = bold(f"{memory_str:>9}", watch_mode)
        row_text = f"{total_str:>9} {memory_col} {shmem_str:>9} {swap_str:>9} {pid_str:>{max_pid_width}} {service:40} {exes_str}"
        if is_tty and service == '[free]':
            available_bytes = row.get('available_bytes')
            total_mem = row.get('total_bytes') or 0
            available_ratio = (available_bytes / total_mem) if available_bytes is not None and total_mem else None
            if available_ratio is not None and available_ratio < 0.10:
                row_text = color(row_text, RED, True)
            elif available_ratio is not None and available_ratio < 0.20:
                row_text = color(row_text, YELLOW, True)
            else:
                row_text = color(row_text, GREEN, True)
        elif is_tty and service == '[tmpfs]':
            row_text = color(row_text, CYAN, True)
        print(row_text, file=file)


def print_processes_table(data, title=None, watch_mode=False, file=sys.stdout):
    """Print process data in aligned table format."""
    if not data:
        return

    is_tty, terminal_width = get_terminal_width()

    if title:
        print(title, file=file)

    # Calculate max widths
    max_pid_width = max(len(str(row['pid'])) for row in data)

    # Print header
    header = f"{'Total':>9} {'RSS':>9} {'Shmem':>9} {'Swap':>9} {'PID':>{max_pid_width}} {'COMMAND'}"
    print(bold(header, watch_mode), file=file)

    # Calculate fixed width for clipping
    fixed_width = 9 + 1 + 9 + 1 + 9 + 1 + 9 + 1 + max_pid_width + 1
    if is_tty:
        cmdline_width = max(20, terminal_width - fixed_width)
    else:
        cmdline_width = None

    for row in data:
        # Include RssFile (file_bytes) as it represents the "active" file cache for the process
        # rss_bytes already includes anon_bytes + file_bytes + shmem_bytes
        total_bytes = row['rss_bytes'] + row['swap_bytes']
        total_str = human_size(total_bytes)
        rss_str = human_size(row['rss_bytes'])
        shmem_str = human_size(row['shmem_bytes'])
        swap_str = human_size(row['swap_bytes'])
        pid_str = str(row['pid'])
        cmdline = row['cmdline']

        if cmdline_width and len(cmdline) > cmdline_width:
            cmdline = cmdline[:cmdline_width-3] + '...'

        memory_col = bold(f"{rss_str:>9}", watch_mode)
        print(f"{total_str:>9} {memory_col} {shmem_str:>9} {swap_str:>9} {pid_str:>{max_pid_width}} {cmdline}", file=file)


def print_jsonl(data, include_processes=False, file=sys.stdout):
    """Print data in JSON Lines format."""
    for row in data:
        output = {k: v for k, v in row.items() if k not in ['cgroup_path', 'pids']}
        # Add a simple exe field (first executable name) to make grouping by executable easier with jq
        exe = None
        exes = row.get('executables') or []
        if exes:
            exe = exes[0]
        output['exe'] = exe
        if include_processes:
            output['processes'] = []
            for pid in row['pids']:
                mem = get_proc_memory(pid)
                output['processes'].append({
                    'pid': pid,
                    'comm': get_proc_comm(pid),
                    **mem
                })
        print(json.dumps(output), file=file)


def show_cgroup_processes(cgroup_name, limit=10, watch_mode=False, file=sys.stdout):
    """Show individual process memory usage for a specific cgroup."""
    for cgroup_dir in Path("/sys/fs/cgroup").rglob("*"):
        if not cgroup_dir.is_dir():
            continue

        if cgroup_dir.name != cgroup_name:
            continue

        procs_file = cgroup_dir / "cgroup.procs"
        if not procs_file.exists():
            continue

        try:
            pids = [int(p) for p in procs_file.read_text().strip().split('\n') if p]
            if not pids:
                continue

            # Collect process data
            processes = []
            for pid in pids:
                cmdline = get_proc_cmdline(pid)
                mem = get_proc_memory(pid)
                if cmdline:
                    processes.append({
                        'pid': pid,
                        'cmdline': cmdline,
                        **mem
                    })

            # Sort by total usage (RSS + Swap).
            processes.sort(key=lambda x: x['rss_bytes'] + x['swap_bytes'], reverse=True)

            # Apply limit
            if limit > 0:
                processes = processes[:limit]

            # Use print_processes_table to display processes
            print_processes_table(processes, title=f"Cgroup: {cgroup_dir}", watch_mode=watch_mode, file=file)

            return True

        except (FileNotFoundError, ValueError, PermissionError):
            continue

    print(f"Cgroup '{cgroup_name}' not found", file=sys.stderr)
    return False


def run_once(args, limit=None, file=sys.stdout):
    if limit is None:
        limit = args.limit
    if args.processes:
        if not show_cgroup_processes(args.processes, limit, watch_mode=args.watch, file=file) and not args.watch:
            sys.exit(1)
        return None

    data = collect_cgroup_data()

    if not args.json and not args.processes:
        data = data + get_synthetic_entries()
        data.sort(key=lambda x: x['memory_bytes'] + (x['swap_bytes'] or 0), reverse=True)

    # Apply limit for table output
    if not args.json and limit > 0:
        data = data[:limit]

    if args.json:
        print_jsonl(data, include_processes=args.include_processes, file=file)
    else:
        print_cgroups_table(data, watch_mode=args.watch, file=file)
    return data


def main():
    parser = argparse.ArgumentParser(description='Display cgroup memory usage')
    parser.add_argument('--json', action='store_true', help='Output in JSON Lines format')
    parser.add_argument('--processes', metavar='CGROUP', help='Show individual process memory for a specific cgroup')
    parser.add_argument('--include-processes', action='store_true', help='Include process details in JSON output')
    parser.add_argument('-n', '--limit', type=int, default=10, help='Number of top entries to show (default: 10, use 0 for all)')
    parser.add_argument('--watch', action='store_true', help='Refresh output periodically')
    parser.add_argument('--interval', type=int, default=2, help='Refresh interval in seconds (default: 2)')
    args = parser.parse_args()

    try:
        if args.watch:
            # Hide cursor while watching
            print("\033[?25l", end="", flush=True)
            try:
                while True:
                    buf = io.StringIO()
                    data = run_once(args, limit=0, file=buf)
                    output = buf.getvalue().rstrip('\n')
                    terminal_height = shutil.get_terminal_size().lines
                    lines = output.split('\n')
                    if len(lines) > terminal_height:
                        lines = lines[:terminal_height]
                    output = '\n'.join(lines)
                    if data and not args.json and not args.processes:
                        print(f"\033]2;{build_watch_title(data)}\007", end="", flush=True)
                    else:
                        print("\033]2;cgroups-top (watch)\007", end="", flush=True)
                    # Clear screen and move to top (ANSI escape codes)
                    # \033[H: Move cursor to home (top-left)
                    # \033[J: Clear from cursor to end of screen
                    print(f"\033[H\033[J{output}", end="", flush=True)
                    time.sleep(args.interval)
            finally:
                # Restore cursor
                print("\033[?25h", flush=True)
        else:
            run_once(args)
    except KeyboardInterrupt:
        pass
    except BrokenPipeError:
        print("Broken pipe", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()