tools/testing/selftests/mm/write_to_hugetlbfs.c

   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         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                     private
 113                         = 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         else
 144                 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         } else
 154                 printf("Using method=%d\n", method);
 155
 156         if (!private)
 157                 printf("Shared mapping.\n");
 158         else
 159                 printf("Private mapping.\n");
 160
 161         if (!reserve)
 162                 printf("NO_RESERVE mapping.\n");
 163         else
 164                 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 }