1 files changed, 30 insertions, 4 deletions
diff --git a/tools/klockstat.py b/tools/klockstat.py
index d157b7be..b8cafd97 100755
--- a/tools/klockstat.py
+++ b/tools/klockstat.py
@@ -367,9 +367,30 @@ KFUNC_PROBE(mutex_lock, void *lock)
 
 """
 
+program_kfunc_nested = """
+KFUNC_PROBE(mutex_unlock, void *lock)
+{
+    return do_mutex_unlock_enter();
+}
+
+KRETFUNC_PROBE(mutex_lock_nested, void *lock, int ret)
+{
+    return do_mutex_lock_return();
+}
+
+KFUNC_PROBE(mutex_lock_nested, void *lock)
+{
+    return do_mutex_lock_enter(ctx, 3);
+}
+
+"""
+
 is_support_kfunc = BPF.support_kfunc()
 if is_support_kfunc:
-    program += program_kfunc
+    if BPF.get_kprobe_functions(b"mutex_lock_nested"):
+        program += program_kfunc_nested
+    else:
+        program += program_kfunc
 else:
     program += program_kprobe
 
@@ -428,9 +449,14 @@ program = program.replace('STACK_STORAGE_SIZE', str(args.stack_storage_size))
 b = BPF(text=program)
 
 if not is_support_kfunc:
-    b.attach_kprobe(event="mutex_unlock",  fn_name="mutex_unlock_enter")
-    b.attach_kretprobe(event="mutex_lock", fn_name="mutex_lock_return")
-    b.attach_kprobe(event="mutex_lock",    fn_name="mutex_lock_enter")
+    b.attach_kprobe(event="mutex_unlock", fn_name="mutex_unlock_enter")
+    # Depending on whether DEBUG_LOCK_ALLOC is set, the proper kprobe may be either mutex_lock or mutex_lock_nested
+    if BPF.get_kprobe_functions(b"mutex_lock_nested"):
+        b.attach_kretprobe(event="mutex_lock_nested", fn_name="mutex_lock_return")
+        b.attach_kprobe(event="mutex_lock_nested", fn_name="mutex_lock_enter")
+    else:
+        b.attach_kretprobe(event="mutex_lock", fn_name="mutex_lock_return")
+        b.attach_kprobe(event="mutex_lock", fn_name="mutex_lock_enter")
 
 enabled = b.get_table("enabled");