Function to get memory used by processes rooted at pid

kitty/child.py

@@ -26,6 +26,8 @@ if is_macos:
     from kitty.fast_data_types import cmdline_of_process as cmdline_
     from kitty.fast_data_types import cwd_of_process as _cwd
     from kitty.fast_data_types import environ_of_process as _environ_of_process
+    from kitty.fast_data_types import memory_of_process as _memory_of_process
+    from kitty.fast_data_types import ppid_of_process as _ppid_of_process
     from kitty.fast_data_types import process_group_map as _process_group_map
 
     def cwd_of_process(pid: int) -> str:
@@ -44,6 +46,30 @@ if is_macos:
 
     def cmdline_of_pid(pid: int) -> list[str]:
         return cmdline_(pid)
+
+    def _get_descendants_of_macos(pid: int) -> set[int]:
+        children_map: DefaultDict[int, list[int]] = defaultdict(list)
+        for p in fast_data_types.get_all_processes():
+            with suppress(Exception):
+                children_map[_ppid_of_process(p)].append(p)
+        result: set[int] = set()
+        stack = list(children_map.get(pid, []))
+        while stack:
+            child = stack.pop()
+            if child not in result:
+                result.add(child)
+                stack.extend(children_map.get(child, []))
+        return result
+
+    def memory_used_by_process_tree_rooted_at(pid: int, check_if_cgroup_root: bool = False) -> int:
+        with suppress(Exception):
+            pids = _get_descendants_of_macos(pid)
+            total = _memory_of_process(pid)  # raises if pid doesn't exist
+            for p in pids:
+                with suppress(Exception):
+                    total += _memory_of_process(p)
+            return total
+        return -1
 else:
 
     def cmdline_of_pid(pid: int) -> list[str]:
@@ -94,6 +120,63 @@ else:
     def abspath_of_exe(pid: int) -> str:
         return os.path.realpath(f'/proc/{pid}/exe', strict=True)
 
+    def _get_descendants_of(pid: int) -> set[int]:
+        result: set[int] = set()
+        stack = [pid]
+        while stack:
+            current = stack.pop()
+            with suppress(OSError):
+                with open(f'/proc/{current}/task/{current}/children') as f:
+                    for child_str in f.read().split():
+                        child = int(child_str)
+                        if child not in result:
+                            result.add(child)
+                            stack.append(child)
+        return result
+
+    def _memory_from_smaps_rollup(pid: int) -> int:
+        with open(f'/proc/{pid}/smaps_rollup') as f:
+            for line in f:
+                if line.startswith('Pss:'):
+                    return int(line.split()[1]) * 1024
+        return 0
+
+    def memory_used_by_process_tree_rooted_at(pid: int, check_if_cgroup_root: bool = False) -> int:
+        with suppress(Exception):
+            with open(f'/proc/{pid}/cgroup') as f:
+                cgroup_line = f.readline().strip()
+            cgroup_path = cgroup_line.split(':')[2].lstrip('/')
+            cgroup_dir = os.path.join('/sys/fs/cgroup', cgroup_path)
+
+            use_cgroup = True
+            if check_if_cgroup_root:
+                with suppress(OSError):
+                    with open(os.path.join(cgroup_dir, 'cgroup.procs')) as f:
+                        cgroup_pids = {int(x) for x in f.read().split() if x}
+                    descendants = _get_descendants_of(pid)
+                    descendants.add(pid)
+                    use_cgroup = cgroup_pids <= descendants
+
+            if use_cgroup:
+                target_keys = {'anon', 'shmem', 'kernel', 'sock', 'zswap'}
+                mem_bytes = 0
+                with open(os.path.join(cgroup_dir, 'memory.stat')) as f:
+                    for line in f:
+                        parts = line.split()
+                        if parts[0] in target_keys:
+                            mem_bytes += int(parts[1])
+                return mem_bytes
+
+            # cgroup contains processes outside our tree; sum PSS per process
+            descendants = _get_descendants_of(pid)
+            descendants.add(pid)
+            mem_bytes = 0
+            for p in descendants:
+                with suppress(OSError):
+                    mem_bytes += _memory_from_smaps_rollup(p)
+            return mem_bytes
+        return -1
+
 
 @run_once
 def checked_terminfo_dir() -> str | None:
@@ -619,3 +702,8 @@ class Child:
                 termios.tcsetattr(self.child_fd, when, self.initial_termios_state)
             except OSError:
                 pass
+
+    def get_memory_used_by_child(self, check_if_cgroup_root: bool = False) -> int:
+        if self.pid is None:
+            return -1
+        return memory_used_by_process_tree_rooted_at(self.pid, check_if_cgroup_root)

kitty/macos_process_info.c

@@ -286,12 +286,36 @@ error:
 }
 
 
