[PATCH] tools/cgroup/slabinfo: Fix use of slab.memcg_data

[Hongfu Li]

After the introduce slabobj_ext to support slab object extensions, the
memcg_slabinfo tool broke. An attempt to run it produces a trace like
this:
Traceback (most recent call last):
  File "/usr/local/bin/drgn", line 8, in <module>
    sys.exit(_main())
             ^^^^^^^
  File "/usr/local/lib64/python3.11/site-packages/drgn/cli.py", line 688, in _main
    runpy.run_path(
  File "<frozen runpy>", line 291, in run_path
  File "<frozen runpy>", line 98, in _run_module_code
  File "<frozen runpy>", line 88, in _run_code
  File "/root/memcg_slabinfo.py", line 225, in <module>
    main()
  File "/root/memcg_slabinfo.py", line 195, in main
    objcg_vec_raw = slab.memcg_data.value_()
AttributeError: 'struct slab' has no member 'memcg_data'

Fixes: 21c690a349ba ("mm: introduce slabobj_ext to support slab object extensions")
Signed-off-by: Hongfu Li <lihongfu@xxxxxxxxxx>
---
 tools/cgroup/memcg_slabinfo.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/tools/cgroup/memcg_slabinfo.py b/tools/cgroup/memcg_slabinfo.py
index 6bf4bde77903..e67dcfae1263 100644
--- a/tools/cgroup/memcg_slabinfo.py
+++ b/tools/cgroup/memcg_slabinfo.py
@@ -33,6 +33,13 @@ def err(s):
     sys.exit(1)
 
 
+def objexts_flags_mask():
+    try:
+        return int(prog.constant('__NR_OBJEXTS_FLAGS')) - 1
+    except:
+        return 0x7
+
+
 def find_memcg_ids(css=prog['root_mem_cgroup'].css, prefix=''):
     if not list_empty(css.children.address_of_()):
         for css in list_for_each_entry('struct cgroup_subsys_state',
@@ -192,23 +199,24 @@ def main():
         # look over all slab folios and look for objects belonging
         # to the given memory cgroup
         for slab in for_each_slab(prog):
-            objcg_vec_raw = slab.memcg_data.value_()
-            if objcg_vec_raw == 0:
+            objext_vec_raw = slab.obj_exts.value_()
+            if objext_vec_raw == 0:
                 continue
             cache = slab.slab_cache
             if not cache:
                 continue
             addr = cache.value_()
             caches[addr] = cache
-            # clear the lowest bit to get the true obj_cgroups
-            objcg_vec = Object(prog, 'struct obj_cgroup **',
-                               value=objcg_vec_raw & ~1)
 
             if addr not in stats:
                 stats[addr] = 0
 
+            # clear OBJEXTS_FLAGS_MASK bits to get the true slabobj_ext
+            mask = objexts_flags_mask()
+            objext_vec = Object(prog, 'struct slabobj_ext *',
+                                value=objext_vec_raw & ~mask)
             for i in range(oo_objects(cache)):
-                if objcg_vec[i].value_() in obj_cgroups:
+                if objext_vec[i].objcg.value_() in obj_cgroups:
                     stats[addr] += 1
 
         for addr in caches:
-- 
2.25.1