1/*
2 * iobw.c - simple I/O bandwidth benchmark
3 *
4 * This program is free software; you can redistribute it and/or
5 * modify it under the terms of the GNU General Public
6 * License as published by the Free Software Foundation; either
7 * version 2 of the License, or (at your option) any later version.
8 *
9 * This program is distributed in the hope that it will be useful,
10 * but WITHOUT ANY WARRANTY; without even the implied warranty of
11 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
12 * General Public License for more details.
13 *
14 * You should have received a copy of the GNU General Public
15 * License along with this program; if not, write to the
16 * Free Software Foundation, Inc., 59 Temple Place - Suite 330,
17 * Boston, MA 021110-1307, USA.
18 *
19 * Copyright (C) 2008 Andrea Righi <righi.andrea@gmail.com>
20 */
21
22#define _GNU_SOURCE
23#define __USE_GNU
24
25#include <errno.h>
26#include <stdio.h>
27#include <stdlib.h>
28#include <fcntl.h>
29#include <signal.h>
30#include <string.h>
31#include <unistd.h>
32#include <limits.h>
33#include <sys/types.h>
34#include <sys/stat.h>
35#include <sys/time.h>
36#include <sys/wait.h>
37
38#ifndef PAGE_SIZE
39#define PAGE_SIZE sysconf(_SC_PAGE_SIZE)
40#endif
41
42#define align(x,a)		__align_mask(x,(typeof(x))(a)-1)
43#define __align_mask(x,mask)	(((x)+(mask))&~(mask))
44#define kb(x)			((x) >> 10)
45
46const char usage[] = "Usage: iobw [-direct] threads chunk_size data_size\n";
47const char child_fmt[] = "(%s) task %3d: time %4lu.%03lu bw %7lu KiB/s (%s)\n";
48const char parent_fmt[] =
49    "(%s) parent %d: time %4lu.%03lu bw %7lu KiB/s (%s)\n";
50
51static int directio = 0;
52static size_t data_size = 0;
53static size_t chunk_size = 0;
54
55typedef enum {
56	OP_WRITE,
57	OP_READ,
58	NUM_IOPS,
59} iops_t;
60
61static const char *iops[] = {
62	"WRITE",
63	"READ ",
64	"TOTAL",
65};
66
67static int threads;
68pid_t *children;
69
70char *mygroup;
71
72static void print_results(int id, iops_t op, size_t bytes, struct timeval *diff)
73{
74	fprintf(stdout, id ? child_fmt : parent_fmt,
75		mygroup, id, diff->tv_sec, diff->tv_usec / 1000,
76		(bytes / (diff->tv_sec * 1000000L + diff->tv_usec))
77		* 1000000L / 1024, iops[op]);
78}
79
80static void thread(int id)
81{
82	struct timeval start, stop, diff;
83	int fd, i, ret;
84	size_t n;
85	void *buf;
86	int flags = O_CREAT | O_RDWR | O_LARGEFILE;
87	char filename[32];
88
89	ret = posix_memalign(&buf, PAGE_SIZE, chunk_size);
90	if (ret < 0) {
91		fprintf(stderr,
92			"ERROR: task %d couldn't allocate %zu bytes (%s)\n",
93			id, chunk_size, strerror(errno));
94		exit(1);
95	}
96	memset(buf, 0xaa, chunk_size);
97
98	snprintf(filename, sizeof(filename), "%s-%d-iobw.tmp", mygroup, id);
99	if (directio)
100		flags |= O_DIRECT;
101	fd = open(filename, flags, 0600);
102	if (fd < 0) {
103		fprintf(stderr, "ERROR: task %d couldn't open %s (%s)\n",
104			id, filename, strerror(errno));
105		free(buf);
106		exit(1);
107	}
108
109	/* Write */
110	lseek(fd, 0, SEEK_SET);
111	n = 0;
112	gettimeofday(&start, NULL);
113	while (n < data_size) {
114		i = write(fd, buf, chunk_size);
115		if (i < 0) {
116			fprintf(stderr, "ERROR: task %d writing to %s (%s)\n",
117				id, filename, strerror(errno));
118			ret = 1;
119			goto out;
120		}
121		n += i;
122	}
123	gettimeofday(&stop, NULL);
124	timersub(&stop, &start, &diff);
125	print_results(id + 1, OP_WRITE, data_size, &diff);
126
127	/* Read */
128	lseek(fd, 0, SEEK_SET);
129	n = 0;
130	gettimeofday(&start, NULL);
131	while (n < data_size) {
132		i = read(fd, buf, chunk_size);
133		if (i < 0) {
134			fprintf(stderr, "ERROR: task %d reading to %s (%s)\n",
135				id, filename, strerror(errno));
136			ret = 1;
137			goto out;
138		}
139		n += i;
140	}
141	gettimeofday(&stop, NULL);
142	timersub(&stop, &start, &diff);
143	print_results(id + 1, OP_READ, data_size, &diff);
144out:
145	close(fd);
146	unlink(filename);
147	free(buf);
148	exit(ret);
149}
150
151static void spawn(int id)
152{
153	pid_t pid;
154
155	pid = fork();
156	switch (pid) {
157	case -1:
158		fprintf(stderr, "ERROR: couldn't fork thread %d\n", id);
159		exit(1);
160	case 0:
161		thread(id);
162	default:
163		children[id] = pid;
164	}
165}
166
167void signal_handler(int sig)
168{
169	char filename[32];
170	int i;
171
172	for (i = 0; i < threads; i++)
173		if (children[i])
174			kill(children[i], SIGKILL);
175
176	for (i = 0; i < threads; i++) {
177		struct stat mystat;
178
179		snprintf(filename, sizeof(filename), "%s-%d-iobw.tmp",
180			 mygroup, i);
181		if (stat(filename, &mystat) < 0)
182			continue;
183		unlink(filename);
184	}
185
186	fprintf(stdout, "received signal %d, exiting\n", sig);
187	exit(0);
188}
189
190unsigned long long memparse(char *ptr, char **retptr)
191{
192	unsigned long long ret = strtoull(ptr, retptr, 0);
193
194	switch (**retptr) {
195	case 'G':
196	case 'g':
197		ret <<= 10;
198	case 'M':
199	case 'm':
200		ret <<= 10;
201	case 'K':
202	case 'k':
203		ret <<= 10;
204		(*retptr)++;
205	default:
206		break;
207	}
208	return ret;
209}
210
211int main(int argc, char *argv[])
212{
213	struct timeval start, stop, diff;
214	char *end;
215	int i;
216
217	if (argv[1] && strcmp(argv[1], "-direct") == 0) {
218		directio = 1;
219		argc--;
220		argv++;
221	}
222	if (argc != 4) {
223		fprintf(stderr, usage);
224		exit(1);
225	}
226	if ((threads = atoi(argv[1])) == 0) {
227		fprintf(stderr, usage);
228		exit(1);
229	}
230	chunk_size = align(memparse(argv[2], &end), PAGE_SIZE);
231	if (*end) {
232		fprintf(stderr, usage);
233		exit(1);
234	}
235	data_size = align(memparse(argv[3], &end), PAGE_SIZE);
236	if (*end) {
237		fprintf(stderr, usage);
238		exit(1);
239	}
240
241	/* retrieve group name */
242	mygroup = getenv("MYGROUP");
243	if (!mygroup) {
244		fprintf(stderr,
245			"ERROR: undefined environment variable MYGROUP\n");
246		exit(1);
247	}
248
249	children = malloc(sizeof(pid_t) * threads);
250	if (!children) {
251		fprintf(stderr, "ERROR: not enough memory\n");
252		exit(1);
253	}
254
255	/* handle user interrupt */
256	signal(SIGINT, signal_handler);
257	/* handle kill from shell */
258	signal(SIGTERM, signal_handler);
259
260	fprintf(stdout, "chunk_size %zuKiB, data_size %zuKiB\n",
261		kb(chunk_size), kb(data_size));
262	fflush(stdout);
263
264	gettimeofday(&start, NULL);
265	for (i = 0; i < threads; i++)
266		spawn(i);
267	for (i = 0; i < threads; i++) {
268		int status;
269		wait(&status);
270		if (!WIFEXITED(status))
271			exit(1);
272	}
273	gettimeofday(&stop, NULL);
274
275	timersub(&stop, &start, &diff);
276	print_results(0, NUM_IOPS, data_size * threads * NUM_IOPS, &diff);
277	fflush(stdout);
278	free(children);
279
280	exit(0);
281}
282