// SPDX-License-Identifier: GPL-2.0
#include <dirent.h>
#include <fcntl.h>
#include <limits.h>
#include <signal.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "vm_util.h"
#include "hugepage_settings.h"

#define THP_SYSFS "/sys/kernel/mm/transparent_hugepage/"
#define MAX_SETTINGS_DEPTH 4
static struct thp_settings settings_stack[MAX_SETTINGS_DEPTH];
static int settings_index;
static struct thp_settings saved_settings;
static char dev_queue_read_ahead_path[PATH_MAX];
static bool thp_settings_saved;

static const char * const thp_enabled_strings[] = {
	"never",
	"always",
	"inherit",
	"madvise",
	NULL
};

static const char * const thp_defrag_strings[] = {
	"always",
	"defer",
	"defer+madvise",
	"madvise",
	"never",
	NULL
};

static const char * const shmem_enabled_strings[] = {
	"never",
	"always",
	"within_size",
	"advise",
	"inherit",
	"deny",
	"force",
	NULL
};

int thp_read_string(const char *name, const char * const strings[])
{
	char path[PATH_MAX];
	char buf[256];
	char *c;
	int ret;

	ret = snprintf(path, PATH_MAX, THP_SYSFS "%s", name);
	if (ret >= PATH_MAX) {
		printf("%s: Pathname is too long\n", __func__);
		exit(EXIT_FAILURE);
	}

	if (!read_file(path, buf, sizeof(buf))) {
		perror(path);
		exit(EXIT_FAILURE);
	}

	c = strchr(buf, '[');
	if (!c) {
		printf("%s: Parse failure\n", __func__);
		exit(EXIT_FAILURE);
	}

	c++;
	memmove(buf, c, sizeof(buf) - (c - buf));

	c = strchr(buf, ']');
	if (!c) {
		printf("%s: Parse failure\n", __func__);
		exit(EXIT_FAILURE);
	}
	*c = '\0';

	ret = 0;
	while (strings[ret]) {
		if (!strcmp(strings[ret], buf))
			return ret;
		ret++;
	}

	printf("Failed to parse %s\n", name);
	exit(EXIT_FAILURE);
}

void thp_write_string(const char *name, const char *val)
{
	char path[PATH_MAX];
	int ret;

	ret = snprintf(path, PATH_MAX, THP_SYSFS "%s", name);
	if (ret >= PATH_MAX) {
		printf("%s: Pathname is too long\n", __func__);
		exit(EXIT_FAILURE);
	}
	write_file(path, val, strlen(val) + 1);
}

unsigned long thp_read_num(const char *name)
{
	char path[PATH_MAX];
	int ret;

	ret = snprintf(path, PATH_MAX, THP_SYSFS "%s", name);
	if (ret >= PATH_MAX) {
		printf("%s: Pathname is too long\n", __func__);
		exit(EXIT_FAILURE);
	}
	return read_num(path);
}

void thp_write_num(const char *name, unsigned long num)
{
	char path[PATH_MAX];
	int ret;

	ret = snprintf(path, PATH_MAX, THP_SYSFS "%s", name);
	if (ret >= PATH_MAX) {
		printf("%s: Pathname is too long\n", __func__);
		exit(EXIT_FAILURE);
	}
	write_num(path, num);
}

