1/*
2 * Test program for memory error handling for hugepages
3 * Author: Naoya Horiguchi <n-horiguchi@ah.jp.nec.com>
4 */
5#define _GNU_SOURCE 1
6#include <stdlib.h>
7#include <stdio.h>
8#include <string.h>
9#include <fcntl.h>
10#include <signal.h>
11#include <unistd.h>
12#include <getopt.h>
13#include <sys/mman.h>
14#include <sys/ipc.h>
15#include <sys/shm.h>
16#include <sys/sem.h>
17#include <sys/types.h>
18#include <sys/prctl.h>
19#include <sys/wait.h>
20#include "hugepage.h"
21
22#define FILE_BASE  "test"
23
24#define HPAGE_SIZE (2UL*1024*1024)
25#define BUF_SIZE   256
26#define PROTECTION (PROT_READ | PROT_WRITE)
27
28#ifndef SHM_HUGETLB
29#define SHM_HUGETLB 04000
30#endif
31
32/* Control early_kill/late_kill */
33#define PR_MCE_KILL 33
34#define PR_MCE_KILL_CLEAR   0
35#define PR_MCE_KILL_SET     1
36#define PR_MCE_KILL_LATE    0
37#define PR_MCE_KILL_EARLY   1
38#define PR_MCE_KILL_DEFAULT 2
39#define PR_MCE_KILL_GET 34
40
41#define MADV_HWPOISON		100
42#define MADV_SOFT_OFFLINE	101
43
44int PS; /* Page size */
45int file_size; /* Memory allocation size (hugepage unit) */
46/* Error injection position (page offset from the first hugepage head) */
47int corrupt_page;
48char filename[BUF_SIZE] = "/test";
49char filepath[BUF_SIZE];
50
51#define DEB printf("DEBUG [%d:%s:%d]\n", getpid(), __FILE__, __LINE__);
52
53static void usage(void)
54{
55	printf(
56"./thugetlb [-m memory] [-o offset] [-f file] [-xOeSAaFpch] hugetlbfs_directory\n"
57"            -m|--memory size(hugepage unit)    Size of hugetlbfs file\n"
58"            -o|--offset offset(page unit)      Position of error injection\n"
59"            -x|--inject                        Error injection switch\n"
60"            -O|--offline                       Soft offline switch\n"
61"            -e|--early-kill                    Set PR_MCE_KILL_EARLY\n"
62"            -S|--shm                           Use shmem with SHM_HUGETLB\n"
63"            -A|--anonymous                     Use MAP_ANONYMOUS\n"
64"            -a|--avoid-touch                   Avoid touching error page\n"
65"            -F|--fork\n"
66"            -p|--private\n"
67"            -c|--cow\n"
68"            -f|--filename string\n"
69"            -h|--help\n"
70"\n"
71	);
72}
73
74/*
75 * semaphore get/put wrapper
76 */
77int get_semaphore(int sem_id, struct sembuf *sembuffer)
78{
79	sembuffer->sem_num = 0;
80	sembuffer->sem_op  = -1;
81	sembuffer->sem_flg = SEM_UNDO;
82	return semop(sem_id, sembuffer, 1);
83}
84
85int put_semaphore(int sem_id, struct sembuf *sembuffer)
86{
87	sembuffer->sem_num = 0;
88	sembuffer->sem_op  = 1;
89	sembuffer->sem_flg = SEM_UNDO;
90	return semop(sem_id, sembuffer, 1);
91}
92
93static struct option opts[] = {
94	{ "memory"          , 1, NULL, 'm' },
95	{ "offset"          , 1, NULL, 'o' },
96	{ "inject"          , 0, NULL, 'x' },
97	{ "offline"         , 0, NULL, 'O' },
98	{ "early_kill"      , 0, NULL, 'e' },
99	{ "shm"             , 0, NULL, 'S' },
100	{ "anonymous"       , 0, NULL, 'A' },
101	{ "avoid-touch"     , 0, NULL, 'a' },
102	{ "fork"            , 0, NULL, 'F' },
103	{ "private"         , 0, NULL, 'p' },
104	{ "cow"             , 0, NULL, 'c' },
105	{ "filename"        , 1, NULL, 'f' },
106	{ "help"            , 0, NULL, 'h' },
107	{ NULL              , 0, NULL,  0  }
108};
109
110int main(int argc, char *argv[])
111{
112	void *addr;
113	int i;
114	int ret;
115	int fd = 0;
116	int semid;
117	int semaphore;
118	int inject = 0;
119	int madvise_code = MADV_HWPOISON;
120	int early_kill = 0;
121	int avoid_touch = 0;
122	int anonflag = 0;
123	int shmflag = 0;
124	int shmkey = 0;
125	int forkflag = 0;
126	int privateflag = 0;
127	int cowflag = 0;
128	char c;
129	pid_t pid = 0;
130	void *expected_addr = NULL;
131	struct sembuf sembuffer;
132
133	PS = getpagesize();
134	HPS = HPAGE_SIZE;
135	file_size = 1;
136	corrupt_page = -1;
137
138	if (argc == 1) {
139		usage();
140		exit(EXIT_FAILURE);
141	}
142
143	while ((c = getopt_long(argc, argv,
144				"m:o:xOeSAaFpcf:h", opts, NULL)) != -1) {
145		switch (c) {
146		case 'm':
147			file_size = strtol(optarg, NULL, 10);
148			break;
149		case 'o':
150			corrupt_page = strtol(optarg, NULL, 10);
151			break;
152		case 'x':
153			inject = 1;
154			break;
155		case 'O':
156			madvise_code = MADV_SOFT_OFFLINE;
157			break;
158		case 'e':
159			early_kill = 1;
160			break;
161		case 'S':
162			shmflag = 1;
163			break;
164		case 'A':
165			anonflag = 1;
166			break;
167		case 'a':
168			avoid_touch = 1;
169			break;
170		case 'F':
171			forkflag = 1;
172			break;
173		case 'p':
174			privateflag = 1;
175			break;
176		case 'c':
177			cowflag = 1;
178			break;
179		case 'f':
180			strcat(filename, optarg);
181			shmkey = strtol(optarg, NULL, 10);
182			break;
183		case 'h':
184			usage();
185			exit(EXIT_SUCCESS);
186		default:
187			usage();
188			exit(EXIT_FAILURE);
189		}
190	}
191
192	if (inject && corrupt_page * PS > file_size * HPAGE_SIZE)
193		errmsg("Target page is out of range.\n");
194
195	if (avoid_touch && corrupt_page == -1)
196		errmsg("Avoid which page?\n");
197
198	/* Construct file name */
199	if (access(argv[argc - 1], F_OK) == -1) {
200		usage();
201		exit(EXIT_FAILURE);
202	} else {
203		strcpy(filepath, argv[argc - 1]);
204		strcat(filepath, filename);
205	}
206
207	if (shmflag) {
208		addr = alloc_shm_hugepage(&shmkey, file_size * HPAGE_SIZE);
209		if (!addr)
210			errmsg("Failed in alloc_shm_hugepage()");
211	} else if (anonflag) {
212		addr = alloc_anonymous_hugepage(file_size * HPAGE_SIZE,
213						privateflag);
214		if (!addr)
215			errmsg("Failed in alloc_anonymous_hugepage()");
216	} else {
217		addr = alloc_filebacked_hugepage(filepath,
218						 file_size * HPAGE_SIZE,
219						 privateflag, &fd);
220		if (!addr)
221			errmsg("Failed in alloc_filebacked_hugepage()");
222	}
223
224	if (corrupt_page != -1 && avoid_touch)
225		expected_addr = (void *)(addr + corrupt_page / 512 * HPAGE_SIZE);
226
227	if (forkflag) {
228		semid = semget(IPC_PRIVATE, 1, 0666|IPC_CREAT);
229		if (semid == -1) {
230			perror("semget");
231			goto cleanout;
232		}
233		semaphore = semctl(semid, 0, SETVAL, 1);
234		if (semaphore == -1) {
235			perror("semctl");
236			goto cleanout;
237		}
238		if (get_semaphore(semid, &sembuffer)) {
239			perror("get_semaphore");
240			goto cleanout;
241		}
242	}
243
244	write_hugepage(addr, file_size, 0);
245	read_hugepage(addr, file_size, 0);
246
247	if (early_kill)
248		prctl(PR_MCE_KILL, PR_MCE_KILL_SET, PR_MCE_KILL_EARLY,
249		      NULL, NULL);
250
251	/*
252	 * Intended order:
253	 *   1. Child COWs
254	 *   2. Parent madvise()s
255	 *   3. Child exit()s
256	 */
257	if (forkflag) {
258		pid = fork();
259		if (!pid) {
260			/* Semaphore is already held */
261			if (cowflag) {
262				write_hugepage(addr, file_size, 0);
263				read_hugepage(addr, file_size, 0);
264			}
265			if (put_semaphore(semid, &sembuffer))
266				err("put_semaphore");
267			usleep(1000);
268			/* Wait for madvise() to be done */
269			if (get_semaphore(semid, &sembuffer))
270				err("put_semaphore");
271			if (put_semaphore(semid, &sembuffer))
272				err("put_semaphore");
273			return 0;
274		}
275	}
276
277	/* Wait for COW */
278	if (forkflag && get_semaphore(semid, &sembuffer)) {
279		perror("get_semaphore");
280		goto cleanout;
281	}
282
283	if (inject && corrupt_page != -1) {
284		ret = madvise(addr + corrupt_page * PS, PS, madvise_code);
285		if (ret) {
286			printf("madivise return %d :", ret);
287			perror("madvise");
288			goto cleanout;
289		}
290	}
291
292	if (forkflag && put_semaphore(semid, &sembuffer)) {
293		perror("put_semaphore");
294		goto cleanout;
295	}
296
297	if (madvise_code != MADV_SOFT_OFFLINE);
298		write_hugepage(addr, file_size, expected_addr);
299	read_hugepage(addr, file_size, expected_addr);
300
301	if (forkflag) {
302		if (wait(&i) == -1)
303			err("wait");
304		if (semctl(semid, 0, IPC_RMID) == -1)
305			err("semctl(IPC_RMID)");
306	}
307cleanout:
308	if (shmflag) {
309		if (free_shm_hugepage(shmkey, addr) == -1)
310			exit(2);
311	} else if (anonflag) {
312		if (free_anonymous_hugepage(addr, file_size * HPAGE_SIZE) == -1)
313			exit(2);
314	} else {
315		if (free_filebacked_hugepage(addr, file_size * HPAGE_SIZE,
316					     fd, filepath) == -1)
317			exit(2);
318	}
319
320	return 0;
321}
322