tools/testing/selftests/vm/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 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 }