Cgroup scripts

contrib/cgroup.py

@@ -4,9 +4,11 @@
 
 """List the paths of all descendants of a cgroup v2"""
 
+import argparse
 from contextlib import contextmanager
 import os
 import sys
+from collections import Counter
 
 from drgn import cast
 from drgn.helpers.common.type import enum_type_to_class
@@ -24,6 +26,11 @@
     exclude=("__MAX_BPF_ATTACH_TYPE",),
 )
 
+CgroupSubsysId = enum_type_to_class(
+    prog.type("enum cgroup_subsys_id"),
+    "CgroupSubsysId",
+    exclude=("CGROUP_SUBSYS_COUNT",),
+)
 
 @contextmanager
 def open_dir(*args, **kwds):
@@ -35,16 +42,6 @@ def open_dir(*args, **kwds):
         os.close(fd)
 
 
-def get_cgroup():
-    if len(sys.argv) == 1:
-        return prog["cgrp_dfl_root"].cgrp
-    task = find_task(prog, os.getpid())
-    with open_dir(sys.argv[1], os.O_RDONLY) as fd:
-        file_ = fget(task, fd)
-        kn = cast("struct kernfs_node *", file_.f_path.dentry.d_inode.i_private)
-        return cast("struct cgroup *", kn.priv)
-
-
 def print_cgroup_bpf_progs(cgrp):
     cgroup_printed = False
     for attach_type in BpfAttachType:
@@ -62,10 +59,74 @@ def print_cgroup_bpf_progs(cgrp):
             )
 
 
-css = get_cgroup().self.address_of_()
+def get_cgroup(path):
+    task = find_task(prog, os.getpid())
+    try:
+        with open_dir(path, os.O_RDONLY) as fd:
+            file_ = fget(task, fd)
+            kn = cast("struct kernfs_node *", file_.f_path.dentry.d_inode.i_private)
+            return cast("struct cgroup *", kn.priv)
+    except FileNotFoundError as e:
+        raise argparse.ArgumentTypeError(e)
+
+
+def cmd_tree(cgroup):
+    css = cgroup.self.address_of_()
 
-for pos in css_for_each_descendant_pre(css):
-    if sys.argv[-1] == "bpf":
-        print_cgroup_bpf_progs(pos.cgroup)
-    else:
+    for pos in css_for_each_descendant_pre(css):
         print(cgroup_path(pos.cgroup).decode())
+
+
+def cmd_bpf(cgroup):
+    css = cgroup.self.address_of_()
+
+    for pos in css_for_each_descendant_pre(css):
+        print_cgroup_bpf_progs(pos.cgroup)
+
+
+def cmd_stat(cgroup):
+    stat = Counter()
+    stat_dying = Counter()
+
+    for ssid in CgroupSubsysId:
+        css = cgroup.subsys[ssid.value]
+        # XXX if subsys of offlined or cgroup rmdir'd under our hands we won't see its subtree
+        if not css:
+            continue
+        for pos in css_for_each_descendant_pre(css, False):
+            stat[ssid] +=1
+            if not pos.flags & prog["CSS_ONLINE"]:
+                stat_dying[ssid] += 1
+
+    for ssid in CgroupSubsysId:
+        if stat[ssid.value] == 0:
+            continue
+        print("nr_{:<30} {:>4}".format(
+            ssid.name,
+            stat[ssid.value]
+            )
+        )
+    for ssid in CgroupSubsysId:
+        if stat_dying[ssid.value] == 0:
+            continue
+        print("nr_dying_{:<24} {:>4}".format(
+            ssid.name,
+            stat_dying[ssid.value]
+            )
+        )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("command", choices=["tree", "bpf", "stat"])
+    parser.add_argument("cgroups", help="Cgroups", nargs="*", type=get_cgroup)
+    args = parser.parse_args()
+
+    if len(args.cgroups) == 0:
+        args.cgroups.append(prog["cgrp_dfl_root"].cgrp)
+
+    for cg in args.cgroups:
+        if len(args.cgroups) > 1:
+            print(cg.kn.name.string_())
+        locals()["cmd_" + args.command](cg)
+

drgn/helpers/linux/cgroup.py

@@ -152,14 +152,14 @@ def css_next_descendant_pre(pos: Object, root: Object) -> Object:
 
 
 def _css_for_each_impl(
-    next_fn: Callable[[Object, Object], Object], css: Object
+    next_fn: Callable[[Object, Object], Object], css: Object, only_online: bool
 ) -> Iterator[Object]:
     pos = NULL(css.prog_, "struct cgroup_subsys_state *")
     while True:
         pos = next_fn(pos, css)
         if not pos:
             break
-        if pos.flags & pos.prog_["CSS_ONLINE"]:
+        if not only_online or pos.flags & pos.prog_["CSS_ONLINE"]:
             yield pos
 
 
@@ -170,14 +170,14 @@ def css_for_each_child(css: Object) -> Iterator[Object]:
     :param css: ``struct cgroup_subsys_state *``
     :return: Iterator of ``struct cgroup_subsys_state *`` objects.
     """
-    return _css_for_each_impl(css_next_child, css)
+    return _css_for_each_impl(css_next_child, css, True)
 
 
-def css_for_each_descendant_pre(css: Object) -> Iterator[Object]:
+def css_for_each_descendant_pre(css: Object, only_online: bool = True) -> Iterator[Object]:
     """
     Iterate through the given css's descendants in pre-order.
 
     :param css: ``struct cgroup_subsys_state *``
     :return: Iterator of ``struct cgroup_subsys_state *`` objects.
     """
-    return _css_for_each_impl(css_next_descendant_pre, css)
+    return _css_for_each_impl(css_next_descendant_pre, css, only_online)