#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/vmalloc.h>
#include <linux/gfp.h>
#include <linux/mm.h>

MODULE_LICENSE("GPL");
MODULE_AUTHOR("mlombard");
MODULE_DESCRIPTION("OOM frag test module");

static struct page_frag_cache nc;

static int memory_size = 1; /* 1 Gb */
static int fragsz = 4096;

static int __init oomk_init(void)
{
	unsigned long i, vsize;
	void **vec;
	struct page *page;

	printk("Test begins, memory size = %d fragsize = %d\n", memory_size, fragsz);

	vsize = ((unsigned long)memory_size * 1073741824) / fragsz;

	vec = vzalloc(vsize * sizeof(void *));
	if (!vec) {
		printk("Could not allocate vector\n");
		return -1;
	}

	for (i = 0; i < vsize; ++i) {
		vec[i] = page_frag_alloc(&nc, fragsz, GFP_ATOMIC);
		if (!vec[i]) {
			printk("allocation number %lu failed!\n", i);
			break;
		}
	}

	for (i = 0; i < vsize; ++i) {
		if (vec[i])
			page_frag_free(vec[i]);
		else
			break;
	}

	if (nc.va) {
		page = virt_to_head_page(nc.va);
		__page_frag_cache_drain(page, nc.pagecnt_bias);
	}

	vfree(vec);

	printk("Test completed!\n");
	return 0;
}

static void __exit oomk_cleanup(void)
{
	printk("Cleaning up module.\n");
}

module_param_named(memory_size_gb, memory_size, int, S_IRUGO | S_IWUSR);
module_param_named(fragsize, fragsz, int, S_IRUGO | S_IWUSR);


module_init(oomk_init);
module_exit(oomk_cleanup);