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 char *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 printf("Deleting the memory\n");
51 if (shmdt((const void *)shmaddr) != 0) {
52 perror("Detach failure");
53 shmctl(shmid, IPC_RMID, NULL);
54 exit(4);
55 }
56
57 shmctl(shmid, IPC_RMID, NULL);
58 printf("Done deleting the memory\n");
59 }
60 exit(2);
61 }
62
main(int argc,char ** argv)63 int main(int argc, char **argv)
64 {
65 int fd = 0;
66 int key = 0;
67 int *ptr = NULL;
68 int c = 0;
69 int size = 0;
70 char path[256] = "";
71 enum method method = MAX_METHOD;
72 int want_sleep = 0, private = 0;
73 int populate = 0;
74 int write = 0;
75 int reserve = 1;
76
77 if (signal(SIGINT, sig_handler) == SIG_ERR)
78 err(1, "\ncan't catch SIGINT\n");
79
80 /* Parse command-line arguments. */
81 setvbuf(stdout, NULL, _IONBF, 0);
82 self = argv[0];
83
84 while ((c = getopt(argc, argv, "s:p:m:owlrn")) != -1) {
85 switch (c) {
86 case 's':
87 size = atoi(optarg);
88 break;
89 case 'p':
90 strncpy(path, optarg, sizeof(path));
91 break;
92 case 'm':
93 if (atoi(optarg) >= MAX_METHOD) {
94 errno = EINVAL;
95 perror("Invalid -m.");
96 exit_usage();
97 }
98 method = atoi(optarg);
99 break;
100 case 'o':
101 populate = 1;
102 break;
103 case 'w':
104 write = 1;
105 break;
106 case 'l':
107 want_sleep = 1;
108 break;
109 case 'r':
110 private
111 = 1;
112 break;
113 case 'n':
114 reserve = 0;
115 break;
116 default:
117 errno = EINVAL;
118 perror("Invalid arg");
119 exit_usage();
120 }
121 }
122
123 if (strncmp(path, "", sizeof(path)) != 0) {
124 printf("Writing to this path: %s\n", path);
125 } else {
126 errno = EINVAL;
127 perror("path not found");
128 exit_usage();
129 }
130
131 if (size != 0) {
132 printf("Writing this size: %d\n", size);
133 } else {
134 errno = EINVAL;
135 perror("size not found");
136 exit_usage();
137 }
138
139 if (!populate)
140 printf("Not populating.\n");
141 else
142 printf("Populating.\n");
143
144 if (!write)
145 printf("Not writing to memory.\n");
146
147 if (method == MAX_METHOD) {
148 errno = EINVAL;
149 perror("-m Invalid");
150 exit_usage();
151 } else
152 printf("Using method=%d\n", method);
153
154 if (!private)
155 printf("Shared mapping.\n");
156 else
157 printf("Private mapping.\n");
158
159 if (!reserve)
160 printf("NO_RESERVE mapping.\n");
161 else
162 printf("RESERVE mapping.\n");
163
164 switch (method) {
165 case HUGETLBFS:
166 printf("Allocating using HUGETLBFS.\n");
167 fd = open(path, O_CREAT | O_RDWR, 0777);
168 if (fd == -1)
169 err(1, "Failed to open file.");
170
171 ptr = mmap(NULL, size, PROT_READ | PROT_WRITE,
172 (private ? MAP_PRIVATE : MAP_SHARED) |
173 (populate ? MAP_POPULATE : 0) |
174 (reserve ? 0 : MAP_NORESERVE),
175 fd, 0);
176
177 if (ptr == MAP_FAILED) {
178 close(fd);
179 err(1, "Error mapping the file");
180 }
181 break;
182 case MMAP_MAP_HUGETLB:
183 printf("Allocating using MAP_HUGETLB.\n");
184 ptr = mmap(NULL, size, PROT_READ | PROT_WRITE,
185 (private ? (MAP_PRIVATE | MAP_ANONYMOUS) :
186 MAP_SHARED) |
187 MAP_HUGETLB | (populate ? MAP_POPULATE : 0) |
188 (reserve ? 0 : MAP_NORESERVE),
189 -1, 0);
190
191 if (ptr == MAP_FAILED)
192 err(1, "mmap");
193
194 printf("Returned address is %p\n", ptr);
195 break;
196 case SHM:
197 printf("Allocating using SHM.\n");
198 shmid = shmget(key, size,
199 SHM_HUGETLB | IPC_CREAT | SHM_R | SHM_W);
200 if (shmid < 0) {
201 shmid = shmget(++key, size,
202 SHM_HUGETLB | IPC_CREAT | SHM_R | SHM_W);
203 if (shmid < 0)
204 err(1, "shmget");
205 }
206 printf("shmid: 0x%x, shmget key:%d\n", shmid, key);
207
208 ptr = shmat(shmid, NULL, 0);
209 if (ptr == (int *)-1) {
210 perror("Shared memory attach failure");
211 shmctl(shmid, IPC_RMID, NULL);
212 exit(2);
213 }
214 printf("shmaddr: %p\n", ptr);
215
216 break;
217 default:
218 errno = EINVAL;
219 err(1, "Invalid method.");
220 }
221
222 if (write) {
223 printf("Writing to memory.\n");
224 memset(ptr, 1, size);
225 }
226
227 if (want_sleep) {
228 /* Signal to caller that we're done. */
229 printf("DONE\n");
230
231 /* Hold memory until external kill signal is delivered. */
232 while (1)
233 sleep(100);
234 }
235
236 if (method == HUGETLBFS)
237 close(fd);
238
239 return 0;
240 }
241