1import os
  2import resource
  3import sys
  4import threading
  5import time
  6from functools import lru_cache
  7
  8
  9@lru_cache(maxsize=1)
 10def _get_cgroup_dir() -> str:
 11    """Resolve the process's own cgroup directory for cgroup v2.
 12
 13    Cached because the cgroup path is static for the lifetime of the process.
 14    """
 15    cgroup_dir = "/sys/fs/cgroup"
 16    try:
 17        with open("/proc/self/cgroup") as f:
 18            for line in f:
 19                # cgroup v2 entries have the form "0::<path>"
 20                parts = line.strip().split(":", 2)
 21                if (
 22                    len(parts) == 3
 23                    and parts[0] == "0"
 24                    and parts[1] == ""
 25                    and parts[2] != "/"
 26                ):
 27                    cgroup_dir = f"/sys/fs/cgroup{parts[2]}"
 28                    break
 29    except FileNotFoundError, IndexError, OSError:
 30        pass
 31    return cgroup_dir
 32
 33
 34def get_rss_bytes() -> int:
 35    """Get the current process's RSS in bytes.
 36
 37    Uses /proc/self/statm on Linux for current RSS,
 38    falls back to ru_maxrss (peak RSS) on macOS.
 39    """
 40    if sys.platform == "linux":
 41        try:
 42            with open("/proc/self/statm") as f:
 43                resident_pages = int(f.read().split()[1])
 44            return resident_pages * os.sysconf("SC_PAGE_SIZE")
 45        except OSError, ValueError, IndexError:
 46            pass
 47
 48    # macOS / fallback: ru_maxrss is peak, not current
 49    rusage = resource.getrusage(resource.RUSAGE_SELF)
 50    if sys.platform == "darwin":
 51        return rusage.ru_maxrss  # bytes on macOS
 52    return rusage.ru_maxrss * 1024  # KB on Linux (fallback)
 53
 54
 55def get_cpu_count() -> int:
 56    """Get the number of CPUs available, respecting cgroup limits in containers.
 57
 58    os.process_cpu_count() only checks sched_getaffinity, not cgroup CPU quotas,
 59    so containers often see the host's full CPU count. This reads the cgroup v2
 60    quota file to detect the actual limit.
 61
 62    Can be removed when minimum Python version is 3.14+ (cpython#120078).
 63    """
 64    cpu_count = os.process_cpu_count() or 1
 65
 66    cgroup_dir = _get_cgroup_dir()
 67
 68    # Check cgroup v2 CPU quota (Docker, Kubernetes, Railway, etc.)
 69    try:
 70        with open(f"{cgroup_dir}/cpu.max") as f:
 71            parts = f.read().strip().split()
 72            if len(parts) >= 2 and parts[0] != "max":
 73                quota = int(parts[0])
 74                period = int(parts[1])
 75                cgroup_cpus = max(1, -(-quota // period))  # ceiling division
 76                cpu_count = min(cpu_count, cgroup_cpus)
 77    except FileNotFoundError, ValueError, OSError:
 78        pass
 79
 80    # Check cgroup v1 CPU quota (Heroku Cedar, older Docker, etc.)
 81    try:
 82        with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") as f:
 83            quota = int(f.read().strip())
 84        if quota > 0:
 85            with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us") as f:
 86                period = int(f.read().strip())
 87            cgroup_cpus = max(1, -(-quota // period))
 88            cpu_count = min(cpu_count, cgroup_cpus)
 89    except FileNotFoundError, ValueError, OSError:
 90        pass
 91
 92    return cpu_count
 93
 94
 95def get_memory_usage() -> tuple[int, int | None]:
 96    """Get memory usage and optional limit in bytes, container-aware.
 97
 98    Returns (usage_bytes, limit_bytes) for containers with cgroup limits,
 99    or (usage_bytes, None) for bare-metal/dev where only process RSS is available.
100
101    Tries cgroup v2 first, then v1, then falls back to process RSS via getrusage.
102    """
103    cgroup_dir = _get_cgroup_dir()
104
105    # cgroup v2
106    try:
107        with open(f"{cgroup_dir}/memory.current") as f:
108            usage = int(f.read().strip())
109        with open(f"{cgroup_dir}/memory.max") as f:
110            content = f.read().strip()
111            limit = None if content == "max" else int(content)
112        return (usage, limit)
113    except FileNotFoundError, ValueError, OSError:
114        pass
115
116    # cgroup v1
117    try:
118        with open("/sys/fs/cgroup/memory/memory.usage_in_bytes") as f:
119            usage = int(f.read().strip())
120        with open("/sys/fs/cgroup/memory/memory.limit_in_bytes") as f:
121            limit = int(f.read().strip())
122            # cgroup v1 uses a very large number to mean "unlimited"
123            if limit >= 2**62:
124                limit = None
125        return (usage, limit)
126    except FileNotFoundError, ValueError, OSError:
127        pass
128
129    # Fallback: process RSS (works on macOS and Linux)
130    return (get_rss_bytes(), None)
131
132
133_cpu_lock = threading.Lock()
134_last_cpu_time: float = 0.0
135_last_wall_time: float = 0.0
136
137
138def get_process_cpu_percent() -> int | None:
139    """Get this process's CPU usage as a percentage of wall-clock time.
140
141    Uses resource.getrusage() which reports actual CPU time consumed,
142    accurate in both containers and on bare metal.
143
144    Returns None on the first call (no baseline to compare against).
145    Subsequent calls return the CPU percent since the previous call.
146    """
147    global _last_cpu_time, _last_wall_time
148
149    usage = resource.getrusage(resource.RUSAGE_SELF)
150    cpu_time = usage.ru_utime + usage.ru_stime
151    wall_time = time.monotonic()
152
153    with _cpu_lock:
154        if _last_wall_time == 0.0:
155            _last_cpu_time = cpu_time
156            _last_wall_time = wall_time
157            return None
158
159        wall_delta = wall_time - _last_wall_time
160        cpu_delta = cpu_time - _last_cpu_time
161
162        _last_cpu_time = cpu_time
163        _last_wall_time = wall_time
164
165    if wall_delta <= 0:
166        return 0
167
168    return min(round(cpu_delta / wall_delta * 100), 100)