1 // SPDX-License-Identifier: GPL-2.0
2 /*
3 * This program reserves and uses hugetlb memory, supporting a bunch of
4 * scenarios needed by the charged_reserved_hugetlb.sh test.
5 */
6
7 #include <err.h>
8 #include <errno.h>
9 #include <signal.h>
10 #include <stdio.h>
11 #include <stdlib.h>
12 #include <string.h>
13 #include <unistd.h>
14 #include <fcntl.h>
15 #include <sys/types.h>
16 #include <sys/shm.h>
17 #include <sys/stat.h>
18 #include <sys/mman.h>
19
20 /* Global definitions. */
21 enum method {
22 HUGETLBFS,
23 MMAP_MAP_HUGETLB,
24 SHM,
25 MAX_METHOD
26 };
27
28
29 /* Global variables. */
30 static const char *self;
31 static int *shmaddr;
32 static int shmid;
33
34 /*
35 * Show usage and exit.
36 */
exit_usage(void)37 static void exit_usage(void)
38 {
39 printf("Usage: %s -p <path to hugetlbfs file> -s <size to map> "
40 "[-m <0=hugetlbfs | 1=mmap(MAP_HUGETLB)>] [-l] [-r] "
41 "[-o] [-w] [-n]\n",
42 self);
43 exit(EXIT_FAILURE);
44 }
45
sig_handler(int signo)46 void sig_handler(int signo)
47 {
48 printf("Received %d.\n", signo);
49 if (signo == SIGINT) {
50 if (shmaddr) {
51 printf("Deleting the memory\n");
52 if (shmdt((const void *)shmaddr) != 0) {
53 perror("Detach failure");
54 shmctl(shmid, IPC_RMID, NULL);
55 exit(4);
56 }
57
58 shmctl(shmid, IPC_RMID, NULL);
59 printf("Done deleting the memory\n");
60 }
61 }
62 exit(2);
63 }
64
main(int argc,char ** argv)65 int main(int argc, char **argv)
66 {
67 int fd = 0;
68 int key = 0;
69 int *ptr = NULL;
70 int c = 0;
71 int size = 0;
72 char path[256] = "";
73 enum method method = MAX_METHOD;
74 int want_sleep = 0, private = 0;
75 int populate = 0;
76 int write = 0;
77 int reserve = 1;
78
79 if (signal(SIGINT, sig_handler) == SIG_ERR)
80 err(1, "\ncan't catch SIGINT\n");
81
82 /* Parse command-line arguments. */
83 setvbuf(stdout, NULL, _IONBF, 0);
84 self = argv[0];
85
86 while ((c = getopt(argc, argv, "s:p:m:owlrn")) != -1) {
87 switch (c) {
88 case 's':
89 size = atoi(optarg);
90 break;
91 case 'p':
92 strncpy(path, optarg, sizeof(path));
93 break;
94 case 'm':
95 if (atoi(optarg) >= MAX_METHOD) {
96 errno = EINVAL;
97 perror("Invalid -m.");
98 exit_usage();
99 }
100 method = atoi(optarg);
101 break;
102 case 'o':
103 populate = 1;
104 break;
105 case 'w':
106 write = 1;
107 break;
108 case 'l':
109 want_sleep = 1;
110 break;
111 case 'r':
112 private
113 = 1;
114 break;
115 case 'n':
116 reserve = 0;
117 break;
118 default:
119 errno = EINVAL;
120 perror("Invalid arg");
121 exit_usage();
122 }
123 }
124
125 if (strncmp(path, "", sizeof(path)) != 0) {
126 printf("Writing to this path: %s\n", path);
127 } else {
128 errno = EINVAL;
129 perror("path not found");
130 exit_usage();
131 }
132
133 if (size != 0) {
134 printf("Writing this size: %d\n", size);
135 } else {
136 errno = EINVAL;
137 perror("size not found");
138 exit_usage();
139 }
140
141 if (!populate)
142 printf("Not populating.\n");
143 else
144 printf("Populating.\n");
145
146 if (!write)
147 printf("Not writing to memory.\n");
148
149 if (method == MAX_METHOD) {
150 errno = EINVAL;
151 perror("-m Invalid");
152 exit_usage();
153 } else
154 printf("Using method=%d\n", method);
155
156 if (!private)
157 printf("Shared mapping.\n");
158 else
159 printf("Private mapping.\n");
160
161 if (!reserve)
162 printf("NO_RESERVE mapping.\n");
163 else
164 printf("RESERVE mapping.\n");
165
166 switch (method) {
167 case HUGETLBFS:
168 printf("Allocating using HUGETLBFS.\n");
169 fd = open(path, O_CREAT | O_RDWR, 0777);
170 if (fd == -1)
171 err(1, "Failed to open file.");
172
173 ptr = mmap(NULL, size, PROT_READ | PROT_WRITE,
174 (private ? MAP_PRIVATE : MAP_SHARED) |
175 (populate ? MAP_POPULATE : 0) |
176 (reserve ? 0 : MAP_NORESERVE),
177 fd, 0);
178
179 if (ptr == MAP_FAILED) {
180 close(fd);
181 err(1, "Error mapping the file");
182 }
183 break;
184 case MMAP_MAP_HUGETLB:
185 printf("Allocating using MAP_HUGETLB.\n");
186 ptr = mmap(NULL, size, PROT_READ | PROT_WRITE,
187 (private ? (MAP_PRIVATE | MAP_ANONYMOUS) :
188 MAP_SHARED) |
189 MAP_HUGETLB | (populate ? MAP_POPULATE : 0) |
190 (reserve ? 0 : MAP_NORESERVE),
191 -1, 0);
192
193 if (ptr == MAP_FAILED)
194 err(1, "mmap");
195
196 printf("Returned address is %p\n", ptr);
197 break;
198 case SHM:
199 printf("Allocating using SHM.\n");
200 shmid = shmget(key, size,
201 SHM_HUGETLB | IPC_CREAT | SHM_R | SHM_W);
202 if (shmid < 0) {
203 shmid = shmget(++key, size,
204 SHM_HUGETLB | IPC_CREAT | SHM_R | SHM_W);
205 if (shmid < 0)
206 err(1, "shmget");
207 }
208 printf("shmid: 0x%x, shmget key:%d\n", shmid, key);
209
210 ptr = shmat(shmid, NULL, 0);
211 if (ptr == (int *)-1) {
212 perror("Shared memory attach failure");
213 shmctl(shmid, IPC_RMID, NULL);
214 exit(2);
215 }
216 shmaddr = ptr;
217 printf("shmaddr: %p\n", shmaddr);
218
219 break;
220 default:
221 errno = EINVAL;
222 err(1, "Invalid method.");
223 }
224
225 if (write) {
226 printf("Writing to memory.\n");
227 memset(ptr, 1, size);
228 }
229
230 if (want_sleep) {
231 /* Signal to caller that we're done. */
232 printf("DONE\n");
233
234 /* Hold memory until external kill signal is delivered. */
235 while (1)
236 sleep(100);
237 }
238
239 if (method == HUGETLBFS)
240 close(fd);
241
242 return 0;
243 }
244