/*

 * libhugetlbfs - Easy use of Linux hugepages

 * Copyright (C) 2005-2006 David Gibson & Adam Litke, IBM Corporation.

 *

 * This library is free software; you can redistribute it and/or

 * modify it under the terms of the GNU Lesser 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 library 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 GNU Lesser General Public

 * License along with this library; if not, write to the Free Software

 * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA

 */

#define _GNU_SOURCE

#include <stdio.h>

#include <stdlib.h>

#include <malloc.h>

#include <unistd.h>

#include <sys/mman.h>

#include <errno.h>

#include <dlfcn.h>

#include <string.h>

#include <fcntl.h>

#include "hugetlbfs.h"

#include "libhugetlbfs_internal.h"

static int heap_fd;

static void *heapbase;

static void *heaptop;

static long mapsize;

static long hpage_size;

static long hugetlbfs_next_addr(long addr)

{

#if defined(__powerpc64__)

	return ALIGN(addr, 1L << SLICE_HIGH_SHIFT);

#elif defined(__powerpc__) && !defined(PPC_NO_SEGMENTS)

	return ALIGN(addr, 1L << SLICE_LOW_SHIFT);

#elif defined(__ia64__)

	if (addr < (1UL << SLICE_HIGH_SHIFT))

		return ALIGN(addr, 1UL << SLICE_HIGH_SHIFT);

	else

		return ALIGN(addr, hpage_size);

#else

	return ALIGN(addr, hpage_size);

#endif

}

/*

 * Our plan is to ask for pages 'roughly' at the BASE.  We expect and

 * require the kernel to offer us sequential pages from wherever it

 * first gave us a page.  If it does not do so, we return the page and

 * pretend there are none this covers us for the case where another

 * map is in the way.  This is required because 'morecore' must have

 * 'sbrk' semantics, ie. return sequential, contigious memory blocks.

 * Luckily, if it does not do so and we error out malloc will happily

 * go back to small pages and use mmap to get them.  Hurrah.

 */

static void *hugetlbfs_morecore(ptrdiff_t increment)

