/*
 * 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 <malloc.h>
#include <unistd.h>
#include <errno.h>
#include <unistd.h>
#include <kstat.h>
#include <inttypes.h>
#include <sys/mman.h>
#include <sys/param.h>


#define MEG  (1024*1024)
#define GIG  1024*1042*1024
#define SECONDS 5
/* don't lock more than 50% of our memory, seems like a good number adjudst as needed */
#define MAX_PERCENT_LOCK 60
/* in megabytes */
#define MEMORY_TO_LOCK 1024
/* we should tune DESIRED_FREEMEM so that its enough to keep VDI happy */
/* but yet not effect other applications */
#define DESIRED_FREEMEM 1500

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);
}

int64_t meg_to_pages(unsigned int megs) {
    return (MEG * megs) / getpagesize();
}

int64_t get_64bit_kstat(kstat_t *ksp, char *name) {
    kstat_named_t *knp;
    knp = kstat_data_lookup(ksp, name);
    if (!knp) {
        printf("kstat_data_lookup  failed for %s\n", name);
        error();
    }
    return knp->value.i64;
}

void print_pages_k(char *title, int64_t num) {
    int kps;
    kps = getpagesize() / 1024;
    printf("\t%s = %" PRId64 " pages (%" PRId64 "k) \n", title, num, num * kps);
}

/* print all we know about memory, this routine fetches all the information again */

void print_memory_data() {
    int psize, err;
    /* kstats */
    /* http://developers.sun.com/solaris/articles/kstatc.html */
    kstat_ctl_t *kc;
    kstat_t *ksp;
    int64_t pageslocked, pagestotal, lotsfree, pagesfree, desired_free_pages;
    float percent_locked;
    psize = getpagesize();
    kc = kstat_open();
    ksp = kstat_lookup(kc, "unix", 0, "system_pages");
    err = kstat_read(kc, ksp, 0);
    if (!err) {
        printf("kstat_read error\n");
        exit(-1);
    }
    pagesfree = get_64bit_kstat(ksp, "pagesfree");
    pageslocked = get_64bit_kstat(ksp, "pageslocked");
    pagestotal = get_64bit_kstat(ksp, "pagestotal");
    lotsfree = get_64bit_kstat(ksp, "lotsfree");
    percent_locked = ((float) pageslocked / (float) pagestotal) * 100;
    desired_free_pages = meg_to_pages(DESIRED_FREEMEM);

    printf("Memory Stats\n");
    print_pages_k("Pages Total", pagestotal);
    print_pages_k("Locked Pages", pageslocked);
    print_pages_k("LotsFree", lotsfree);
    print_pages_k("Pagesize", 1);
    print_pages_k("FreePages", pagesfree);
    printf("\tPercent Locked  %.2f%%\n", percent_locked);
    print_pages_k("Total Memory", pagestotal);
    print_pages_k("Locked Memory", pageslocked);
    print_pages_k("Free Memory", pagesfree);
    print_pages_k("Desired Free Memory", desired_free_pages);
    return;
}

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" \
         "-M show memory statisics"
         "-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, err, retval;
    /* kstats */
    /* http://developers.sun.com/solaris/articles/kstatc.html */
    kstat_ctl_t *kc;
    kstat_t *ksp;
    int64_t pageslocked, pagestotal, pagesfree, pages_to_lock, desired_free_pages;
    float percent_locked;
    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;
        count = argc;
        i = 1;
        if (argv[i][0] == '-') {
            switch (argv[i][1]) {
                case 'h':
                    help();
                    break;
                case 'q':
                    quiet = 1;
                    break;
                case 'M':
                    print_memory_data();
                    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();
    /* open kstat */
    kc = kstat_open();
    ksp = kstat_lookup(kc, "unix", 0, "system_pages");
    err = kstat_read(kc, ksp, 0);
    if (!err) {
        printf("kstat_read error\n");
        exit(-1);
    }
    pagesfree = get_64bit_kstat(ksp, "pagesfree");
    pageslocked = get_64bit_kstat(ksp, "pageslocked");
    pagestotal = get_64bit_kstat(ksp, "pagestotal");
    percent_locked = ((float) pageslocked / (float) pagestotal) * 100;
    if (percent_locked > MAX_PERCENT_LOCK) {
        print("Sorry, system allready exceeded max locked page threshold\n");
        print_memory_data();
        exit(-1);
    }
    desired_free_pages = meg_to_pages(DESIRED_FREEMEM);
    if (desired_free_pages <= pagesfree) {
        print_memory_data();
        print("already enough free memory, bye\n");
        exit(0);
    }
    if (desired_free_pages > pagestotal) {
        print("Desired free memory greater than Physical memory\n");
        exit(-1);
    }
    ml = memory_to_lock * MEG;
    pages_to_lock = meg_to_pages(memory_to_lock);
    buffer = memalign(psize, ml);
    if (!buffer)
        error();
    print("Memory Allocated\n");
    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);
}