1// SPDX-License-Identifier: GPL-2.0
  2// Copyright (c) 2019 Facebook
  3#include "vmlinux.h"
  4#include <bpf/bpf_helpers.h>
  5#include "runqslower.h"
  6
  7#define TASK_RUNNING 0
  8#define BPF_F_CURRENT_CPU 0xffffffffULL
  9
 10const volatile __u64 min_us = 0;
 11const volatile pid_t targ_pid = 0;
 12
 13struct {
 14	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
 15	__uint(map_flags, BPF_F_NO_PREALLOC);
 16	__type(key, int);
 17	__type(value, u64);
 18} start SEC(".maps");
 19
 20struct {
 21	__uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
 22	__uint(key_size, sizeof(u32));
 23	__uint(value_size, sizeof(u32));
 24} events SEC(".maps");
 25
 26/* record enqueue timestamp */
 27__always_inline
 28static int trace_enqueue(struct task_struct *t)
 29{
 30	u32 pid = t->pid;
 31	u64 *ptr;
 32
 33	if (!pid || (targ_pid && targ_pid != pid))
 34		return 0;
 35
 36	ptr = bpf_task_storage_get(&start, t, 0,
 37				   BPF_LOCAL_STORAGE_GET_F_CREATE);
 38	if (!ptr)
 39		return 0;
 40
 41	*ptr = bpf_ktime_get_ns();
 42	return 0;
 43}
 44
 45SEC("tp_btf/sched_wakeup")
 46int handle__sched_wakeup(u64 *ctx)
 47{
 48	/* TP_PROTO(struct task_struct *p) */
 49	struct task_struct *p = (void *)ctx[0];
 50
 51	return trace_enqueue(p);
 52}
 53
 54SEC("tp_btf/sched_wakeup_new")
 55int handle__sched_wakeup_new(u64 *ctx)
 56{
 57	/* TP_PROTO(struct task_struct *p) */
 58	struct task_struct *p = (void *)ctx[0];
 59
 60	return trace_enqueue(p);
 61}
 62
 63SEC("tp_btf/sched_switch")
 64int handle__sched_switch(u64 *ctx)
 65{
 66	/* TP_PROTO(bool preempt, struct task_struct *prev,
 67	 *	    struct task_struct *next)
 68	 */
 69	struct task_struct *prev = (struct task_struct *)ctx[1];
 70	struct task_struct *next = (struct task_struct *)ctx[2];
 71	struct runq_event event = {};
 72	u64 *tsp, delta_us;
 73	u32 pid;
 74
 75	/* ivcsw: treat like an enqueue event and store timestamp */
 76	if (prev->__state == TASK_RUNNING)
 77		trace_enqueue(prev);
 78
 79	pid = next->pid;
 80
 81	/* For pid mismatch, save a bpf_task_storage_get */
 82	if (!pid || (targ_pid && targ_pid != pid))
 83		return 0;
 84
 85	/* fetch timestamp and calculate delta */
 86	tsp = bpf_task_storage_get(&start, next, 0, 0);
 87	if (!tsp)
 88		return 0;   /* missed enqueue */
 89
 90	delta_us = (bpf_ktime_get_ns() - *tsp) / 1000;
 91	if (min_us && delta_us <= min_us)
 92		return 0;
 93
 94	event.pid = pid;
 95	event.delta_us = delta_us;
 96	bpf_get_current_comm(&event.task, sizeof(event.task));
 97
 98	/* output */
 99	bpf_perf_event_output(ctx, &events, BPF_F_CURRENT_CPU,
100			      &event, sizeof(event));
101
102	bpf_task_storage_delete(&start, next);
103	return 0;
104}
105
106char LICENSE[] SEC("license") = "GPL";