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

#define CHUNK_SIZE (20UL<<20)
#ifndef PAGE_SIZE
#define PAGE_SIZE (4096)
#endif

/*
 * For native build, use aarch64-linux-android-clang -pie -o memhog memhog.c
 */
void usage(char *exe)
{
    fprintf(stderr,
            "Usage: %s [options] size[k|m|g]\n"
            "    -c|--compress ratio fill memory with ratio compressible data\n"
            "    -m|--memory SIZE    allocate memory in SIZE byte chunks\n"
            "    -M|--mlock          mlock() the memory\n"
            "    -s|--sleep SEC      sleep SEC seconds during repeat cycle\n"
            "    -r|--repeat N       repeat read/write N times\n"
            "    -h|--help           show this message\n",
            exe);

    exit(1);
}

static const struct option opts[] = {
    { "compress", 1, NULL, 'c' },
    { "memory"  , 1, NULL, 'm' },
    { "mlock"   , 0, NULL, 'M' },
    { "sleep"   , 1, NULL, 's' },
    { "repeat"  , 1, NULL, 'r' },
    { "help"    , 0, NULL, 'h' },
    { NULL      , 0, NULL, 0 }
};

unsigned long long memparse(const char *ptr, char **retptr)
{
    char *endptr;

    unsigned long long ret = strtoull(ptr, &endptr, 0);

    switch (*endptr) {
        case 'G':
        case 'g':
            ret <<= 10;
        case 'M':
        case 'm':
            ret <<= 10;
        case 'K':
        case 'k':
            ret <<= 10;
            endptr++;
        default:
            break;
    }

    if (retptr)
        *retptr = endptr;

    return ret;
}

void allocate_mem(unsigned long long size, void *alloc_ptr[], int *len)
{
    int i;
    void *ptr;
    unsigned long nr_chunk = size / CHUNK_SIZE;
    int allocated = 0;

    for (i = 0; i < nr_chunk; i++) {
        ptr = mmap(NULL, CHUNK_SIZE, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANON, 0, 0);
        if (ptr == MAP_FAILED) {
            printf("fail to allocate %d\n", i);
            break;
        }

        alloc_ptr[allocated++] = ptr;
    }

    *len = allocated;
}

void free_mem(void *alloc_ptr[], int len)
{
    int i;

    for (i = 0; i < len; i++)
        munmap(alloc_ptr[i], CHUNK_SIZE);
}

void fill_mem(void *ptr, void *rand_page, long int comp_ratio)
{
    int i;
    static int nr_page = CHUNK_SIZE / PAGE_SIZE;
    int zero_size = PAGE_SIZE * comp_ratio / 100;

    for (i = 0; i < nr_page; i++, ptr += PAGE_SIZE) {
        memset(ptr, 0, zero_size);
        memcpy(ptr + zero_size, rand_page, PAGE_SIZE - zero_size);
    }
}

int fill_chunk(void *alloc_ptr[], int len, long int comp_ratio)
{
    int i, ret;
    char rand_buf[PAGE_SIZE];
    int fd = open("/dev/urandom", O_RDONLY);

    if (fd < 0) {
        perror("Fail to open /dev/urandom\n");
        return 1;
    }

    ret = read(fd, rand_buf, PAGE_SIZE);
    if (ret != PAGE_SIZE) {
        perror("Fail to read /dev/urandom\n");
        return 1;
    }

    for (i = 0; i < len; i++)
        fill_mem(alloc_ptr[i], rand_buf, comp_ratio);

    close(fd);
    return 0;
}

int main(int argc, char *argv[])
{
    char buf[256] = {0,};
    unsigned long long opt_mem = 100 << 20;
    long int opt_sleep = 1;
    long int opt_reps = 10000;
    long int opt_comp_ratio = 30;
    unsigned long opt_mlock = 0;
    long int loops;
    int pid = getpid();
    int err, c, count;

    while ((c = getopt_long(argc, argv,
                    "m:Ms:r:c:", opts, NULL)) != -1) {
        switch (c) {
            case 'c':
                opt_comp_ratio = strtol(optarg, NULL, 10);
                break;
            case 'm':
                opt_mem = memparse(optarg, NULL);
                break;
            case 'M':
                opt_mlock = 1;
                break;
            case 's':
                opt_sleep = strtol(optarg, NULL, 10);
                break;
            case 'r':
                opt_reps = strtol(optarg, NULL, 10);
                break;
            case 'h':
                usage(argv[0]);
                break;
            default:
                usage(argv[0]);
        }
    }

    if (opt_mem < CHUNK_SIZE) {
        printf("memory size should be greater than %lu\n", CHUNK_SIZE);
        return 1;
    }

    /* Disable LMK/OOM killer */
    sprintf(buf, "echo -1000 > /proc/%d/oom_score_adj\n", pid);
    if (WEXITSTATUS(system(buf))) {
        fprintf(stderr, "fail to disable OOM. Maybe you need root permission\n");
        return 1;
    }

    if (opt_mlock) {
        err = mlockall(MCL_CURRENT|MCL_FUTURE|MCL_ONFAULT);
        if (err) {
            perror("Fail to mlockall\n");
            return err;
        }
    }

    printf("%llu MB allocated\n", opt_mem >> 20);
    printf("%lu loop\n", opt_reps);
    printf("%lu sleep\n", opt_sleep);
    printf("%lu comp_ratio\n", opt_comp_ratio);
    count = 0;
    loops = opt_reps;

    while (loops) {
        /* 20M * 4096 = 80G is enouggh */
        void *alloc_ptr[PAGE_SIZE];
        int len;

        count++;
retry:
        allocate_mem(opt_mem, alloc_ptr, &len);
        if (len == 0) {
            /*
             * If we couldn't allocate any memory, let's try again
             * after a while
             */
            sleep(1);
            goto retry;
        }

        if (fill_chunk(alloc_ptr, len, opt_comp_ratio)) {
            printf("Fail to fill chunck\n");
            return 1;
        }

        if (opt_sleep == -1) {
            while (1) {
                printf("Forever sleep, Bye\n");
                sleep(100000);
            }
        }

        sleep(opt_sleep);
        free_mem(alloc_ptr, len);
        if (loops != -1)
            loops--;
        printf("[%d] Pass %d\n", pid, count);
    }

    return 0;
}