{

	int ret;

	void *p;

	long delta;

	int mmap_reserve = __hugetlb_opts.no_reserve ? MAP_NORESERVE : 0;

	int mmap_hugetlb = 0;

	int using_default_pagesize =

		(hpage_size == kernel_default_hugepage_size());

	INFO("hugetlbfs_morecore(%ld) = ...\n", (long)increment);

	/*

	 * how much to grow the heap by =

	 * 	(size of heap) + malloc request - mmap'd space

	 */

	delta = (heaptop-heapbase) + increment - mapsize;

	INFO("heapbase = %p, heaptop = %p, mapsize = %lx, delta=%ld\n",

	      heapbase, heaptop, mapsize, delta);

	/* align to multiple of hugepagesize. */

	delta = ALIGN(delta, hpage_size);

#ifdef MAP_HUGETLB

	mmap_hugetlb = MAP_HUGETLB;

#endif

	if (delta > 0) {

		/* growing the heap */

		INFO("Attempting to map %ld bytes\n", delta);

		/* map in (extend) more of the file at the end of our last map */

		if (__hugetlb_opts.map_hugetlb && using_default_pagesize)

			p = mmap(heapbase + mapsize, delta, PROT_READ|PROT_WRITE,

				 mmap_hugetlb|MAP_ANONYMOUS|MAP_PRIVATE|mmap_reserve,

				 heap_fd, mapsize);

		else

			p = mmap(heapbase + mapsize, delta, PROT_READ|PROT_WRITE,

				 MAP_PRIVATE|mmap_reserve, heap_fd, mapsize);

		if (p == MAP_FAILED) {

			WARNING("New heap segment map at %p failed: %s\n",

				heapbase+mapsize, strerror(errno));

			return NULL;

		}

		/* if this is the first map */

		if (! mapsize) {

			if (heapbase && (heapbase != p)) {

				WARNING("Heap originates at %p instead of %p\n",

					p, heapbase);

				if (__hugetlbfs_debug)

					dump_proc_pid_maps();

			}

			/* then setup the heap variables */

			heapbase = heaptop = p;

		} else if (p != (heapbase + mapsize)) {

			/* Couldn't get the mapping where we wanted */

			munmap(p, delta);

			WARNING("New heap segment mapped at %p instead of %p\n",

			      p, heapbase + mapsize);

			if (__hugetlbfs_debug)

				dump_proc_pid_maps();

			return NULL;

		}

		/* Fault the region to ensure accesses succeed */

		if (hugetlbfs_prefault(p, delta) != 0) {

			munmap(p, delta);

			return NULL;

		}

		/* we now have mmap'd further */

		mapsize += delta;

	} else if (delta < 0) {

		/* shrinking the heap */

		if (!__hugetlb_opts.shrink_ok) {

			/* shouldn't ever get here */

			WARNING("Heap shrinking is turned off\n");

			return NULL;

		}

		if (!mapsize) {

			WARNING("Can't shrink empty heap!\n");

			return NULL;

		}

		/*

		 * If we are forced to change the heapaddr from the

		 * original brk() value we have violated brk semantics

		 * (which we are not supposed to do).  This shouldn't

		 * pose a problem until glibc tries to trim the heap to an

		 * address lower than what we aligned heapaddr to.  At that

		 * point the alignment "gap" causes heap corruption.

		 * So we don't allow the heap to shrink below heapbase.

		 */

		if (mapsize + delta < 0) {  /* remember: delta is negative */

			WARNING("Unable to shrink heap below %p\n", heapbase);

			/* unmap just what is currently mapped */

			delta = -mapsize;

			/* we need heaptop + increment == heapbase, so: */

			increment = heapbase - heaptop;

		}

		INFO("Attempting to unmap %ld bytes @ %p\n", -delta,

			heapbase + mapsize + delta);

		ret = munmap(heapbase + mapsize + delta, -delta);

		if (ret) {

			WARNING("Unmapping failed while shrinking heap: "

				"%s\n", strerror(errno));

		} else {

			mapsize += delta;

			/*

			* the glibc assumes by default that newly allocated

			* memory by morecore() will be zeroed.  It would be

			* wasteful to do it for allocation so we only shrink

			* the top by the size of a page.

			*/

			increment = heapbase - heaptop + mapsize;

			if (!__hugetlb_opts.map_hugetlb && !using_default_pagesize){

				/*

				* Now shrink the hugetlbfs file.

				*/

				ret = ftruncate(heap_fd, mapsize);

				if (ret) {

					WARNING("Could not truncate hugetlbfs file to "

						"shrink heap: %s\n", strerror(errno));

				}

			}

		}

	}

	else if (increment < 0) {

		/* Don't shrink by less than a page to avoid having to zero

		 * the memory.  There is no point in lying to glibc since

		 * we're not freeing any memory.

		 */

		increment = 0;

	}

	/* heap is continuous */

	p = heaptop;

	/* and we now have added this much more space to the heap */

	heaptop = heaptop + increment;

	INFO("... = %p\n", p);

	return p;

}

static void *thp_morecore(ptrdiff_t increment)

{

	void *p;

	long delta;

	INFO("thp_morecore(%ld) = ...\n", (long)increment);

	delta = (heaptop - heapbase) + increment - mapsize;

	delta = ALIGN(delta, hpage_size);

	if (delta > 0) {

		/*

		 * This first time we expand the mapping we need to account for

		 * the initial heap mapping not necessarily being huge page

		 * aligned

		 */

		if (!mapsize)

			delta = hugetlbfs_next_addr((long)heapbase + delta) -

					(unsigned long)heapbase;

		INFO("Adding %ld bytes to heap\n", delta);

		p = sbrk(delta);

		if (p == (void *)-1) {

			WARNING("sbrk returned ENOMEM\n");

			return NULL;

		}

		if (!mapsize) {

			if (heapbase && (heapbase != p)) {

				WARNING("Heap was expected at %p instead of %p, "

					"heap has been modified by someone else!\n",

					heapbase, p);

				if (__hugetlbfs_debug)

					dump_proc_pid_maps();

			}

			heapbase = heaptop = p;

		}

		mapsize += delta;

#ifdef MADV_HUGEPAGE

		madvise(p, delta, MADV_HUGEPAGE);

#endif

	} else if (delta < 0) {

		/* shrinking the heap */

		if (!mapsize) {

			WARNING("Can't shrink an empty heap\n");

			return NULL;

		}

		INFO("Attempting to shrink heap by %ld bytes with sbrk\n",

			-delta);

		p = sbrk(delta);

		if (p == (void *)-1) {

			WARNING("Unable to shrink heap\n");

			return heaptop;

		}

		mapsize += delta;

	}

	p = heaptop;

	heaptop += increment;

	INFO("... = %p\n", p);

	return p;

}

