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 | */ |
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 | |
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 | |
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 | |