// SPDX-License-Identifier: BSD-2-Clause

/* Copyright (C) 2014 - 2020 Intel Corporation. */

#include <memkind/internal/memkind_hugetlb.h>

#include <memkind/internal/memkind_default.h>

#include <memkind/internal/memkind_arena.h>

#include <memkind/internal/memkind_private.h>

#include <memkind/internal/memkind_log.h>

#include <sys/mman.h>

#ifndef MAP_HUGETLB

#define MAP_HUGETLB 0x40000

#endif

#ifndef MAP_HUGE_2MB

#define MAP_HUGE_2MB (21 << 26)

#endif

#include <stdio.h>

#include <errno.h>

#include <numa.h>

#include <pthread.h>

#include <dirent.h>

MEMKIND_EXPORT struct memkind_ops MEMKIND_HUGETLB_OPS = {

    .create = memkind_arena_create,

    .destroy = memkind_default_destroy,

    .malloc = memkind_arena_malloc,

    .calloc = memkind_arena_calloc,

    .posix_memalign = memkind_arena_posix_memalign,

    .realloc = memkind_arena_realloc,

    .free = memkind_arena_free,

    .check_available = memkind_hugetlb_check_available_2mb,

    .get_mmap_flags = memkind_hugetlb_get_mmap_flags,

    .get_arena = memkind_thread_get_arena,

    .init_once = memkind_hugetlb_init_once,

    .malloc_usable_size = memkind_default_malloc_usable_size,

    .finalize = memkind_arena_finalize,

    .get_stat = memkind_arena_get_kind_stat,

    .defrag_reallocate = memkind_arena_defrag_reallocate

};

static int get_nr_overcommit_hugepages_cached(size_t pagesize, size_t *out);

static int get_nr_hugepages_cached(size_t pagesize, struct bitmask *nodemask,

                                   size_t *out);

static int memkind_hugetlb_check_available(struct memkind *kind,

                                           size_t huge_size);

MEMKIND_EXPORT int memkind_hugetlb_get_mmap_flags(struct memkind *kind,

                                                  int *flags)

{

    *flags = MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB | MAP_HUGE_2MB;

    return 0;

}

MEMKIND_EXPORT void memkind_hugetlb_init_once(void)

{

    memkind_init(MEMKIND_HUGETLB, true);

}

MEMKIND_EXPORT int memkind_hugetlb_check_available_2mb(struct memkind *kind)

{

    return memkind_hugetlb_check_available(kind, 2097152);

}

/* huge_size: the huge page size in bytes */

static int memkind_hugetlb_check_available(struct memkind *kind,

                                           size_t huge_size)

{

    int err = 0;

    nodemask_t nodemask;

    struct bitmask nodemask_bm = {NUMA_NUM_NODES, nodemask.n};

    /* on x86_64 default huge page size is 2MB */

    if (huge_size == 0) {

        huge_size = 2097152;

    }

    if (kind->ops->get_mbind_nodemask) {

        err = kind->ops->get_mbind_nodemask(kind, nodemask.n, NUMA_NUM_NODES);

    } else {

        numa_bitmask_setall(&nodemask_bm);

    }

    size_t nr_persistent_hugepages, nr_overcommit_hugepages;

    err = get_nr_hugepages_cached(huge_size, &nodemask_bm,

                                  &nr_persistent_hugepages);

    if(err) {

        return err;

    }

    err = get_nr_overcommit_hugepages_cached(huge_size, &nr_overcommit_hugepages);

    if(err) {

        return err;

    }

    if (!nr_overcommit_hugepages && !nr_persistent_hugepages) {

        log_err("Persistent hugepages and overcommit hugepages are not available.");

        return MEMKIND_ERROR_HUGETLB;

    }

    return err;

}

struct hugepage_size_info {

    size_t size;

    size_t *nr_hugepages_per_node_array;

    size_t  nr_overcommit;

};

struct memkind_hugepages_config_t {

    struct hugepage_size_info **hugepages_info_array;

    int hugepages_info_array_len;

    int err; // 0 if sysfs parsing successful, appropriate memkind_error otherwise

} memkind_hugepages_config;

