162306a36Sopenharmony_ci// SPDX-License-Identifier: (GPL-2.0-only OR BSD-2-Clause)
262306a36Sopenharmony_ci// Copyright (c) 2021 Google
362306a36Sopenharmony_ci#include "vmlinux.h"
462306a36Sopenharmony_ci#include <bpf/bpf_helpers.h>
562306a36Sopenharmony_ci#include <bpf/bpf_tracing.h>
662306a36Sopenharmony_ci
762306a36Sopenharmony_ci// This should be in sync with "util/ftrace.h"
862306a36Sopenharmony_ci#define NUM_BUCKET  22
962306a36Sopenharmony_ci
1062306a36Sopenharmony_cistruct {
1162306a36Sopenharmony_ci	__uint(type, BPF_MAP_TYPE_HASH);
1262306a36Sopenharmony_ci	__uint(key_size, sizeof(__u64));
1362306a36Sopenharmony_ci	__uint(value_size, sizeof(__u64));
1462306a36Sopenharmony_ci	__uint(max_entries, 10000);
1562306a36Sopenharmony_ci} functime SEC(".maps");
1662306a36Sopenharmony_ci
1762306a36Sopenharmony_cistruct {
1862306a36Sopenharmony_ci	__uint(type, BPF_MAP_TYPE_HASH);
1962306a36Sopenharmony_ci	__uint(key_size, sizeof(__u32));
2062306a36Sopenharmony_ci	__uint(value_size, sizeof(__u8));
2162306a36Sopenharmony_ci	__uint(max_entries, 1);
2262306a36Sopenharmony_ci} cpu_filter SEC(".maps");
2362306a36Sopenharmony_ci
2462306a36Sopenharmony_cistruct {
2562306a36Sopenharmony_ci	__uint(type, BPF_MAP_TYPE_HASH);
2662306a36Sopenharmony_ci	__uint(key_size, sizeof(__u32));
2762306a36Sopenharmony_ci	__uint(value_size, sizeof(__u8));
2862306a36Sopenharmony_ci	__uint(max_entries, 1);
2962306a36Sopenharmony_ci} task_filter SEC(".maps");
3062306a36Sopenharmony_ci
3162306a36Sopenharmony_cistruct {
3262306a36Sopenharmony_ci	__uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
3362306a36Sopenharmony_ci	__uint(key_size, sizeof(__u32));
3462306a36Sopenharmony_ci	__uint(value_size, sizeof(__u64));
3562306a36Sopenharmony_ci	__uint(max_entries, NUM_BUCKET);
3662306a36Sopenharmony_ci} latency SEC(".maps");
3762306a36Sopenharmony_ci
3862306a36Sopenharmony_ci
3962306a36Sopenharmony_ciint enabled = 0;
4062306a36Sopenharmony_ciint has_cpu = 0;
4162306a36Sopenharmony_ciint has_task = 0;
4262306a36Sopenharmony_ciint use_nsec = 0;
4362306a36Sopenharmony_ci
4462306a36Sopenharmony_ciSEC("kprobe/func")
4562306a36Sopenharmony_ciint BPF_PROG(func_begin)
4662306a36Sopenharmony_ci{
4762306a36Sopenharmony_ci	__u64 key, now;
4862306a36Sopenharmony_ci
4962306a36Sopenharmony_ci	if (!enabled)
5062306a36Sopenharmony_ci		return 0;
5162306a36Sopenharmony_ci
5262306a36Sopenharmony_ci	key = bpf_get_current_pid_tgid();
5362306a36Sopenharmony_ci
5462306a36Sopenharmony_ci	if (has_cpu) {
5562306a36Sopenharmony_ci		__u32 cpu = bpf_get_smp_processor_id();
5662306a36Sopenharmony_ci		__u8 *ok;
5762306a36Sopenharmony_ci
5862306a36Sopenharmony_ci		ok = bpf_map_lookup_elem(&cpu_filter, &cpu);
5962306a36Sopenharmony_ci		if (!ok)
6062306a36Sopenharmony_ci			return 0;
6162306a36Sopenharmony_ci	}
6262306a36Sopenharmony_ci
6362306a36Sopenharmony_ci	if (has_task) {
6462306a36Sopenharmony_ci		__u32 pid = key & 0xffffffff;
6562306a36Sopenharmony_ci		__u8 *ok;
6662306a36Sopenharmony_ci
6762306a36Sopenharmony_ci		ok = bpf_map_lookup_elem(&task_filter, &pid);
6862306a36Sopenharmony_ci		if (!ok)
6962306a36Sopenharmony_ci			return 0;
7062306a36Sopenharmony_ci	}
7162306a36Sopenharmony_ci
7262306a36Sopenharmony_ci	now = bpf_ktime_get_ns();
7362306a36Sopenharmony_ci
7462306a36Sopenharmony_ci	// overwrite timestamp for nested functions
7562306a36Sopenharmony_ci	bpf_map_update_elem(&functime, &key, &now, BPF_ANY);
7662306a36Sopenharmony_ci	return 0;
7762306a36Sopenharmony_ci}
7862306a36Sopenharmony_ci
7962306a36Sopenharmony_ciSEC("kretprobe/func")
8062306a36Sopenharmony_ciint BPF_PROG(func_end)
8162306a36Sopenharmony_ci{
8262306a36Sopenharmony_ci	__u64 tid;
8362306a36Sopenharmony_ci	__u64 *start;
8462306a36Sopenharmony_ci	__u64 cmp_base = use_nsec ? 1 : 1000;
8562306a36Sopenharmony_ci
8662306a36Sopenharmony_ci	if (!enabled)
8762306a36Sopenharmony_ci		return 0;
8862306a36Sopenharmony_ci
8962306a36Sopenharmony_ci	tid = bpf_get_current_pid_tgid();
9062306a36Sopenharmony_ci
9162306a36Sopenharmony_ci	start = bpf_map_lookup_elem(&functime, &tid);
9262306a36Sopenharmony_ci	if (start) {
9362306a36Sopenharmony_ci		__s64 delta = bpf_ktime_get_ns() - *start;
9462306a36Sopenharmony_ci		__u32 key;
9562306a36Sopenharmony_ci		__u64 *hist;
9662306a36Sopenharmony_ci
9762306a36Sopenharmony_ci		bpf_map_delete_elem(&functime, &tid);
9862306a36Sopenharmony_ci
9962306a36Sopenharmony_ci		if (delta < 0)
10062306a36Sopenharmony_ci			return 0;
10162306a36Sopenharmony_ci
10262306a36Sopenharmony_ci		// calculate index using delta
10362306a36Sopenharmony_ci		for (key = 0; key < (NUM_BUCKET - 1); key++) {
10462306a36Sopenharmony_ci			if (delta < (cmp_base << key))
10562306a36Sopenharmony_ci				break;
10662306a36Sopenharmony_ci		}
10762306a36Sopenharmony_ci
10862306a36Sopenharmony_ci		hist = bpf_map_lookup_elem(&latency, &key);
10962306a36Sopenharmony_ci		if (!hist)
11062306a36Sopenharmony_ci			return 0;
11162306a36Sopenharmony_ci
11262306a36Sopenharmony_ci		*hist += 1;
11362306a36Sopenharmony_ci	}
11462306a36Sopenharmony_ci
11562306a36Sopenharmony_ci	return 0;
11662306a36Sopenharmony_ci}
117