void thp_read_settings(struct thp_settings *settings)
{
	unsigned long orders = thp_supported_orders();
	unsigned long shmem_orders = thp_shmem_supported_orders();
	char path[PATH_MAX];
	int i;

	*settings = (struct thp_settings) {
		.thp_enabled = thp_read_string("enabled", thp_enabled_strings),
		.thp_defrag = thp_read_string("defrag", thp_defrag_strings),
		.shmem_enabled =
			thp_read_string("shmem_enabled", shmem_enabled_strings),
		.use_zero_page = thp_read_num("use_zero_page"),
	};
	settings->khugepaged = (struct khugepaged_settings) {
		.defrag = thp_read_num("khugepaged/defrag"),
		.alloc_sleep_millisecs =
			thp_read_num("khugepaged/alloc_sleep_millisecs"),
		.scan_sleep_millisecs =
			thp_read_num("khugepaged/scan_sleep_millisecs"),
		.max_ptes_none = thp_read_num("khugepaged/max_ptes_none"),
		.max_ptes_swap = thp_read_num("khugepaged/max_ptes_swap"),
		.max_ptes_shared = thp_read_num("khugepaged/max_ptes_shared"),
		.pages_to_scan = thp_read_num("khugepaged/pages_to_scan"),
	};
	if (dev_queue_read_ahead_path[0])
		settings->read_ahead_kb = read_num(dev_queue_read_ahead_path);

	for (i = 0; i < NR_ORDERS; i++) {
		if (!((1 << i) & orders)) {
			settings->hugepages[i].enabled = THP_NEVER;
			continue;
		}
		snprintf(path, PATH_MAX, "hugepages-%ukB/enabled",
			(getpagesize() >> 10) << i);
		settings->hugepages[i].enabled =
			thp_read_string(path, thp_enabled_strings);
	}

	for (i = 0; i < NR_ORDERS; i++) {
		if (!((1 << i) & shmem_orders)) {
			settings->shmem_hugepages[i].enabled = SHMEM_NEVER;
			continue;
		}
		snprintf(path, PATH_MAX, "hugepages-%ukB/shmem_enabled",
			(getpagesize() >> 10) << i);
		settings->shmem_hugepages[i].enabled =
			thp_read_string(path, shmem_enabled_strings);
	}
}

void thp_write_settings(struct thp_settings *settings)
{
	struct khugepaged_settings *khugepaged = &settings->khugepaged;
	unsigned long orders = thp_supported_orders();
	unsigned long shmem_orders = thp_shmem_supported_orders();
	char path[PATH_MAX];
	int enabled;
	int i;

	thp_write_string("enabled", thp_enabled_strings[settings->thp_enabled]);
	thp_write_string("defrag", thp_defrag_strings[settings->thp_defrag]);
	thp_write_string("shmem_enabled",
			shmem_enabled_strings[settings->shmem_enabled]);
	thp_write_num("use_zero_page", settings->use_zero_page);

	thp_write_num("khugepaged/defrag", khugepaged->defrag);
	thp_write_num("khugepaged/alloc_sleep_millisecs",
			khugepaged->alloc_sleep_millisecs);
	thp_write_num("khugepaged/scan_sleep_millisecs",
			khugepaged->scan_sleep_millisecs);
	thp_write_num("khugepaged/max_ptes_none", khugepaged->max_ptes_none);
	thp_write_num("khugepaged/max_ptes_swap", khugepaged->max_ptes_swap);
	thp_write_num("khugepaged/max_ptes_shared", khugepaged->max_ptes_shared);
	thp_write_num("khugepaged/pages_to_scan", khugepaged->pages_to_scan);

	if (dev_queue_read_ahead_path[0])
		write_num(dev_queue_read_ahead_path, settings->read_ahead_kb);

	for (i = 0; i < NR_ORDERS; i++) {
		if (!((1 << i) & orders))
			continue;
		snprintf(path, PATH_MAX, "hugepages-%ukB/enabled",
			(getpagesize() >> 10) << i);
		enabled = settings->hugepages[i].enabled;
		thp_write_string(path, thp_enabled_strings[enabled]);
	}

	for (i = 0; i < NR_ORDERS; i++) {
		if (!((1 << i) & shmem_orders))
			continue;
		snprintf(path, PATH_MAX, "hugepages-%ukB/shmem_enabled",
			(getpagesize() >> 10) << i);
		enabled = settings->shmem_hugepages[i].enabled;
		thp_write_string(path, shmem_enabled_strings[enabled]);
	}
}

struct thp_settings *thp_current_settings(void)
{
	if (!settings_index) {
		printf("Fail: No settings set");
		exit(EXIT_FAILURE);
	}
	return settings_stack + settings_index - 1;
}

void thp_push_settings(struct thp_settings *settings)
{
	if (settings_index >= MAX_SETTINGS_DEPTH) {
		printf("Fail: Settings stack exceeded");
		exit(EXIT_FAILURE);
	}
	settings_stack[settings_index++] = *settings;
	thp_write_settings(thp_current_settings());
}

void thp_pop_settings(void)
{
	if (settings_index <= 0) {
		printf("Fail: Settings stack empty");
		exit(EXIT_FAILURE);
	}
	--settings_index;
	thp_write_settings(thp_current_settings());
}

void thp_restore_settings(void)
{
	if (thp_settings_saved)
		thp_write_settings(&saved_settings);
}

static void __thp_save_settings(void)
{
	if (!thp_available())
		return;

	if (thp_settings_saved)
		return;

	thp_read_settings(&saved_settings);
	thp_settings_saved = true;
}

