tools: fix klockstat when DEBUG_LOCK_ALLOC is set

When DEBUG_LOCK_ALLOC is set, mutex_lock is changed to
mutex_lock_nested, which means that attaching the kprobe statically to
mutex_lock won't work for all kernel configurations.

This change detects the available kprobe functions exposed by the kernel
and ensures that the kprobe is attached correctly.

Fixes iovisor#3882

tools/klockstat.py

@@ -367,9 +367,30 @@ def stack_id_err(stack_id):
 
 """
 
+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 @@ def display(sort, maxs, totals, counts):
 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");