#include <stdlib.h>
#include <stdio.h>
#include <errno.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>

/*
 * Tries to fragment memory and then executes the given program.
 * Usage:
 * 	fragment-mem-and-run program_to_run mem_to_allocate mem_to_free
 *
 * First tries to allocate mem_to_allocate[KMG] amount of memory which,
 * then tries to free mem_to_free[KMG] in a way to maximize the fragmentation
 * of the page allocator. It is advisable to run compaction before starting
 * to get reproducible behavior.
 *
 * Copyright Michal Hocko 2016
 */
#define PAGE_SIZE 4096UL
#define MAX_ORDER 11
#define ALIGN(x, a) (((x) + (a) - 1) & ~((a) - 1))
#define PAGE_ALIGN(addr) ALIGN(addr, PAGE_SIZE)

size_t parse_size(const char *value)
{
	char *endptr;
	size_t size = strtoul(value, &endptr, 10);

	if (*endptr == 'K')
		size *= 1024;
	else if (*endptr == 'M')
		size *= 1024*1024;
	else if (*endptr == 'G')
		size *= 1024*1024*1024;
	else if (*endptr)
		size = -1UL;

	return size;
}

void dump_file(const char *filename)
{
	char buffer[BUFSIZ];
	int fd;

	fd = open(filename, O_RDONLY);
	if (fd == -1)
		return;

	while (read(fd, buffer, sizeof(buffer)))
		printf("%s", buffer);

	printf("\n");
	close(fd);
}

int main(int argc, char **argv)
{
	size_t size = 10<<20;
	size_t to_free, freed = 0;
	size_t i, step = PAGE_SIZE*((1UL<<MAX_ORDER)-1);
	unsigned char *addr;
	int buddy_fd;
	const char *to_run;

	if (argc > 1) {
		to_run = argv[1];
	} else {
		fprintf(stderr, "Didn't tell me what to run");
		return 1;
	}
	if (argc > 2) {
		size = parse_size(argv[2]);
		if (size == -1UL) {
			fprintf(stderr, "Number expected \"%s\" given.\n", argv[0]);
			return 1;
		}
	}
	if (argc > 3) {
		to_free = parse_size(argv[3]);
		if (to_free == -1UL) {
			fprintf(stderr, "Number expected \"%s\" given.\n", argv[0]);
			return 1;
		}
	} else {
		to_free = size;
	}

	dump_file("/proc/buddyinfo");
	addr = mmap(NULL, size, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
	if (addr == MAP_FAILED) {
		perror("mmap");
		return 3;
	}
	madvise(addr, size, MADV_NOHUGEPAGE);
	for (i = 0; i < size; i += PAGE_SIZE)
		addr[i] = 1;

	while (freed < to_free) {
		for (i = step; (freed < to_free) && (i < size); i = (i + step) % size) {
			i = PAGE_ALIGN(i);
			if (madvise(&addr[i], PAGE_SIZE, MADV_DONTNEED))
				continue;
			freed += PAGE_SIZE;
		}
		step = (step / 2) + 1;
	}

	printf("Done fragmenting. size=%lu freed=%lu\n", size, freed);
	dump_file("/proc/buddyinfo");
	printf("Executing \"%s\"\n", to_run);
	fflush(stdout);
	return system(to_run);
}