+static PyObject*
+memory_of_process(PyObject *self UNUSED, PyObject *pid_) {
+    if (!PyLong_Check(pid_)) { PyErr_SetString(PyExc_TypeError, "pid must be an int"); return NULL; }
+    pid_t pid = (pid_t)PyLong_AsLong(pid_);
+    if (pid < 0) { PyErr_SetString(PyExc_TypeError, "pid cannot be negative"); return NULL; }
+    struct proc_taskinfo ti;
+    int ret = proc_pidinfo(pid, PROC_PIDTASKINFO, 0, &ti, sizeof(ti));
+    if (ret <= 0) { PyErr_SetFromErrno(PyExc_OSError); return NULL; }
+    return PyLong_FromUnsignedLongLong(ti.pti_resident_size);
+}
+
+static PyObject*
+ppid_of_process(PyObject *self UNUSED, PyObject *pid_) {
+    if (!PyLong_Check(pid_)) { PyErr_SetString(PyExc_TypeError, "pid must be an int"); return NULL; }
+    pid_t pid = (pid_t)PyLong_AsLong(pid_);
+    if (pid < 0) { PyErr_SetString(PyExc_TypeError, "pid cannot be negative"); return NULL; }
+    struct proc_bsdshortinfo si;
+    int ret = proc_pidinfo(pid, PROC_PIDT_SHORTBSDINFO, 0, &si, sizeof(si));
+    if (ret <= 0) { PyErr_SetFromErrno(PyExc_OSError); return NULL; }
+    return PyLong_FromUnsignedLong(si.pbsi_ppid);
+}
+
 static PyMethodDef module_methods[] = {
     {"cwd_of_process", (PyCFunction)cwd_of_process, METH_O, ""},
     {"abspath_of_process", (PyCFunction)abspath_of_process, METH_O, ""},
     {"cmdline_of_process", (PyCFunction)cmdline_of_process, METH_O, ""},
     {"environ_of_process", (PyCFunction)environ_of_process, METH_O, ""},
     {"get_all_processes", (PyCFunction)get_all_processes, METH_NOARGS, ""},
+    {"memory_of_process", (PyCFunction)memory_of_process, METH_O, ""},
+    {"ppid_of_process", (PyCFunction)ppid_of_process, METH_O, ""},
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };
 

kitty_tests/child.py

@@ -0,0 +1,100 @@
+#!/usr/bin/env python
+# License: GPL v3 Copyright: 2026, Kovid Goyal <kovid at kovidgoyal.net>
+
+import os
+import subprocess
+
+from kitty.child import memory_used_by_process_tree_rooted_at
+from kitty.constants import is_macos, kitty_exe
+
+from . import BaseTest
+
+
+class ChildMemoryTest(BaseTest):
+
+    def _spawn_allocating_child(self, alloc_bytes: int) -> subprocess.Popen:
+        p = subprocess.Popen(
+            [kitty_exe(), '+runpy', f'''\
+import sys, time
+buf = bytearray({alloc_bytes})
+for i in range(0, {alloc_bytes}, 4096):
+    buf[i] = 1
+sys.stdout.write("ready\\n")
+sys.stdout.flush()
+time.sleep(300)
+'''],
+            stdout=subprocess.PIPE,
+        )
+        line = p.stdout.readline().strip()
+        p.stdout.close()
+        if line != b'ready':
+            p.kill()
+            p.wait()
+            raise AssertionError(f'Unexpected output from allocating child: {line!r}')
+        return p
+
+    def _terminate(self, p: subprocess.Popen) -> None:
+        p.terminate()
+        p.wait()
+
+    def test_memory_returns_positive_for_live_process(self):
+        mem = memory_used_by_process_tree_rooted_at(os.getpid())
+        self.assertGreater(mem, 0)
+
+    def test_memory_returns_minus_one_for_nonexistent_pid(self):
+        self.ae(memory_used_by_process_tree_rooted_at(99999999), -1)
+
+    def test_memory_accounts_for_child_allocation(self):
+        # Verify that a child's known resident allocation shows up in the
+        # measurement.  check_if_cgroup_root=True falls back to per-process
+        # tree walk when pid is not the cgroup root (the common case when
+        # running under a shared session cgroup), so this exercises the tree
+        # walk path on Linux and the always-tree-walk path on macOS.
+        alloc = 20 * 1024 * 1024  # 20 MiB
+        child = self._spawn_allocating_child(alloc)
+        try:
+            mem = memory_used_by_process_tree_rooted_at(child.pid, check_if_cgroup_root=True)
+            self.assertGreater(
+                mem, alloc // 2,
+                f'Expected at least {alloc // 2} bytes for a {alloc}-byte allocation, got {mem}',
+            )
+        finally:
+            self._terminate(child)
+
+    def test_memory_of_parent_tree_includes_child(self):
+        # Measuring the current process tree must yield more than measuring
+        # the child alone, because the test runner itself occupies memory.
+        alloc = 20 * 1024 * 1024  # 20 MiB
+        child = self._spawn_allocating_child(alloc)
+        try:
+            mem_child = memory_used_by_process_tree_rooted_at(child.pid, check_if_cgroup_root=True)
+            mem_tree = memory_used_by_process_tree_rooted_at(os.getpid(), check_if_cgroup_root=True)
+            self.assertGreater(
+                mem_tree, mem_child,
+                'Parent tree memory should exceed child-only memory',
+            )
+        finally:
+            self._terminate(child)
+
+    def test_memory_cgroup_path_returns_positive(self):
+        # The fast cgroup path (check_if_cgroup_root=False, the default) must
+        # return a usable value on Linux.
+        if is_macos:
+            self.skipTest('cgroup not available on macOS')
+        mem = memory_used_by_process_tree_rooted_at(os.getpid())
+        self.assertGreater(mem, 0)
+
+    def test_memory_cgroup_and_tree_walk_both_positive(self):
+        # Both the cgroup path and the tree-walk path should give positive
+        # results for a process that is alive.
+        if is_macos:
+            self.skipTest('cgroup path not applicable on macOS')
+        alloc = 10 * 1024 * 1024  # 10 MiB
+        child = self._spawn_allocating_child(alloc)
+        try:
+            mem_cgroup = memory_used_by_process_tree_rooted_at(child.pid, check_if_cgroup_root=False)
+            mem_walk = memory_used_by_process_tree_rooted_at(child.pid, check_if_cgroup_root=True)
+            self.assertGreater(mem_cgroup, 0)
+            self.assertGreater(mem_walk, 0)
+        finally:
+            self._terminate(child)