1/* Copyright (c) 2016 Facebook
  2 *
  3 * This program is free software; you can redistribute it and/or
  4 * modify it under the terms of version 2 of the GNU General Public
  5 * License as published by the Free Software Foundation.
  6 */
  7#define _GNU_SOURCE
  8#include <sched.h>
  9#include <stdio.h>
 10#include <sys/types.h>
 11#include <asm/unistd.h>
 12#include <unistd.h>
 13#include <assert.h>
 14#include <sys/wait.h>
 15#include <stdlib.h>
 16#include <signal.h>
 17#include <linux/bpf.h>
 18#include <string.h>
 19#include <time.h>
 20#include <sys/resource.h>
 21#include "libbpf.h"
 22#include "bpf_load.h"
 23
 24#define MAX_CNT 1000000
 25
 26static __u64 time_get_ns(void)
 27{
 28	struct timespec ts;
 29
 30	clock_gettime(CLOCK_MONOTONIC, &ts);
 31	return ts.tv_sec * 1000000000ull + ts.tv_nsec;
 32}
 33
 34#define HASH_PREALLOC		(1 << 0)
 35#define PERCPU_HASH_PREALLOC	(1 << 1)
 36#define HASH_KMALLOC		(1 << 2)
 37#define PERCPU_HASH_KMALLOC	(1 << 3)
 38
 39static int test_flags = ~0;
 40
 41static void test_hash_prealloc(int cpu)
 42{
 43	__u64 start_time;
 44	int i;
 45
 46	start_time = time_get_ns();
 47	for (i = 0; i < MAX_CNT; i++)
 48		syscall(__NR_getuid);
 49	printf("%d:hash_map_perf pre-alloc %lld events per sec\n",
 50	       cpu, MAX_CNT * 1000000000ll / (time_get_ns() - start_time));
 51}
 52
 53static void test_percpu_hash_prealloc(int cpu)
 54{
 55	__u64 start_time;
 56	int i;
 57
 58	start_time = time_get_ns();
 59	for (i = 0; i < MAX_CNT; i++)
 60		syscall(__NR_geteuid);
 61	printf("%d:percpu_hash_map_perf pre-alloc %lld events per sec\n",
 62	       cpu, MAX_CNT * 1000000000ll / (time_get_ns() - start_time));
 63}
 64
 65static void test_hash_kmalloc(int cpu)
 66{
 67	__u64 start_time;
 68	int i;
 69
 70	start_time = time_get_ns();
 71	for (i = 0; i < MAX_CNT; i++)
 72		syscall(__NR_getgid);
 73	printf("%d:hash_map_perf kmalloc %lld events per sec\n",
 74	       cpu, MAX_CNT * 1000000000ll / (time_get_ns() - start_time));
 75}
 76
 77static void test_percpu_hash_kmalloc(int cpu)
 78{
 79	__u64 start_time;
 80	int i;
 81
 82	start_time = time_get_ns();
 83	for (i = 0; i < MAX_CNT; i++)
 84		syscall(__NR_getegid);
 85	printf("%d:percpu_hash_map_perf kmalloc %lld events per sec\n",
 86	       cpu, MAX_CNT * 1000000000ll / (time_get_ns() - start_time));
 87}
 88
 89static void loop(int cpu)
 90{
 91	cpu_set_t cpuset;
 92
 93	CPU_ZERO(&cpuset);
 94	CPU_SET(cpu, &cpuset);
 95	sched_setaffinity(0, sizeof(cpuset), &cpuset);
 96
 97	if (test_flags & HASH_PREALLOC)
 98		test_hash_prealloc(cpu);
 99
100	if (test_flags & PERCPU_HASH_PREALLOC)
101		test_percpu_hash_prealloc(cpu);
102
103	if (test_flags & HASH_KMALLOC)
104		test_hash_kmalloc(cpu);
105
106	if (test_flags & PERCPU_HASH_KMALLOC)
107		test_percpu_hash_kmalloc(cpu);
108}
109
110static void run_perf_test(int tasks)
111{
112	pid_t pid[tasks];
113	int i;
114
115	for (i = 0; i < tasks; i++) {
116		pid[i] = fork();
117		if (pid[i] == 0) {
118			loop(i);
119			exit(0);
120		} else if (pid[i] == -1) {
121			printf("couldn't spawn #%d process\n", i);
122			exit(1);
123		}
124	}
125	for (i = 0; i < tasks; i++) {
126		int status;
127
128		assert(waitpid(pid[i], &status, 0) == pid[i]);
129		assert(status == 0);
130	}
131}
132
133int main(int argc, char **argv)
134{
135	struct rlimit r = {RLIM_INFINITY, RLIM_INFINITY};
136	char filename[256];
137	int num_cpu = 8;
138
139	snprintf(filename, sizeof(filename), "%s_kern.o", argv[0]);
140	setrlimit(RLIMIT_MEMLOCK, &r);
141
142	if (argc > 1)
143		test_flags = atoi(argv[1]) ? : test_flags;
144
145	if (argc > 2)
146		num_cpu = atoi(argv[2]) ? : num_cpu;
147
148	if (load_bpf_file(filename)) {
149		printf("%s", bpf_log_buf);
150		return 1;
151	}
152
153	run_perf_test(num_cpu);
154
155	return 0;
156}