From 2d9b918b862ee94c21136e630f0dd9d29b887246 Mon Sep 17 00:00:00 2001
From: Imran Khan <imran.f.khan@oracle.com>
Date: Thu, 1 Aug 2024 14:46:58 +1000
Subject: [PATCH] kernfs,memcg: Add helpers to get memcg info

Add kernfs based helpers to extract memcg related information,
like number of active and inactive memcgroups, page cache pages
pinning memcgroups etc.

Signed-off-by: Imran Khan <imran.f.khan@oracle.com>
---
 drgn_tools/kernfs_memcg.py | 123 +++++++++++++++++++++++++++++++++++++
 1 file changed, 123 insertions(+)
 create mode 100644 drgn_tools/kernfs_memcg.py

diff --git a/drgn_tools/kernfs_memcg.py b/drgn_tools/kernfs_memcg.py
new file mode 100644
index 00000000..c70ed62d
--- /dev/null
+++ b/drgn_tools/kernfs_memcg.py
@@ -0,0 +1,123 @@
+from drgn import cast
+from drgn import FaultError
+from drgn import Object
+from drgn import Program
+from drgn.helpers.linux import cgroup_path
+from drgn.helpers.linux import css_for_each_descendant_pre
+from drgn.helpers.linux import dentry_path
+from drgn.helpers.linux import find_slab_cache
+from drgn.helpers.linux import for_each_page
+from drgn.helpers.linux import kernfs_path
+from drgn.helpers.linux import slab_cache_for_each_allocated_object
+
+
+def for_each_kernfs_node(prog: Program):
+    kernfs_node_cache = find_slab_cache(prog, "kernfs_node_cache")
+    for kn in slab_cache_for_each_allocated_object(
+        kernfs_node_cache, "struct kernfs_node"
+    ):
+        yield kn
+
+
+def dump_memcgroup_hierarchy(prog: Program):
+    cgroup_subsys = prog["cgroup_subsys"][4]
+    css = cgroup_subsys.root.cgrp.self.address_of_()
+    print(f"dumping: {cgroup_subsys.name.string_().decode()} hierarchy")
+    for pos in css_for_each_descendant_pre(css):
+        print(
+            f"path: {cgroup_path(pos.cgroup).decode()} flags: 0x{pos.flags.value_():x}"
+        )
+
+
+def kernfs_node_of_cgroup(kn: Object) -> bool:
+    if (kn.flags.value_() & 0xF) == 0x1:
+        try:
+            cgrp = Object(kn.prog_, "struct cgroup", address=kn.priv.value_())
+            return cgrp.kn == kn
+        except FaultError:
+            return False
+    else:
+        return False
+
+
+def kernfs_node_of_memcgroup(kn: Object) -> bool:
+    if kernfs_node_of_cgroup(kn):
+        prog = kn.prog_
+        cgrp = Object(prog, "struct cgroup", address=kn.priv.value_())
+        return prog["cgroup_subsys"][4].root == cgrp.root
+    else:
+        return False
+
+
+def dump_memcg_kernfs_nodes(prog):
+    count = 0
+    for kn in for_each_kernfs_node(prog):
+        if kernfs_node_of_memcgroup(kn):
+            count = count + 1
+            path = kernfs_path(kn).decode()
+            print("kernfs_node: ", hex(kn.value_()), "  ", path)
+
+    print("Total number of memcg kernfs_node objects: ", count)
+
+
+def get_num_active_mem_cgroups(prog: Program) -> int:
+    mem_cgroup_subsys = prog["cgroup_subsys"][4]
+    # add 1 to number of active memcgroups to account for root memcgroup
+    return mem_cgroup_subsys.root.cgrp.nr_descendants.value_() + 1
+
+
+def get_num_dying_mem_cgroups(prog: Program) -> int:
+    mem_cgroup_subsys = prog["cgroup_subsys"][4]
+    return mem_cgroup_subsys.root.cgrp.nr_dying_descendants.value_()
+
+
+# By default we scan all pages, that have memcg ref
+# but if max_pages is specified then we bail out
+# after getting those many pages or scanning all
+# pages , whichever happens first
+def dump_page_cache_pages_pinning_cgroups(
+    prog: Program, file: str, max_pages=0
+):
+    PG_slab_mask = 1 << prog.constant("PG_slab")
+    mem_cgroup_root = prog["cgroup_subsys"][4].root
+    total_count = 0
+    found_count = 0
+    for page in for_each_page(prog):
+        total_count = total_count + 1
+        try:
+            # Ignore slab pages
+            if page.flags & PG_slab_mask:
+                continue
+            # Ignore non page-cache pages
+            if not page.mapping:
+                continue
+            try:
+                mem_cgroup = page.mem_cgroup
+            except AttributeError:
+                mem_cgroup = page.memcg_data
+
+            if not mem_cgroup.value_() or mem_cgroup.value_() & 3:
+                continue
+            cgroup_subsys_state = cast(
+                "struct cgroup_subsys_state *", mem_cgroup
+            )
+            if cgroup_subsys_state.cgroup.root == mem_cgroup_root:
+                found_count = found_count + 1
+                cgrp = cgroup_subsys_state.cgroup
+                address_space = page.mapping
+                inode = address_space.host
+                dentry_addr = inode.i_dentry.first.value_() - 0xB0
+                if dentry_addr <= 0:
+                    continue
+                dentry = Object(prog, "struct dentry", address=dentry_addr)
+                print(
+                    f"page: 0x{page.value_():x} cgroup: {cgroup_path(cgrp).decode()} flags: {cgrp.self.flags.value_()} dpath: {dentry_path(dentry.address_of_()).decode()}\n"
+                )
+                if max_pages and found_count == max_pages:
+                    break
+        except FaultError:
+            continue
+
+    print(
+        f"Scanned {total_count} pages, found {found_count} pages with memory cgroup refs."
+    )