/*
 * Mlock memory to flush caches and create availible memory
 * for Sun's VDI on undersized servers.
 *
 * Based on James Litchfield's blog entry on locking memory
 *          http://blogs.sun.com/thejel/entry/locking_memory
 */

#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <sys/param.h>
#include <malloc.h>
#include <unistd.h>
#include <errno.h>
#include <unistd.h>

#define MEG  1024*1024
#define GIG  1024*1042*1024
#define SECONDS 5
/* in megabytes */
#define MEMORY_TO_LOCK 1024

int quiet;

int error() {
    switch (errno) {
        case ENOMEM:
            perror("Not enough Memory");
            break;
        case EAGAIN:
            perror("Not Enough Free Memory");
            break;
        case ENOSYS:
            perror("Sorry system doesn't support memory Locking");
            break;
        case EPERM:
            perror("Permission Denied");
            break;
        case ERANGE:
            perror("Exceed Range in -s argument");
            exit(errno);
            break;
        case EINVAL:
            perror("We dont know that crazy base you are asking for");
            exit(errno);
        default:
            perror("Unknown Error!");
    }
    exit(-1);
}

void fatal() {
    exit(errno);
}

void print(char *string) {
    if (quiet) return;
    printf("%s", string);
    return;
}

void printflush(char *string) {
    if (quiet) return;
    printf("%s", string);
    fflush(stdout);
}

void help() {
    print("mLock's locks lots of ram sleeps for 5 seconds unlocks the ram, exits\n\n" \
         "-q quiet\t -h help (this message)\n" \
         "-m <value> amount of memory to lock in megabytes\n" \
         "-s <value> time to sleep holding the memory lock\n");
    exit(0);
}

int main(int argc, char** argv) {
    char *buffer;
    int psize, i, seconds, memory_to_lock, ml;
    seconds = SECONDS;
    memory_to_lock = MEMORY_TO_LOCK;
    /* printf("Argc == %d\n", argc);
    printf("Argv 1 == %s\n", argv[1]); */
    if (argc > 1) {
        int count = argc;
        i = 1;
        if (argv[i][0] == '-') {
            switch (argv[i][1]) {
                case 'h':
                    help();
                    break;
                case 'q':
                    quiet = 1;
                    break;
                case 'm':
                    errno = 0;
                    memory_to_lock = atoi(argv[++i]);
                    if (errno) error();
                    if (memory_to_lock < 1) {
                        printf("Please enter a number greater than 0 we can't unlock memory we dont own\n");
                        fatal();
                    }
                    printf("locking %d megabytes of memory\n", memory_to_lock);
                    break;
                case 's':
                    errno = 0;
                    seconds = atoi(argv[++i]);
                    if (errno) error();
                    printf("seconds == %d\n", seconds);
                    break;
                default:
                    print(argv[i]);
                    perror(" has not been implemented yet!");
                    exit(-1);
            }
        }
    }

    psize = getpagesize();
    ml = memory_to_lock * MEG;
    buffer = memalign(psize, ml);
    if (!buffer)
        error();
    print("Memory Allocated\n");
    int retval;
    if (retval = mlock(buffer, ml))
        error();
    print("Memory is now locked\n");
    if (seconds > 0) {
        print("Sleeping");
        for ( i = 0; i < seconds; i++) {
            printflush(".");
            sleep(1);
        }
        print("\n");
    } else {
        print("okay no sleep for us\n");
    }
    if (retval = munlock(buffer, ml))
        error();
    print("Memory is now unlocked\n");
    free(buffer);
    print("Memory freed\n");
    return (EXIT_SUCCESS);
}