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 */ 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 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 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 size_t 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 if (sscanf(optarg, "%zu", &size) != 1) { 90 perror("Invalid -s."); 91 exit_usage(); 92 } 93 break; 94 case 'p': 95 strncpy(path, optarg, sizeof(path) - 1); 96 break; 97 case 'm': 98 if (atoi(optarg) >= MAX_METHOD) { 99 errno = EINVAL; 100 perror("Invalid -m."); 101 exit_usage(); 102 } 103 method = atoi(optarg); 104 break; 105 case 'o': 106 populate = 1; 107 break; 108 case 'w': 109 write = 1; 110 break; 111 case 'l': 112 want_sleep = 1; 113 break; 114 case 'r': 115 private 116 = 1; 117 break; 118 case 'n': 119 reserve = 0; 120 break; 121 default: 122 errno = EINVAL; 123 perror("Invalid arg"); 124 exit_usage(); 125 } 126 } 127 128 if (strncmp(path, "", sizeof(path)) != 0) { 129 printf("Writing to this path: %s\n", path); 130 } else { 131 errno = EINVAL; 132 perror("path not found"); 133 exit_usage(); 134 } 135 136 if (size != 0) { 137 printf("Writing this size: %zu\n", size); 138 } else { 139 errno = EINVAL; 140 perror("size not found"); 141 exit_usage(); 142 } 143 144 if (!populate) 145 printf("Not populating.\n"); 146 else 147 printf("Populating.\n"); 148 149 if (!write) 150 printf("Not writing to memory.\n"); 151 152 if (method == MAX_METHOD) { 153 errno = EINVAL; 154 perror("-m Invalid"); 155 exit_usage(); 156 } else 157 printf("Using method=%d\n", method); 158 159 if (!private) 160 printf("Shared mapping.\n"); 161 else 162 printf("Private mapping.\n"); 163 164 if (!reserve) 165 printf("NO_RESERVE mapping.\n"); 166 else 167 printf("RESERVE mapping.\n"); 168 169 switch (method) { 170 case HUGETLBFS: 171 printf("Allocating using HUGETLBFS.\n"); 172 fd = open(path, O_CREAT | O_RDWR, 0777); 173 if (fd == -1) 174 err(1, "Failed to open file."); 175 176 ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, 177 (private ? MAP_PRIVATE : MAP_SHARED) | 178 (populate ? MAP_POPULATE : 0) | 179 (reserve ? 0 : MAP_NORESERVE), 180 fd, 0); 181 182 if (ptr == MAP_FAILED) { 183 close(fd); 184 err(1, "Error mapping the file"); 185 } 186 break; 187 case MMAP_MAP_HUGETLB: 188 printf("Allocating using MAP_HUGETLB.\n"); 189 ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, 190 (private ? (MAP_PRIVATE | MAP_ANONYMOUS) : 191 MAP_SHARED) | 192 MAP_HUGETLB | (populate ? MAP_POPULATE : 0) | 193 (reserve ? 0 : MAP_NORESERVE), 194 -1, 0); 195 196 if (ptr == MAP_FAILED) 197 err(1, "mmap"); 198 199 printf("Returned address is %p\n", ptr); 200 break; 201 case SHM: 202 printf("Allocating using SHM.\n"); 203 shmid = shmget(key, size, 204 SHM_HUGETLB | IPC_CREAT | SHM_R | SHM_W); 205 if (shmid < 0) { 206 shmid = shmget(++key, size, 207 SHM_HUGETLB | IPC_CREAT | SHM_R | SHM_W); 208 if (shmid < 0) 209 err(1, "shmget"); 210 } 211 printf("shmid: 0x%x, shmget key:%d\n", shmid, key); 212 213 ptr = shmat(shmid, NULL, 0); 214 if (ptr == (int *)-1) { 215 perror("Shared memory attach failure"); 216 shmctl(shmid, IPC_RMID, NULL); 217 exit(2); 218 } 219 shmaddr = ptr; 220 printf("shmaddr: %p\n", shmaddr); 221 222 break; 223 default: 224 errno = EINVAL; 225 err(1, "Invalid method."); 226 } 227 228 if (write) { 229 printf("Writing to memory.\n"); 230 memset(ptr, 1, size); 231 } 232 233 if (want_sleep) { 234 /* Signal to caller that we're done. */ 235 printf("DONE\n"); 236 237 /* Hold memory until external kill signal is delivered. */ 238 while (1) 239 sleep(100); 240 } 241 242 if (method == HUGETLBFS) 243 close(fd); 244 245 return 0; 246 } 247