/***************************************************************************
 *   User front end for using huge pages Copyright (C) 2008, IBM           *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the Lesser GNU General Public License as        *
 *   published by the Free Software Foundation; either version 2.1 of the  *
 *   License, or at your option) any later version.                        *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU Lesser General Public License for more details.                   *
 *                                                                         *
 *   You should have received a copy of the Lesser GNU General Public      *
 *   License along with this program; if not, write to the                 *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/

/*
 * hugectl is inspired by numactl as a single front end to a large number of
 * options for controlling a very specific environment.  Eventually it will
 * have support for controlling the all of the environment variables for
 * libhugetlbfs, but options will only be added after they have been in the
 * library for some time and are throughly tested and stable.
 *
 * This program should be treated as an ABI for using libhugetlbfs.
 */

#include <stdlib.h>
#include <stdio.h>
#include <errno.h>
#include <string.h>
#include <limits.h>

#define _GNU_SOURCE /* for getopt_long */
#include <unistd.h>
#include <getopt.h>

#define REPORT(level, prefix, format, ...)				      \
	do {								      \
		if (verbose_level >= level)				      \
			fprintf(stderr, "hugectl: " prefix ": " format,	      \
				##__VA_ARGS__);				      \
	} while (0);

#include "libhugetlbfs_debug.h"

extern int errno;
extern int optind;
extern char *optarg;

#define OPTION(opts, text)	fprintf(stderr, " %-25s  %s\n", opts, text)
#define CONT(text) 		fprintf(stderr, " %-25s  %s\n", "", text)

void print_usage()
{
	fprintf(stderr, "hugectl [options] target\n");
	fprintf(stderr, "options:\n");

	OPTION("--help, -h", "Prints this message");
	OPTION("--verbose <level>, -v", "Increases/sets tracing levels");

	OPTION("--text[=<size>]", "Requests remapping of the program text");
	OPTION("--data[=<size>]", "Requests remapping of the program data");
	OPTION("--bss[=<size>]", "Requests remapping of the program bss");
	OPTION("--heap[=<size>]", "Requests remapping of the program heap");
	CONT("(malloc space)");
	OPTION("--shm", "Requests remapping of shared memory segments");
	OPTION("--thp", "Setup the heap space to be aligned for merging");
	CONT("by khugepaged into huge pages.  This requires");
	CONT("kernel support for transparent huge pages to be");
	CONT("enabled");

	OPTION("--no-preload", "Disable preloading the libhugetlbfs library");
	OPTION("--no-reserve", "Disable huge page reservation for segments");
	OPTION("--force-preload", "Force preloading the libhugetlbfs library");

	OPTION("--dry-run", "describe what would be done without doing it");

	OPTION("--library-use-path", "Use the system library path");
	OPTION("--share-text", "Share text segments between multiple");
	CONT("application instances");
	OPTION("--library-path <path>", "Select a library prefix");
	CONT("(Default: "
#ifdef LIBDIR32
		LIBDIR32 ":"
#endif
#ifdef LIBDIR32
		LIBDIR32 ":"
#endif
		")");
}

int opt_dry_run = 0;
int opt_force_preload = 0;
int verbose_level = VERBOSITY_DEFAULT;

void verbose_init(void)
{
	char *env;

	env = getenv("HUGETLB_VERBOSE");
	if (env)
		verbose_level = atoi(env);
	env = getenv("HUGETLB_DEBUG");
	if (env)
		verbose_level = VERBOSITY_MAX;
}

void verbose(char *which)
{
	int new_level;

	if (which) {
		new_level = atoi(which);
		if (new_level < 0 || new_level > 99) {
			ERROR("%d: verbosity out of range 0-99\n",
				new_level);
			exit(EXIT_FAILURE);
		}
	} else {
		new_level = verbose_level + 1;
		if (new_level == 100) {
			WARNING("verbosity limited to 99\n");
			new_level--;
		}
	}
	verbose_level = new_level;
}

void quiet(void)
{
	int new_level = verbose_level - 1;
	if (new_level < 0) {
		WARNING("verbosity must be at least 0\n");
		new_level = 0;
	}
	verbose_level = new_level;
}

void setup_environment(char *var, char *val)
{
	setenv(var, val, 1);
	INFO("%s='%s'\n", var, val);

	if (opt_dry_run)
		printf("%s='%s'\n", var, val);
}

void verbose_expose(void)
{
	char level[3];

	if (verbose_level == 99) {
		setup_environment("HUGETLB_DEBUG", "yes");
	}
	snprintf(level, sizeof(level), "%d", verbose_level);
	setup_environment("HUGETLB_VERBOSE", level);
}

/*
 * getopts return values for options which are long only.
 */
#define MAP_BASE	0x1000
#define LONG_BASE	0x2000

#define LONG_NO_PRELOAD		(LONG_BASE | 'p')
#define LONG_NO_RESERVE		(LONG_BASE | 'r')
#define LONG_FORCE_PRELOAD	(LONG_BASE | 'F')

#define LONG_DRY_RUN		(LONG_BASE | 'd')

#define LONG_SHARE		(LONG_BASE | 's')
#define LONG_NO_LIBRARY		(LONG_BASE | 'L')
#define LONG_LIBRARY		(LONG_BASE | 'l')

#define LONG_THP_HEAP		('t')

/*
 * Mapping selectors, one per remappable/backable area as requested
 * by the user.  These are also used as returns from getopts where they
 * are offset from MAP_BASE, which must be removed before they are compared.
 */
enum {
	MAP_TEXT,
	MAP_DATA,
	MAP_BSS,
	MAP_HEAP,
	MAP_SHM,
	MAP_DISABLE,

	MAP_COUNT,
};
char *map_size[MAP_COUNT];

char default_size[] = "the default hugepage size";
#define DEFAULT_SIZE default_size

#define available(buf, ptr) ((int)(sizeof(buf) - (ptr - buf)))
void setup_mappings(int count)
{
	char value[128];
	char *ptr = value;
	int needed;

	/*
	 * HUGETLB_ELFMAP should be set to either a combination of 'R' and 'W'
	 * which indicate which segments should be remapped.  Each may take
	 * an optional page size.  It may also be set to 'no' to prevent
	 * remapping.
	 */

	/*
	 * Accumulate sections each with a ':' prefix to simplify later
	 * handling.  We will elide the initial ':' before use.
	 */
	if (map_size[MAP_TEXT]) {
		if (map_size[MAP_TEXT] == DEFAULT_SIZE)
			needed = snprintf(ptr, available(value, ptr), ":R");
		else
			needed = snprintf(ptr, available(value, ptr),
						":R=%s", map_size[MAP_TEXT]);
		ptr += needed;
		if (needed < 0 || available(value, ptr) < 0) {
			ERROR("%s: bad size specification\n", map_size[MAP_TEXT]);
			exit(EXIT_FAILURE);
		}
	}
	if (map_size[MAP_DATA] != 0 || map_size[MAP_BSS] != 0) {
		char *size = map_size[MAP_BSS];
		if (map_size[MAP_DATA])
			size = map_size[MAP_DATA];
		if (map_size[MAP_DATA] != map_size[MAP_BSS])
			WARNING("data and bss remapped together in %s\n", size);

		if (size == DEFAULT_SIZE)
			needed = snprintf(ptr, available(value, ptr), ":W");
		else
			needed = snprintf(ptr, available(value, ptr),
						":W=%s", size);
		ptr += needed;
		if (needed < 0 || available(value, ptr) < 0) {
			ERROR("%s: bad size specification\n", size);
			exit(EXIT_FAILURE);
		}
	}
	*ptr = '\0';
	if (ptr != value)
		setup_environment("HUGETLB_ELFMAP", &value[1]);

	if (map_size[MAP_DISABLE]) {
		if (ptr != value)
			WARNING("--disable masks requested remap\n");
		setup_environment("HUGETLB_ELFMAP", "no");
	}

	if (map_size[MAP_HEAP] == DEFAULT_SIZE)
		setup_environment("HUGETLB_MORECORE", "yes");
	else if (map_size[MAP_HEAP])
		setup_environment("HUGETLB_MORECORE", map_size[MAP_HEAP]);

	if (map_size[MAP_SHM] && map_size[MAP_SHM] != DEFAULT_SIZE)
		WARNING("shm segments may only be mapped in the "
			"default hugepage size\n");
	if (map_size[MAP_SHM])
		setup_environment("HUGETLB_SHM", "yes");
}

#define LIBRARY_DISABLE ((void *)-1)

void library_path(char *path)
{
	char val[PATH_MAX] = "";
	char *env;

	env = getenv("LD_LIBRARY_PATH");

	/*
	 * Select which libraries we wish to use.  If the path is NULL
	 * use the libraries included with hugectl.  If the path is valid
	 * and points to a directory including a libhugetlbfs.so use it
	 * directly.  Else path is assumed to be a prefix to the 32/64 bit
	 * directories both of which are added, where available.
	 */
	if (path) {
		snprintf(val, sizeof(val), "%s/libhugetlbfs.so", path);
		if (access(val, F_OK) == 0) {
			/* $PATH */
			snprintf(val, sizeof(val), "%s:%s",
				path, env ? env : "");

		} else {
			/* [$PATH/LIB32:][$PATH/LIB64:]$LD_LIBRARY_PATH */
			snprintf(val, sizeof(val), ""
#ifdef LIBDIR32
				"%s/" LIB32 ":"
#endif
#ifdef LIBDIR64
				"%s/" LIB64 ":"
#endif
				"%s",
#ifdef LIBDIR32
				path,
#endif
#ifdef LIBDIR64
				path,
#endif
				env ? env : "");
		}

	} else {
		/* [LIBDIR32:][LIBDIR64:]$LD_LIBRARY_PATH */
		snprintf(val, sizeof(val), ""
#ifdef LIBDIR32
			LIBDIR32 ":"
#endif
#ifdef LIBDIR64
			LIBDIR64 ":"
#endif
			"%s", env ? env : "");
	}
	setup_environment("LD_LIBRARY_PATH", val);
}

void ldpreload(int count)
{
	int allowed = 0;

	if (map_size[MAP_HEAP])
		allowed++;
	if (map_size[MAP_SHM])
		allowed++;

	if ((allowed == count) || opt_force_preload) {
		setup_environment("LD_PRELOAD", "libhugetlbfs.so");
		if (allowed == count)
			INFO("LD_PRELOAD in use for lone --heap/--shm\n");
	} else {
		WARNING("LD_PRELOAD not appropriate for this map combination\n");
	}
}

int main(int argc, char** argv)
{
	int opt_mappings = 0;
	int opt_preload = 1;
	int opt_no_reserve = 0;
	int opt_share = 0;
	int opt_thp_heap = 0;
	char *opt_library = NULL;

	char opts[] = "+hvq";
	int ret = 0, index = 0;
	struct option long_opts[] = {
		{"help",       no_argument, NULL, 'h'},
		{"verbose",    required_argument, NULL, 'v' },
		{"no-preload", no_argument, NULL, LONG_NO_PRELOAD},
		{"no-reserve", no_argument, NULL, LONG_NO_RESERVE},
		{"force-preload",
			       no_argument, NULL, LONG_FORCE_PRELOAD},
		{"dry-run",    no_argument, NULL, LONG_DRY_RUN},
		{"library-path",
			       required_argument, NULL, LONG_LIBRARY},
		{"library-use-path",
			       no_argument, NULL, LONG_NO_LIBRARY},
		{"share-text", no_argument, NULL, LONG_SHARE},

		{"disable",    optional_argument, NULL, MAP_BASE|MAP_DISABLE},
		{"text",       optional_argument, NULL, MAP_BASE|MAP_TEXT},
		{"data",       optional_argument, NULL, MAP_BASE|MAP_DATA},
		{"bss",        optional_argument, NULL, MAP_BASE|MAP_BSS},
		{"heap",       optional_argument, NULL, MAP_BASE|MAP_HEAP},
		{"shm",        optional_argument, NULL, MAP_BASE|MAP_SHM},
		{"thp",        no_argument, NULL, LONG_THP_HEAP},
		{0},
	};

	verbose_init();

	while (ret != -1) {
		ret = getopt_long(argc, argv, opts, long_opts, &index);
		if (ret > 0 && (ret & MAP_BASE)) {
			if (optarg)
				map_size[ret & ~MAP_BASE] = optarg;
			else
				map_size[ret & ~MAP_BASE] = DEFAULT_SIZE;
			opt_mappings++;
			continue;
		}
		switch (ret) {
		case '?':
			print_usage();
			exit(EXIT_FAILURE);

		case 'h':
			print_usage();
			exit(EXIT_SUCCESS);

		case 'v':
			verbose(optarg);
			break;

		case 'q':
			quiet();
			break;

		case LONG_THP_HEAP:
			opt_thp_heap = 1;
			INFO("Aligning heap for use with THP\n");
			break;

		case LONG_NO_PRELOAD:
			opt_preload = 0;
			INFO("LD_PRELOAD disabled\n");
			break;

		case LONG_NO_RESERVE:
			opt_no_reserve = 1;
			INFO("MAP_NORESERVE used for huge page mappings\n");
			break;

		case LONG_FORCE_PRELOAD:
			opt_preload = 1;
			opt_force_preload = 1;
			INFO("Forcing ld preload\n");
			break;

		case LONG_DRY_RUN:
			opt_dry_run = 1;
			break;

		case LONG_NO_LIBRARY:
			opt_library = LIBRARY_DISABLE;
			INFO("using LD_LIBRARY_PATH to find library\n");
			break;

		case LONG_LIBRARY:
			opt_library = optarg;
			break;

		case LONG_SHARE:
			opt_share = 1;
			break;

		case -1:
			break;

		default:
			WARNING("unparsed option %08x\n", ret);
			ret = -1;
			break;
		}
	}
	index = optind;

	if (!opt_dry_run && (argc - index) < 1) {
		print_usage();
		exit(EXIT_FAILURE);
	}

	verbose_expose();

	if (opt_library != LIBRARY_DISABLE)
		library_path(opt_library);

	if (opt_mappings)
		setup_mappings(opt_mappings);

	if (opt_preload)
		ldpreload(opt_mappings);

	if (opt_no_reserve)
		setup_environment("HUGETLB_NO_RESERVE", "yes");

	if (opt_share)
		setup_environment("HUGETLB_SHARE", "1");

	if (opt_thp_heap)
		setup_environment("HUGETLB_MORECORE", "thp");

	if (opt_dry_run)
		exit(EXIT_SUCCESS);

	execvp(argv[index], &argv[index]);
	ERROR("exec failed: %s\n", strerror(errno));
	exit(EXIT_FAILURE);
}