static pthread_once_t memkind_hugepages_config_once_g = PTHREAD_ONCE_INIT;

static struct hugepage_size_info *allocate_hugepage_size_info()

{

    struct hugepage_size_info *newInfo = malloc(sizeof(struct hugepage_size_info));

    if(newInfo == NULL) {

        log_err("malloc() failed.");

        return NULL;

    }

    newInfo->nr_hugepages_per_node_array = calloc(NUMA_NUM_NODES, sizeof(size_t));

    if(newInfo->nr_hugepages_per_node_array == NULL) {

        free(newInfo);

        log_err("calloc() failed.");

        return NULL;

    }

    return newInfo;

}

static size_t get_sysfs_entry_value(const char *entry_path)

{

    int errno_before;

    FILE *fid;

    int num_read;

    size_t value_read, ret = 0;

    errno_before = errno;

    fid = fopen(entry_path, "r");

    if (fid) {

        num_read = fscanf(fid, "%zud", &value_read);

        if(num_read) {

            ret  = value_read;

        }

        fclose(fid);

    } else {

        errno = errno_before;

    }

    return ret;

}

// construct hugepage_size_info object and fill it with data for provided pagesize

static void init_hugepage_size_info(size_t pagesize,

                                    struct hugepage_size_info *newInfo)

{

    char formatted_path[128];

    const char *nr_path_fmt =

        "/sys/devices/system/node/node%u/hugepages/hugepages-%zukB/nr_hugepages";

    const char *nr_overcommit_path_fmt =

        "/sys/kernel/mm/hugepages/hugepages-%zukB/nr_overcommit_hugepages";

    int snprintf_ret = 0;

    size_t node;

    size_t pagesize_kb = pagesize >> 10;

    newInfo->size = pagesize;

    //read overcommit hugepages limit for this pagesize

    snprintf_ret = snprintf(formatted_path, sizeof(formatted_path),

                            nr_overcommit_path_fmt, pagesize_kb);

    if (snprintf_ret > 0 && snprintf_ret < sizeof(formatted_path)) {

        newInfo->nr_overcommit = get_sysfs_entry_value(formatted_path);

        log_info("Overcommit limit for %zu kB hugepages is %zu.", pagesize,

                 newInfo->nr_overcommit);

    }

    //read every node nr_hugepages for this pagesize

    for (node = 0; node < NUMA_NUM_NODES; ++node) {

        snprintf_ret = snprintf(formatted_path, sizeof(formatted_path), nr_path_fmt,

                                node, pagesize_kb);

        if(snprintf_ret > 0 && snprintf_ret < sizeof(formatted_path)) {

            newInfo->nr_hugepages_per_node_array[node] = get_sysfs_entry_value(

                                                             formatted_path);

            if(node < numa_num_configured_nodes()) {

                log_info("Number of %zu kB hugepages on node %zu equals %zu.", pagesize, node,

                         newInfo->nr_hugepages_per_node_array[node]);

            }

        }

    }

}

// get hugepage size in bytes out of sysfs dir name

static int parse_pagesize_from_sysfs_entry(const char *entry, size_t *out)

{

    size_t pagesize;

    int ret = sscanf(entry, "hugepages-%zukB", &pagesize);

    if(ret == 1) {

        *out = pagesize << 10; //we are using bytes but kernel is using kB

        return 0;

    }

    return -1;

}

static void hugepages_config_init_once()

