244 lines · c
1// SPDX-License-Identifier: GPL-2.02/*3 * This program reserves and uses hugetlb memory, supporting a bunch of4 * 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. */21enum method {22 HUGETLBFS,23 MMAP_MAP_HUGETLB,24 SHM,25 MAX_METHOD26};27 28 29/* Global variables. */30static const char *self;31static int *shmaddr;32static int shmid;33 34/*35 * Show usage and exit.36 */37static 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 46void 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 65int 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 private113 = 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 else144 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 } else154 printf("Using method=%d\n", method);155 156 if (!private)157 printf("Shared mapping.\n");158 else159 printf("Private mapping.\n");160 161 if (!reserve)162 printf("NO_RESERVE mapping.\n");163 else164 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