// SPDX-License-Identifier: GPL-2.0
/*
* This program reserves and uses hugetlb memory, supporting a bunch of
* scenarios needed by the charged_reserved_hugetlb.sh test.
*/
#include <err.h>
#include <errno.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/shm.h>
#include <sys/stat.h>
#include <sys/mman.h>
/* Global definitions. */
enum method {
HUGETLBFS,
MMAP_MAP_HUGETLB,
SHM,
MAX_METHOD
};
/* Global variables. */
static const char *self;
static int *shmaddr;
static int shmid;
/*
* Show usage and exit.
*/
static void exit_usage(void)
{
printf("Usage: %s -p <path to hugetlbfs file> -s <size to map> "
"[-m <0=hugetlbfs | 1=mmap(MAP_HUGETLB)>] [-l] [-r] "
"[-o] [-w] [-n]\n",
self);
exit(EXIT_FAILURE);
}
void sig_handler(int signo)
{
printf("Received %d.\n", signo);
if (signo == SIGINT) {
if (shmaddr) {
printf("Deleting the memory\n");
if (shmdt((const void *)shmaddr) != 0) {
perror("Detach failure");
shmctl(shmid, IPC_RMID, NULL);
exit(4);
}
shmctl(shmid, IPC_RMID, NULL);
printf("Done deleting the memory\n");
}
}
exit(2);
}
int main(int argc, char **argv)
{
int fd = 0;
int key = 0;
int *ptr = NULL;
int c = 0;
int size = 0;
char path[256] = "";
enum method method = MAX_METHOD;
int want_sleep = 0, private = 0;
int populate = 0;
int write = 0;
int reserve = 1;
if (signal(SIGINT, sig_handler) == SIG_ERR)
err(1, "\ncan't catch SIGINT\n");
/* Parse command-line arguments. */
setvbuf(stdout, NULL, _IONBF, 0);
self = argv[0];
while ((c = getopt(argc, argv, "s:p:m:owlrn")) != -1) {
switch (c) {
case 's':
size = atoi(optarg);
break;
case 'p':
strncpy(path, optarg, sizeof(path));
break;
case 'm':
if (atoi(optarg) >= MAX_METHOD) {
errno = EINVAL;
perror("Invalid -m.");
exit_usage();
}
method = atoi(optarg);
break;
case 'o':
populate = 1;
break;
case 'w':
write = 1;
break;
case 'l':
want_sleep = 1;
break;
case 'r':
private
= 1;
break;
case 'n':
reserve = 0;
break;
default:
errno = EINVAL;
perror("Invalid arg");
exit_usage();
}
}
if (strncmp(path, "", sizeof(path)) != 0) {
printf("Writing to this path: %s\n", path);
} else {
errno = EINVAL;
perror("path not found");
exit_usage();
}
if (size != 0) {
printf("Writing this size: %d\n", size);
} else {
errno = EINVAL;
perror("size not found");
exit_usage();
}
if (!populate)
printf("Not populating.\n");
else
printf("Populating.\n");
if (!write)
printf("Not writing to memory.\n");
if (method == MAX_METHOD) {
errno = EINVAL;
perror("-m Invalid");
exit_usage();
} else
printf("Using method=%d\n", method);
if (!private)
printf("Shared mapping.\n");
else
printf("Private mapping.\n");
if (!reserve)
printf("NO_RESERVE mapping.\n");
else
printf("RESERVE mapping.\n");
switch (method) {
case HUGETLBFS:
printf("Allocating using HUGETLBFS.\n");
fd = open(path, O_CREAT | O_RDWR, 0777);
if (fd == -1)
err(1, "Failed to open file.");
ptr = mmap(NULL, size, PROT_READ | PROT_WRITE,
(private ? MAP_PRIVATE : MAP_SHARED) |
(populate ? MAP_POPULATE : 0) |
(reserve ? 0 : MAP_NORESERVE),
fd, 0);
if (ptr == MAP_FAILED) {
close(fd);
err(1, "Error mapping the file");
}
break;
case MMAP_MAP_HUGETLB:
printf("Allocating using MAP_HUGETLB.\n");
ptr = mmap(NULL, size, PROT_READ | PROT_WRITE,
(private ? (MAP_PRIVATE | MAP_ANONYMOUS) :
MAP_SHARED) |
MAP_HUGETLB | (populate ? MAP_POPULATE : 0) |
(reserve ? 0 : MAP_NORESERVE),
-1, 0);
if (ptr == MAP_FAILED)
err(1, "mmap");
printf("Returned address is %p\n", ptr);
break;
case SHM:
printf("Allocating using SHM.\n");
shmid = shmget(key, size,
SHM_HUGETLB | IPC_CREAT | SHM_R | SHM_W);
if (shmid < 0) {
shmid = shmget(++key, size,
SHM_HUGETLB | IPC_CREAT | SHM_R | SHM_W);
if (shmid < 0)
err(1, "shmget");
}
printf("shmid: 0x%x, shmget key:%d\n", shmid, key);
ptr = shmat(shmid, NULL, 0);
if (ptr == (int *)-1) {
perror("Shared memory attach failure");
shmctl(shmid, IPC_RMID, NULL);
exit(2);
}
shmaddr = ptr;
printf("shmaddr: %p\n", shmaddr);
break;
default:
errno = EINVAL;
err(1, "Invalid method.");
}
if (write) {
printf("Writing to memory.\n");
memset(ptr, 1, size);
}
if (want_sleep) {
/* Signal to caller that we're done. */
printf("DONE\n");
/* Hold memory until external kill signal is delivered. */
while (1)
sleep(100);
}
if (method == HUGETLBFS)
close(fd);
return 0;
}