void hugetlbfs_setup_morecore(void)

{

	char *ep;

	unsigned long heapaddr;

	if (! __hugetlb_opts.morecore)

		return;

	if (strcasecmp(__hugetlb_opts.morecore, "no") == 0) {

		INFO("HUGETLB_MORECORE=%s, not setting up morecore\n",

						__hugetlb_opts.morecore);

		return;

	}

	/*

	 * Determine the page size that will be used for the heap.

	 * This can be set explicitly by setting HUGETLB_MORECORE to a valid

	 * page size string or by setting HUGETLB_DEFAULT_PAGE_SIZE.

	 */

	if (strncasecmp(__hugetlb_opts.morecore, "y", 1) == 0)

		hpage_size = gethugepagesize();

	else if (__hugetlb_opts.thp_morecore)

		hpage_size = kernel_default_hugepage_size();

	else

		hpage_size = parse_page_size(__hugetlb_opts.morecore);

	if (hpage_size <= 0) {

		if (errno == ENOSYS)

			WARNING("Hugepages unavailable\n");

		else if (errno == EOVERFLOW || errno == ERANGE)

			WARNING("Hugepage size too large\n");

		else if (errno == EINVAL)

			WARNING("Invalid huge page size\n");

		else

			WARNING("Hugepage size (%s)\n", strerror(errno));

		return;

	}

	/*

	 * We won't need an fd for the heap mmaps if we are using MAP_HUGETLB

	 * or we are depending on transparent huge pages

	 */

	if(__hugetlb_opts.thp_morecore || (__hugetlb_opts.map_hugetlb &&

			hpage_size == kernel_default_hugepage_size())) {

		heap_fd = -1;

	} else {

		if (!hugetlbfs_find_path_for_size(hpage_size)) {

			WARNING("Hugepage size %li unavailable", hpage_size);

			return;

		}

		heap_fd = hugetlbfs_unlinked_fd_for_size(hpage_size);

		if (heap_fd < 0) {

			WARNING("Couldn't open hugetlbfs file for morecore\n");

			return;

		}

	}

	/*

	 * THP morecore uses sbrk to allocate more heap space, counting on the

	 * kernel to back the area with THP.  So setting heapbase is

	 * meaningless if thp_morecore is used.

	 */

	if (!__hugetlb_opts.thp_morecore && __hugetlb_opts.heapbase) {

		heapaddr = strtoul(__hugetlb_opts.heapbase, &ep, 16);

		if (*ep != '\0') {

			WARNING("Can't parse HUGETLB_MORECORE_HEAPBASE: %s\n",

			      __hugetlb_opts.heapbase);

			return;

		}

	} else {

		heapaddr = (unsigned long)sbrk(0);

		if (!__hugetlb_opts.thp_morecore)

			heapaddr = hugetlbfs_next_addr(heapaddr);

	}

	INFO("setup_morecore(): heapaddr = 0x%lx\n", heapaddr);

	heaptop = heapbase = (void *)heapaddr;

	if (__hugetlb_opts.thp_morecore)

		__morecore = &thp_morecore;

	else

		__morecore = &hugetlbfs_morecore;

	/* Set some allocator options more appropriate for hugepages */

	if (__hugetlb_opts.shrink_ok)

		mallopt(M_TRIM_THRESHOLD, hpage_size + hpage_size / 2);

	else

		mallopt(M_TRIM_THRESHOLD, -1);

	mallopt(M_TOP_PAD, hpage_size / 2);

	/* we always want to use our morecore, not ordinary mmap().

	 * This doesn't appear to prohibit malloc() from falling back

	 * to mmap() if we run out of hugepages. */

	mallopt(M_MMAP_MAX, 0);

}