{

    unsigned j, i = 0;

    size_t pagesize;

    struct hugepage_size_info **hugepages_info_array = NULL;

    struct dirent *dir;

    DIR *hugepages_sysfs = opendir("/sys/kernel/mm/hugepages");

    if(hugepages_sysfs == NULL) {

        memkind_hugepages_config.err = MEMKIND_ERROR_HUGETLB;

        log_err("/sys/kernel/mm/hugepages directory is not available.");

        return;

    }

    unsigned hugepages_info_array_len = 2; //initial size of array

    hugepages_info_array = malloc(hugepages_info_array_len * sizeof(

                                      struct hugepage_size_info *));

    if (hugepages_info_array == NULL) {

        memkind_hugepages_config.err = MEMKIND_ERROR_MALLOC;

        closedir(hugepages_sysfs);

        log_err("malloc() failed.");

        return;

    }

    while ((dir = readdir(hugepages_sysfs)) != NULL) {

        if(dir->d_type == DT_DIR &&

           parse_pagesize_from_sysfs_entry(dir->d_name, &pagesize) == 0) {

            struct hugepage_size_info *new_hugepage_info = allocate_hugepage_size_info();

            if(new_hugepage_info == NULL) {

                memkind_hugepages_config.err = MEMKIND_ERROR_MALLOC;

                break;

            }

            init_hugepage_size_info(pagesize, new_hugepage_info);

            //there is more hugepage sizes than expected, reallocation of array needed

            if(i == hugepages_info_array_len) {

                hugepages_info_array_len *= 2;

                struct hugepage_size_info **swap_tmp = realloc(hugepages_info_array,

                                                               hugepages_info_array_len * sizeof(struct hugepage_size_info *));

                if(swap_tmp == NULL) {

                    free(new_hugepage_info);

                    memkind_hugepages_config.err = MEMKIND_ERROR_MALLOC;

                    log_err("realloc() failed.");

                    break;

                }

                hugepages_info_array = swap_tmp;

            }

            hugepages_info_array[i] = new_hugepage_info;

            i++;

        }

    }

    closedir(hugepages_sysfs);

    if(memkind_hugepages_config.err == 0) {

        memkind_hugepages_config.hugepages_info_array = hugepages_info_array;

        memkind_hugepages_config.hugepages_info_array_len = i;

    } else {

        for(j=0; j

            free(hugepages_info_array[i]);

        }

        free(hugepages_info_array);

    }

    return;

}

#ifdef __GNUC__

__attribute__((destructor))

#endif

static void destroy_hugepages_per_node()

{

    int i;

    for(i=0; i

        free(memkind_hugepages_config.hugepages_info_array[i]);

    }

    free(memkind_hugepages_config.hugepages_info_array);

}

// helper function that find and return hugepage_size_info object for specified pagesize

static struct hugepage_size_info *get_hugepage_info_for_pagesize(

    size_t pagesize)

{

    int i;

    for(i=0; i

        if(memkind_hugepages_config.hugepages_info_array[i]->size == pagesize) {

            return memkind_hugepages_config.hugepages_info_array[i];

        }

    }

    return NULL;

}

// returns sum of pre-allocated hugepage for specified pagesize and set of nodes

static int get_nr_hugepages_cached(size_t pagesize, struct bitmask *nodemask,

                                   size_t *out)

{

    int i;

    size_t nr_hugepages = 0;

    int num_node = numa_num_configured_nodes();

    pthread_once(&memkind_hugepages_config_once_g,

                 hugepages_config_init_once);

    if(memkind_hugepages_config.err != 0) {

        return memkind_hugepages_config.err;

    }

    struct hugepage_size_info *info = get_hugepage_info_for_pagesize(pagesize);

    if(info == NULL) {

        log_err("Unable to allocate hugepages, because info about pre-allocated hugepages is not available.");

        return MEMKIND_ERROR_HUGETLB;

    }

    for(i=0; i

        if(numa_bitmask_isbitset(nodemask, i)) {

            nr_hugepages += info->nr_hugepages_per_node_array[i];

        }

    }

    *out = nr_hugepages;

    return 0;

}

// returns hugepages overcommit limit for specified pagesize

static int get_nr_overcommit_hugepages_cached(size_t pagesize, size_t *out)

{

    pthread_once(&memkind_hugepages_config_once_g,

                 hugepages_config_init_once);

    if(memkind_hugepages_config.err != 0) {

        return memkind_hugepages_config.err;

    }

    struct hugepage_size_info *info = get_hugepage_info_for_pagesize(pagesize);

    if(info == NULL) {

        log_err("Unable to allocate hugepages, because info about overcommit hugepages is not available.");

        return MEMKIND_ERROR_HUGETLB;

    }

    *out = info->nr_overcommit;

    return 0;

}