lockstat.py

from __future__ import print_function

from bcc import BPF
import datetime


prog = """
#include <linux/sched.h>
#include <linux/spinlock_types.h>

// define struct for key

struct key_t {
    u64 pid;
    raw_spinlock_t* lock;
};
// define output data structure in C
struct data_t {
    u32 pid;
    u32 tid;
    u64 ts;
    char comm[TASK_COMM_LEN];
    u64 lock;
    u64 lock_time;
    u64 present_time;
    u64 diff;
    u64 stack_id;
    u32 lock_count;
    u32 type;
};

BPF_STACK_TRACE(stack_traces, 102400);

BPF_PERF_OUTPUT(output);
BPF_HASH(lock_events, struct key_t, struct data_t, 102400);

int lock(struct pt_regs *ctx, raw_spinlock_t *lock) {

    u32 current_pid = bpf_get_current_pid_tgid();
        
    struct data_t data = {};
    struct key_t key = {bpf_get_current_pid_tgid(), lock};
    struct data_t *data_ptr;
    data_ptr = lock_events.lookup(&key);
    if(data_ptr)
    {
        data_ptr->ts = bpf_ktime_get_ns();
        data_ptr->lock_count += 1;
        data_ptr->stack_id = stack_traces.get_stackid(ctx, BPF_F_REUSE_STACKID);
    }
    else
    {
        data.pid = bpf_get_current_pid_tgid();
        data.tid = bpf_get_current_pid_tgid() >> 32;
        data.ts = bpf_ktime_get_ns();
        bpf_get_current_comm(&data.comm, sizeof(data.comm));
        data.lock = (u64)lock;
        data.lock_count = 1;
        data.stack_id = stack_traces.get_stackid(ctx, BPF_F_REUSE_STACKID);
        lock_events.insert(&key, &data);
    }
    return 0;
}

int release(struct pt_regs *ctx, raw_spinlock_t *lock) {
    u64 present = bpf_ktime_get_ns();
    
    u32 current_pid = bpf_get_current_pid_tgid();
        
    struct data_t *data;
    struct key_t key = {bpf_get_current_pid_tgid(), lock};
    data = lock_events.lookup(&key);
    if(data)
    {
        data->lock_time += (present - data->ts);
        data->present_time = present;
        data->diff = present - data->ts;
        data->type = 1;
        output.perf_submit(ctx, data, sizeof(struct data_t));
        //lock_events.delete(&key);
    }
    return 0;
}
"""

def print_event(cpu, data, size):
    global start
    event = b['output'].event(data)
    # for lock in locks:
    if start == 0:
        start = event.ts
    time_s = (float(event.ts - start)) / 1000000000
    print("%-18.9f %-16s %-6d %-6d %-6d %-6f %-15f %-6d" % (time_s, event.comm, event.pid, event.tid,
                                                                 event.lock,
                                                                 (float(event.present_time - start)) / 1000000000,
                                                                 event.lock_time, event.diff))


b = BPF(text=prog)
b.attach_kprobe(event="_raw_read_lock", fn_name="lock")
b.attach_kretprobe(event="_raw_read_lock", fn_name="release")


print("Tracing locks for %d seconds" % 10)

# process event
start = 0

# loop with callback to print_event
b['output'].open_perf_buffer(print_event, page_cnt=16384)
start_time = datetime.datetime.now()

while 1:
    b.perf_buffer_poll()
    time_elapsed = datetime.datetime.now() - start_time
    if time_elapsed.seconds > 5:
        exit()