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