void thp_set_read_ahead_path(char *path)
{
	if (!path) {
		dev_queue_read_ahead_path[0] = '\0';
		return;
	}

	strncpy(dev_queue_read_ahead_path, path,
		sizeof(dev_queue_read_ahead_path));
	dev_queue_read_ahead_path[sizeof(dev_queue_read_ahead_path) - 1] = '\0';
}

static unsigned long __thp_supported_orders(bool is_shmem)
{
	unsigned long orders = 0;
	char path[PATH_MAX];
	char buf[256];
	int ret, i;
	char anon_dir[] = "enabled";
	char shmem_dir[] = "shmem_enabled";

	for (i = 0; i < NR_ORDERS; i++) {
		ret = snprintf(path, PATH_MAX, THP_SYSFS "hugepages-%ukB/%s",
			       (getpagesize() >> 10) << i, is_shmem ? shmem_dir : anon_dir);
		if (ret >= PATH_MAX) {
			printf("%s: Pathname is too long\n", __func__);
			exit(EXIT_FAILURE);
		}

		ret = read_file(path, buf, sizeof(buf));
		if (ret)
			orders |= 1UL << i;
	}

	return orders;
}

unsigned long thp_supported_orders(void)
{
	return __thp_supported_orders(false);
}

unsigned long thp_shmem_supported_orders(void)
{
	return __thp_supported_orders(true);
}

bool thp_available(void)
{
	if (access(THP_SYSFS, F_OK) != 0)
		return false;
	return true;
}

bool thp_is_enabled(void)
{
	if (!thp_available())
		return false;

	int mode = thp_read_string("enabled", thp_enabled_strings);

	/* THP is considered enabled if it's either "always" or "madvise" */
	return mode == 1 || mode == 3;
}

#define HUGETLB_MAX_NR_PAGESIZES 10
struct hugetlb_settings {
	unsigned long nr_hugepages[HUGETLB_MAX_NR_PAGESIZES];
	unsigned long sizes[HUGETLB_MAX_NR_PAGESIZES];
	unsigned long default_size;
	int nr_sizes;
};

static struct hugetlb_settings hugetlb_saved_settings;
static bool hugetlb_settings_saved;

int detect_hugetlb_page_sizes(unsigned long sizes[], int max)
{
	static struct hugetlb_settings *settings = &hugetlb_saved_settings;
	DIR *dir;
	int count = 0;

	if (settings->nr_sizes) {
		if (settings->nr_sizes < max)
			max = settings->nr_sizes;
		for (count = 0; count < max; count++)
			sizes[count] = settings->sizes[count];
		return count;
	}

	dir = opendir("/sys/kernel/mm/hugepages/");
	if (!dir)
		return 0;

	while (count < max) {
		struct dirent *entry = readdir(dir);
		size_t kb;

		if (!entry)
			break;
		if (entry->d_type != DT_DIR)
			continue;
		if (sscanf(entry->d_name, "hugepages-%zukB", &kb) != 1)
			continue;
		sizes[count++] = kb * 1024;
		ksft_print_msg("[INFO] detected hugetlb page size: %zu KiB\n",
			       kb);
	}
	closedir(dir);
	return count;
}

unsigned long default_huge_page_size(void)
{
	static struct hugetlb_settings *settings = &hugetlb_saved_settings;
	unsigned long hps = 0;
	char *line = NULL;
	size_t linelen = 0;
	FILE *f;

	if (settings->default_size)
		return settings->default_size;

	f = fopen("/proc/meminfo", "r");
	if (!f)
		return 0;
	while (getline(&line, &linelen, f) > 0) {
		if (sscanf(line, "Hugepagesize:       %lu kB", &hps) == 1) {
			hps <<= 10;
			break;
		}
	}

	free(line);
	fclose(f);
	return hps;
}

static void hugetlb_sysfs_path(char *buf, size_t buflen,
			       unsigned long size, const char *attr)
{
	snprintf(buf, buflen, "/sys/kernel/mm/hugepages/hugepages-%lukB/%s",
		 size / 1024, attr);
}

unsigned long hugetlb_nr_pages(unsigned long size)
{
	char path[PATH_MAX];

	hugetlb_sysfs_path(path, sizeof(path), size, "nr_hugepages");

	return read_num(path);
}

