19c01546dSAndrii Nakryiko // SPDX-License-Identifier: GPL-2.0
29c01546dSAndrii Nakryiko // Copyright (c) 2019 Facebook
39c01546dSAndrii Nakryiko #include "vmlinux.h"
45b554ce5SToke Høiland-Jørgensen #include <bpf/bpf_helpers.h>
59c01546dSAndrii Nakryiko #include "runqslower.h"
69c01546dSAndrii Nakryiko 
79c01546dSAndrii Nakryiko #define TASK_RUNNING 0
837ccc12bSAndrii Nakryiko #define BPF_F_CURRENT_CPU 0xffffffffULL
99c01546dSAndrii Nakryiko 
109c01546dSAndrii Nakryiko const volatile __u64 min_us = 0;
119c01546dSAndrii Nakryiko const volatile pid_t targ_pid = 0;
129c01546dSAndrii Nakryiko 
139c01546dSAndrii Nakryiko struct {
14ced47e30SSong Liu 	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
15ced47e30SSong Liu 	__uint(map_flags, BPF_F_NO_PREALLOC);
16ced47e30SSong Liu 	__type(key, int);
179c01546dSAndrii Nakryiko 	__type(value, u64);
189c01546dSAndrii Nakryiko } start SEC(".maps");
199c01546dSAndrii Nakryiko 
209c01546dSAndrii Nakryiko struct {
219c01546dSAndrii Nakryiko 	__uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
229c01546dSAndrii Nakryiko 	__uint(key_size, sizeof(u32));
239c01546dSAndrii Nakryiko 	__uint(value_size, sizeof(u32));
249c01546dSAndrii Nakryiko } events SEC(".maps");
259c01546dSAndrii Nakryiko 
269c01546dSAndrii Nakryiko /* record enqueue timestamp */
279c01546dSAndrii Nakryiko __always_inline
trace_enqueue(struct task_struct * t)28ced47e30SSong Liu static int trace_enqueue(struct task_struct *t)
299c01546dSAndrii Nakryiko {
30ced47e30SSong Liu 	u32 pid = t->pid;
31ced47e30SSong Liu 	u64 *ptr;
329c01546dSAndrii Nakryiko 
339c01546dSAndrii Nakryiko 	if (!pid || (targ_pid && targ_pid != pid))
349c01546dSAndrii Nakryiko 		return 0;
359c01546dSAndrii Nakryiko 
36ced47e30SSong Liu 	ptr = bpf_task_storage_get(&start, t, 0,
37ced47e30SSong Liu 				   BPF_LOCAL_STORAGE_GET_F_CREATE);
38ced47e30SSong Liu 	if (!ptr)
39ced47e30SSong Liu 		return 0;
40ced47e30SSong Liu 
41ced47e30SSong Liu 	*ptr = bpf_ktime_get_ns();
429c01546dSAndrii Nakryiko 	return 0;
439c01546dSAndrii Nakryiko }
449c01546dSAndrii Nakryiko 
459c01546dSAndrii Nakryiko SEC("tp_btf/sched_wakeup")
handle__sched_wakeup(u64 * ctx)469c01546dSAndrii Nakryiko int handle__sched_wakeup(u64 *ctx)
479c01546dSAndrii Nakryiko {
489c01546dSAndrii Nakryiko 	/* TP_PROTO(struct task_struct *p) */
499c01546dSAndrii Nakryiko 	struct task_struct *p = (void *)ctx[0];
509c01546dSAndrii Nakryiko 
51ced47e30SSong Liu 	return trace_enqueue(p);
529c01546dSAndrii Nakryiko }
539c01546dSAndrii Nakryiko 
549c01546dSAndrii Nakryiko SEC("tp_btf/sched_wakeup_new")
handle__sched_wakeup_new(u64 * ctx)559c01546dSAndrii Nakryiko int handle__sched_wakeup_new(u64 *ctx)
569c01546dSAndrii Nakryiko {
579c01546dSAndrii Nakryiko 	/* TP_PROTO(struct task_struct *p) */
589c01546dSAndrii Nakryiko 	struct task_struct *p = (void *)ctx[0];
599c01546dSAndrii Nakryiko 
60ced47e30SSong Liu 	return trace_enqueue(p);
619c01546dSAndrii Nakryiko }
629c01546dSAndrii Nakryiko 
639c01546dSAndrii Nakryiko SEC("tp_btf/sched_switch")
handle__sched_switch(u64 * ctx)649c01546dSAndrii Nakryiko int handle__sched_switch(u64 *ctx)
659c01546dSAndrii Nakryiko {
669c01546dSAndrii Nakryiko 	/* TP_PROTO(bool preempt, struct task_struct *prev,
679c01546dSAndrii Nakryiko 	 *	    struct task_struct *next)
689c01546dSAndrii Nakryiko 	 */
699c01546dSAndrii Nakryiko 	struct task_struct *prev = (struct task_struct *)ctx[1];
709c01546dSAndrii Nakryiko 	struct task_struct *next = (struct task_struct *)ctx[2];
71*88a71086SNaveen N. Rao 	struct runq_event event = {};
729c01546dSAndrii Nakryiko 	u64 *tsp, delta_us;
739c01546dSAndrii Nakryiko 	long state;
749c01546dSAndrii Nakryiko 	u32 pid;
759c01546dSAndrii Nakryiko 
769c01546dSAndrii Nakryiko 	/* ivcsw: treat like an enqueue event and store timestamp */
775616e895SSanjayKumar Jeyakumar 	if (prev->__state == TASK_RUNNING)
78ced47e30SSong Liu 		trace_enqueue(prev);
799c01546dSAndrii Nakryiko 
809c01546dSAndrii Nakryiko 	pid = next->pid;
819c01546dSAndrii Nakryiko 
82ced47e30SSong Liu 	/* For pid mismatch, save a bpf_task_storage_get */
83ced47e30SSong Liu 	if (!pid || (targ_pid && targ_pid != pid))
84ced47e30SSong Liu 		return 0;
85ced47e30SSong Liu 
869c01546dSAndrii Nakryiko 	/* fetch timestamp and calculate delta */
87ced47e30SSong Liu 	tsp = bpf_task_storage_get(&start, next, 0, 0);
889c01546dSAndrii Nakryiko 	if (!tsp)
899c01546dSAndrii Nakryiko 		return 0;   /* missed enqueue */
909c01546dSAndrii Nakryiko 
919c01546dSAndrii Nakryiko 	delta_us = (bpf_ktime_get_ns() - *tsp) / 1000;
929c01546dSAndrii Nakryiko 	if (min_us && delta_us <= min_us)
939c01546dSAndrii Nakryiko 		return 0;
949c01546dSAndrii Nakryiko 
959c01546dSAndrii Nakryiko 	event.pid = pid;
969c01546dSAndrii Nakryiko 	event.delta_us = delta_us;
979c01546dSAndrii Nakryiko 	bpf_get_current_comm(&event.task, sizeof(event.task));
989c01546dSAndrii Nakryiko 
999c01546dSAndrii Nakryiko 	/* output */
1009c01546dSAndrii Nakryiko 	bpf_perf_event_output(ctx, &events, BPF_F_CURRENT_CPU,
1019c01546dSAndrii Nakryiko 			      &event, sizeof(event));
1029c01546dSAndrii Nakryiko 
103ced47e30SSong Liu 	bpf_task_storage_delete(&start, next);
1049c01546dSAndrii Nakryiko 	return 0;
1059c01546dSAndrii Nakryiko }
1069c01546dSAndrii Nakryiko 
1079c01546dSAndrii Nakryiko char LICENSE[] SEC("license") = "GPL";
108