memcg_process.c revision 6c4d1c6d4b3c116ff88f9209b4860940fb8659a8 
1/******************************************************************************/
2/*                                                                            */
3/* Copyright (c) 2009 FUJITSU LIMITED                                         */
4/*                                                                            */
5/* This program is free software;  you can redistribute it and/or modify      */
6/* it under the terms of the GNU General Public License as published by       */
7/* the Free Software Foundation; either version 2 of the License, or          */
8/* (at your option) any later version.                                        */
9/*                                                                            */
10/* This program is distributed in the hope that it will be useful,            */
11/* but WITHOUT ANY WARRANTY;  without even the implied warranty of            */
12/* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See                  */
13/* the GNU General Public License for more details.                           */
14/*                                                                            */
15/* You should have received a copy of the GNU General Public License          */
16/* along with this program;  if not, write to the Free Software               */
17/* Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA    */
18/*                                                                            */
19/* Author: Li Zefan <lizf@cn.fujitsu.com>                                     */
20/*                                                                            */
21/******************************************************************************/
22
23#include <sys/types.h>
24#include <sys/mman.h>
25#include <sys/shm.h>
26#include <sys/stat.h>
27#include <err.h>
28#include <errno.h>
29#include <fcntl.h>
30#include <getopt.h>
31#include <limits.h>
32#include <signal.h>
33#include <stdio.h>
34#include <stdlib.h>
35#include <string.h>
36#include <unistd.h>
37
38int fd;
39
40int flag_exit;
41int flag_allocated;
42
43int opt_mmap_anon;
44int opt_mmap_file;
45int opt_mmap_lock1;
46int opt_mmap_lock2;
47int opt_shm;
48int opt_hugepage;
49
50int key_id;		/* used with opt_shm */
51unsigned long memsize;
52
53#define FILE_HUGEPAGE	"/hugetlb/hugepagefile"
54
55#define MMAP_ANON	(SCHAR_MAX + 1)
56#define MMAP_FILE	(SCHAR_MAX + 2)
57#define MMAP_LOCK1	(SCHAR_MAX + 3)
58#define MMAP_LOCK2	(SCHAR_MAX + 4)
59#define SHM		(SCHAR_MAX + 5)
60#define HUGEPAGE	(SCHAR_MAX + 6)
61
62const struct option long_opts[] = {
63	{ "mmap-anon",	0, NULL, MMAP_ANON	},
64	{ "mmap-file",	0, NULL, MMAP_FILE	},
65	{ "mmap-lock1",	0, NULL, MMAP_LOCK1	},
66	{ "mmap-lock2",	0, NULL, MMAP_LOCK2	},
67	{ "shm",	0, NULL, SHM		},
68	{ "hugepage",	0, NULL, HUGEPAGE	},
69	{ "size",	1, NULL, 's'		},
70	{ "key",	1, NULL, 'k'		},
71	{ NULL,		0, NULL, 0		},
72};
73
74/*
75 * process_options: read options from user input
76 */
77void process_options(int argc, char *argv[])
78{
79	int c;
80	char *end;
81
82	while ((c = getopt_long(argc, argv, "k:s:", long_opts, NULL)) != -1) {
83		switch (c) {
84		case 'k':
85			key_id = atoi(optarg);
86			break;
87		case 's':
88			memsize = strtoul(optarg, &end, 10);
89			if (*end != '\0')
90				errx(1, "wrong -s argument!");
91			break;
92		case MMAP_ANON:
93			opt_mmap_anon = 1;
94			break;
95		case MMAP_FILE:
96			opt_mmap_file = 1;
97			break;
98		case MMAP_LOCK1:
99			opt_mmap_lock1 = 1;
100			break;
101		case MMAP_LOCK2:
102			opt_mmap_lock2 = 1;
103			break;
104		case SHM:
105			opt_shm = 1;
106			break;
107		case HUGEPAGE:
108			opt_hugepage = 1;
109			break;
110		default:
111			errx(1, "unknown option: %c", c);
112			break;
113		}
114	}
115}
116
117/*
118 * touch_memory: force allocating phy memory
119 */
120void touch_memory(char *p, int size)
121{
122	int i;
123	int pagesize = getpagesize();
124
125	for (i = 0; i < size; i += pagesize)
126		p[i] = 0xef;
127}
128
129void mmap_anon()
130{
131	static char *p;
132
133	if (!flag_allocated) {
134		p = mmap(NULL, memsize, PROT_WRITE | PROT_READ,
135			 MAP_PRIVATE|MAP_ANONYMOUS, 0, 0);
136		if (p == MAP_FAILED)
137			err(1, "mmap(anonymous) failed");
138		touch_memory(p, memsize);
139	} else {
140		if (munmap(p, memsize) == -1)
141			err(1, "munmap(anonymous) failed");
142	}
143}
144
145void mmap_file()
146{
147	static char *p;
148	static int fd_hugepage;
149	int fd_tmp;
150
151	if (!flag_allocated) {
152		if (opt_hugepage) {
153			fd_hugepage = open(FILE_HUGEPAGE,
154					   O_CREAT | O_RDWR, 0755);
155			if (fd_hugepage < 0)
156				err(1, "open hugepage file failed");
157			fd_tmp = fd_hugepage;
158		} else
159			fd_tmp = fd;
160
161		p = mmap(NULL, memsize, PROT_WRITE | PROT_READ,
162			 MAP_SHARED, fd_tmp, 0);
163		if (p == MAP_FAILED) {
164			if (opt_hugepage)
165				unlink(FILE_HUGEPAGE);
166			err(1, "mmap(file) failed");
167		}
168		touch_memory(p, memsize);
169	} else {
170		if (!munmap(p, memsize) == -1)
171			err(1, "munmap(file) failed");
172
173		if (opt_hugepage) {
174			close(fd_hugepage);
175			unlink(FILE_HUGEPAGE);
176		}
177	}
178}
179
180void mmap_lock1()
181{
182	static char *p;
183
184	if (!flag_allocated) {
185		p = mmap(NULL, memsize, PROT_WRITE | PROT_READ,
186			 MAP_PRIVATE | MAP_ANONYMOUS | MAP_LOCKED, 0, 0);
187		if (p == MAP_FAILED)
188			err(1, "mmap(lock) failed");
189	} else {
190		if (munmap(p, memsize) == -1)
191			err(1, "munmap(lock) failed");
192	}
193}
194
195void mmap_lock2()
196{
197	static char *p;
198
199	if (!flag_allocated) {
200		p = mmap(NULL, memsize, PROT_WRITE | PROT_READ,
201			 MAP_PRIVATE | MAP_ANONYMOUS, 0, 0);
202		if (p == MAP_FAILED)
203			err(1, "mmap failed");
204
205		if (!mlock(p, memsize))
206			err(1, "mlock failed");
207	} else {
208		if (!munmap(p, memsize) == -1)
209			err(1, "munmap failed");
210	}
211}
212
213void shm()
214{
215	static char *p;
216	static int shmid;
217	unsigned long flag;
218
219	key_t key;
220
221	if (!flag_allocated) {
222		flag = IPC_CREAT | SHM_R | SHM_W;
223		if (opt_hugepage)
224			flag |= SHM_HUGETLB;
225
226		key = ftok("/dev/null", key_id);
227		if (key == -1)
228			err(1, "ftok() failed");
229
230		shmid = shmget(key, memsize, flag);
231		if (shmid == -1)
232			err(1, "shmget() failed");
233		shmctl(shmid, IPC_RMID, NULL);
234
235		shmid = shmget(key, memsize, flag);
236		if (shmid == -1)
237			err(1, "shmget() failed");
238
239		p = shmat(shmid, NULL, 0);
240		if (p == (void *)-1) {
241			shmctl(shmid, IPC_RMID, NULL);
242			err(1, "shmat() failed");
243		}
244		touch_memory(p, memsize);
245	} else {
246		if (shmdt(p) == -1)
247			err(1, "shmdt() failed");
248		if (shmctl(shmid, IPC_RMID, NULL) == -1)
249			err(1, "shmctl() failed");
250	}
251}
252
253/*
254 * sigint_handler: handle SIGINT by set the exit flag.
255 */
256void sigint_handler(int __attribute__((unused)) signo)
257{
258	flag_exit = 1;
259}
260
261/*
262 * sigusr_handler: handler SIGUSR
263 *
264 * When we receive SIGUSR, we allocate some memory according
265 * to the user input when the process started.
266 *
267 * When we receive SIGUSR again, we will free all the allocated
268 * memory.
269 */
270void sigusr_handler(int __attribute__((unused)) signo)
271{
272	if (opt_mmap_anon)
273		mmap_anon();
274
275	if (opt_mmap_file)
276		mmap_file();
277
278	if (opt_mmap_lock1)
279		mmap_lock1();
280
281	if (opt_mmap_lock2)
282		mmap_lock2();
283
284	if (opt_shm)
285		shm();
286
287	flag_allocated = !flag_allocated;
288}
289
290int main(int argc, char *argv[])
291{
292	struct sigaction sigint_action;
293	struct sigaction sigusr_action;
294
295	if ((fd = open("/dev/zero", O_RDWR)) == -1)
296		err(1, "open /dev/zero failed");
297
298	/* TODO: writer error handling here. */
299	sigemptyset(&sigint_action.sa_mask);
300	sigint_action.sa_handler = &sigint_handler;
301	sigaction(SIGINT, &sigint_action, NULL);
302
303	sigemptyset(&sigint_action.sa_mask);
304	sigusr_action.sa_handler = &sigusr_handler;
305	sigaction(SIGUSR1, &sigusr_action, NULL);
306
307	sigemptyset(&sigusr_action.sa_mask);
308
309	process_options(argc, argv);
310
311	while (!flag_exit)
312		sleep(1);
313
314	close(fd);
315
316	return 0;
317}
318