xref: /aosp_15_r20/external/bcc/examples/local_storage/task_storage.py (revision 387f9dfdfa2baef462e92476d413c7bc2470293e)
1#!/usr/bin/python3
2
3from bcc import BPF
4
5source = r"""
6BPF_TASK_STORAGE(task_storage_map, __u64);
7
8KFUNC_PROBE(inet_listen)
9{
10	__u64 ts = bpf_ktime_get_ns();
11
12	/* save timestamp to local storage on function entry */
13	task_storage_map.task_storage_get(bpf_get_current_task_btf(), &ts, BPF_LOCAL_STORAGE_GET_F_CREATE);
14
15	bpf_trace_printk("inet_listen entry: store timestamp %lld", ts);
16	return 0;
17}
18
19KRETFUNC_PROBE(inet_listen)
20{
21	__u64 *ts;
22
23	/* retrieve timestamp stored at local storage on function exit */
24	ts = task_storage_map.task_storage_get(bpf_get_current_task_btf(), 0, 0);
25	if (!ts)
26		return 0;
27
28	/* delete timestamp from local storage */
29	task_storage_map.task_storage_delete(bpf_get_current_task_btf());
30
31	/* calculate latency */
32	bpf_trace_printk("inet_listen exit: cost %lldus", (bpf_ktime_get_ns() - *ts) / 1000);
33	return 0;
34}
35"""
36
37b = BPF(text=source)
38try:
39    b.trace_print()
40except KeyboardInterrupt:
41    pass
42