1// SPDX-License-Identifier: GPL-2.0
  2#include <stdio.h>
  3#include <unistd.h>
  4#include <stdlib.h>
  5#include <signal.h>
  6#include <linux/bpf.h>
  7#include <string.h>
  8#include <sys/resource.h>
  9
 10#include <bpf/bpf.h>
 11#include "bpf_load.h"
 12#include "bpf_util.h"
 13
 14#define MAX_INDEX	64
 15#define MAX_STARS	38
 16
 17static void stars(char *str, long val, long max, int width)
 18{
 19	int i;
 20
 21	for (i = 0; i < (width * val / max) - 1 && i < width - 1; i++)
 22		str[i] = '*';
 23	if (val > max)
 24		str[i - 1] = '+';
 25	str[i] = '\0';
 26}
 27
 28struct task {
 29	char comm[16];
 30	__u64 pid_tgid;
 31	__u64 uid_gid;
 32};
 33
 34struct hist_key {
 35	struct task t;
 36	__u32 index;
 37};
 38
 39#define SIZE sizeof(struct task)
 40
 41static void print_hist_for_pid(int fd, void *task)
 42{
 43	unsigned int nr_cpus = bpf_num_possible_cpus();
 44	struct hist_key key = {}, next_key;
 45	long values[nr_cpus];
 46	char starstr[MAX_STARS];
 47	long value;
 48	long data[MAX_INDEX] = {};
 49	int max_ind = -1;
 50	long max_value = 0;
 51	int i, ind;
 52
 53	while (bpf_map_get_next_key(fd, &key, &next_key) == 0) {
 54		if (memcmp(&next_key, task, SIZE)) {
 55			key = next_key;
 56			continue;
 57		}
 58		bpf_map_lookup_elem(fd, &next_key, values);
 59		value = 0;
 60		for (i = 0; i < nr_cpus; i++)
 61			value += values[i];
 62		ind = next_key.index;
 63		data[ind] = value;
 64		if (value && ind > max_ind)
 65			max_ind = ind;
 66		if (value > max_value)
 67			max_value = value;
 68		key = next_key;
 69	}
 70
 71	printf("           syscall write() stats\n");
 72	printf("     byte_size       : count     distribution\n");
 73	for (i = 1; i <= max_ind + 1; i++) {
 74		stars(starstr, data[i - 1], max_value, MAX_STARS);
 75		printf("%8ld -> %-8ld : %-8ld |%-*s|\n",
 76		       (1l << i) >> 1, (1l << i) - 1, data[i - 1],
 77		       MAX_STARS, starstr);
 78	}
 79}
 80
 81static void print_hist(int fd)
 82{
 83	struct hist_key key = {}, next_key;
 84	static struct task tasks[1024];
 85	int task_cnt = 0;
 86	int i;
 87
 88	while (bpf_map_get_next_key(fd, &key, &next_key) == 0) {
 89		int found = 0;
 90
 91		for (i = 0; i < task_cnt; i++)
 92			if (memcmp(&tasks[i], &next_key, SIZE) == 0)
 93				found = 1;
 94		if (!found)
 95			memcpy(&tasks[task_cnt++], &next_key, SIZE);
 96		key = next_key;
 97	}
 98
 99	for (i = 0; i < task_cnt; i++) {
100		printf("\npid %d cmd %s uid %d\n",
101		       (__u32) tasks[i].pid_tgid,
102		       tasks[i].comm,
103		       (__u32) tasks[i].uid_gid);
104		print_hist_for_pid(fd, &tasks[i]);
105	}
106
107}
108
109static void int_exit(int sig)
110{
111	print_hist(map_fd[1]);
112	exit(0);
113}
114
115int main(int ac, char **argv)
116{
117	struct rlimit r = {1024*1024, RLIM_INFINITY};
118	char filename[256];
119	long key, next_key, value;
120	FILE *f;
121	int i;
122
123	snprintf(filename, sizeof(filename), "%s_kern.o", argv[0]);
124
125	if (setrlimit(RLIMIT_MEMLOCK, &r)) {
126		perror("setrlimit(RLIMIT_MEMLOCK)");
127		return 1;
128	}
129
130	signal(SIGINT, int_exit);
131	signal(SIGTERM, int_exit);
132
133	/* start 'ping' in the background to have some kfree_skb events */
134	f = popen("ping -4 -c5 localhost", "r");
135	(void) f;
136
137	/* start 'dd' in the background to have plenty of 'write' syscalls */
138	f = popen("dd if=/dev/zero of=/dev/null count=5000000", "r");
139	(void) f;
140
141	if (load_bpf_file(filename)) {
142		printf("%s", bpf_log_buf);
143		return 1;
144	}
145
146	for (i = 0; i < 5; i++) {
147		key = 0;
148		while (bpf_map_get_next_key(map_fd[0], &key, &next_key) == 0) {
149			bpf_map_lookup_elem(map_fd[0], &next_key, &value);
150			printf("location 0x%lx count %ld\n", next_key, value);
151			key = next_key;
152		}
153		if (key)
154			printf("\n");
155		sleep(1);
156	}
157	print_hist(map_fd[1]);
158
159	return 0;
160}