void hugetlb_set_nr_pages(unsigned long size, unsigned long nr)
{
	char path[PATH_MAX];

	hugetlb_sysfs_path(path, sizeof(path), size, "nr_hugepages");

	write_num(path, nr);
}

unsigned long hugetlb_free_pages(unsigned long size)
{
	char path[PATH_MAX];

	hugetlb_sysfs_path(path, sizeof(path), size, "free_hugepages");

	return read_num(path);
}

static bool __hugetlb_setup(unsigned long size, unsigned long nr)
{
	unsigned long free = hugetlb_free_pages(size);
	unsigned long total = hugetlb_nr_pages(size);

	if (free >= nr)
		return true;

	hugetlb_set_nr_pages(size, total + (nr - free));

	return hugetlb_free_pages(size) >= nr;
}

bool hugetlb_setup_default(unsigned long nr)
{
	unsigned long size;

	hugetlb_save_settings();
	size = default_huge_page_size();
	if (!size)
		return false;

	return __hugetlb_setup(size, nr);
}

bool hugetlb_setup_default_exact(unsigned long nr)
{
	unsigned long size;

	hugetlb_save_settings();
	size = default_huge_page_size();
	if (!size)
		return false;

	hugetlb_set_nr_pages(size, nr);

	return hugetlb_free_pages(size) == nr;
}

unsigned long hugetlb_setup(unsigned long nr, unsigned long sizes[],
			    int max)
{
	unsigned long enabled[10];
	int nr_sizes = 0;
	int nr_enabled;

	hugetlb_save_settings();

	nr_enabled = detect_hugetlb_page_sizes(enabled, ARRAY_SIZE(enabled));
	if (!nr_enabled)
		return 0;

	if (nr_enabled > max) {
		ksft_print_msg("detected %d huge page sizes, will only test %d\n", nr_enabled, max);
		nr_enabled = max;
	}

	/* request nr HugeTLB pages of every size. */
	for (int i = 0; i < nr_enabled; i++) {
		if (!__hugetlb_setup(enabled[i], nr))
			continue;
		sizes[nr_sizes++] = enabled[i];
	}

	return nr_sizes;
}

static void __hugetlb_save_settings(void)
{
	struct hugetlb_settings *settings = &hugetlb_saved_settings;
	int nr_sizes;

	if (hugetlb_settings_saved)
		return;

	settings->default_size = default_huge_page_size();
	if (!settings->default_size)
		return;

	nr_sizes = detect_hugetlb_page_sizes(settings->sizes,
					     HUGETLB_MAX_NR_PAGESIZES);
	if (!nr_sizes) {
		settings->default_size = 0;
		return;
	}

	for (int i = 0; i < nr_sizes; i++) {
		unsigned long sz = settings->sizes[i];

		if (!sz)
			continue;
		settings->nr_hugepages[i] = hugetlb_nr_pages(sz);
	}

	settings->nr_sizes = nr_sizes;
	hugetlb_settings_saved = true;
}

void hugetlb_restore_settings(void)
{
	struct hugetlb_settings *settings = &hugetlb_saved_settings;

	if (!hugetlb_settings_saved || !settings->default_size)
		return;

	for (int i = 0; i < HUGETLB_MAX_NR_PAGESIZES; i++) {
		unsigned long sz = settings->sizes[i];

		if (!sz)
			continue;

		hugetlb_set_nr_pages(sz, settings->nr_hugepages[i]);
	}
}

static void hugepage_restore_settings_atexit(void)
{
	if (thp_settings_saved)
		thp_restore_settings();
	if (hugetlb_settings_saved)
		hugetlb_restore_settings();
}

static void hugepage_restore_settings_sighandler(int sig)
{
	/* exit() will invoke the hugepage_restore_settings_atexit handler. */
	exit(KSFT_FAIL);
}

void hugepage_save_settings(bool thp, bool hugetlb)
{
	if (!thp && !hugetlb)
		return;

	if (thp)
		__thp_save_settings();
	if (hugetlb)
		__hugetlb_save_settings();

	/*
	 * setup exit hooks to make sure THP and HugeTLB settings are
	 * restored on graceful and error exits and signals
	 */
	atexit(hugepage_restore_settings_atexit);
	signal(SIGTERM, hugepage_restore_settings_sighandler);
	signal(SIGINT, hugepage_restore_settings_sighandler);
	signal(SIGHUP, hugepage_restore_settings_sighandler);
	signal(SIGQUIT, hugepage_restore_settings_sighandler);
}