// SPDX-License-Identifier: BSD-2-Clause
/* Copyright (C) 2017 - 2021 Intel Corporation. */

#include "allocator_perf_tool/Allocation_info.hpp"
#include "allocator_perf_tool/GTestAdapter.hpp"
#include "common.h"
#include "proc_stat.h"
#include "random_sizes_allocator.h"

#include <memkind.h>

#include <condition_variable>
#include <functional>
#include <mutex>
#include <random>
#include <thread>

extern const char *PMEM_DIR;
class Worker
{
public:
    Worker(RandomSizesAllocator &&allocator, double malloc_probability)
        : allocator(std::move(allocator)),
          malloc_probability(malloc_probability)
    {}

    void work()
    {
        if (allocator.empty() || get_random_bool(malloc_probability)) {
            requested_memory_sum += allocator.malloc_random_memory();
        } else {
            requested_memory_sum -= allocator.free_random_memory();
        }
    }

    size_t get_requested_memory_sum_bytes() const
    {
        return requested_memory_sum;
    }

private:
    bool get_random_bool(double probability)
    {
        std::bernoulli_distribution distribution(probability);
        return distribution(generator);
    }

    std::default_random_engine generator;
    RandomSizesAllocator allocator;
    size_t requested_memory_sum = 0;
    double malloc_probability;
};

class MemoryFootprintStats
{
public:
    void reset()
    {
        std::lock_guard<std::mutex> lk(sample_guard);
        initial_virtual_memory = proc_stat.get_virtual_memory_size_bytes();
        initial_physical_memory = proc_stat.get_physical_memory_size_bytes();
        vm_overhead_sum = 0;
        current_vm_overhead = 0;
        max_vm_overhead = 0;

        current_phys_overhead = 0;
        phys_overhead_sum = 0;
        max_phys_overhead = 0;

        requested_memory = 0;
        sample_count = 0;
    }

    void sample(long long requested_memory_bytes)
    {
        std::lock_guard<std::mutex> lk(sample_guard);
        current_vm = proc_stat.get_virtual_memory_size_bytes();
        current_phys = proc_stat.get_physical_memory_size_bytes();

        sample_count++;

        requested_memory = requested_memory_bytes;

        current_vm_overhead =
            current_vm - initial_virtual_memory - requested_memory;
        vm_overhead_sum += current_vm_overhead;
        max_vm_overhead = std::max(max_vm_overhead, current_vm_overhead);

        current_phys_overhead = current_phys - initial_physical_memory;
        phys_overhead_sum += current_phys_overhead;
        max_phys_overhead = std::max(max_phys_overhead, current_phys_overhead);
    }

    void log_data() const
    {
        std::lock_guard<std::mutex> lk(sample_guard);
        GTestAdapter::RecordProperty("avg_vm_overhead_per_operation_mb",
                                     convert_bytes_to_mb(vm_overhead_sum) /
                                         sample_count);
        GTestAdapter::RecordProperty("avg_vm_overhead_growth_per_operation_mb",
                                     convert_bytes_to_mb(current_vm_overhead) /
                                         sample_count);
        GTestAdapter::RecordProperty("max_vm_overhead_mb",
                                     convert_bytes_to_mb(max_vm_overhead));
        GTestAdapter::RecordProperty("avg_phys_overhead_per_operation_mb",
                                     convert_bytes_to_mb(phys_overhead_sum) /
                                         sample_count);
        GTestAdapter::RecordProperty("max_phys_overhead_mb",
                                     convert_bytes_to_mb(max_phys_overhead));
        GTestAdapter::RecordProperty(
            "overhead_to_requested_memory_ratio_percent",
            100.f * current_vm_overhead / requested_memory);
        GTestAdapter::RecordProperty("requested_memory_mb",
                                     convert_bytes_to_mb(requested_memory));
    }

private:
    long long initial_virtual_memory;
    long long vm_overhead_sum = 0;
    long long current_vm_overhead = 0;
    long long max_vm_overhead = 0;

    long long initial_physical_memory;
    long long current_phys_overhead = 0;
    long long phys_overhead_sum = 0;
    long long max_phys_overhead = 0;

    long long requested_memory = 0;
    long long sample_count = 0;

    long long current_vm;
    long long current_phys;

    ProcStat proc_stat;
    mutable std::mutex sample_guard;
};

/* Execute func calling it n_calls times in n_threads threads.
 * The execution is multithreaded but the func calls order is sequential.
 * func takes thread id, and operation id as an argument,
 * and must return thread id of the next thread, where thread ids are in range
 * <0, n_threads-1>. init_thread_id specify the initial thread id.
 */
void run_multithreaded_seq_exec(
    unsigned n_threads, unsigned init_thread_id, unsigned n_calls,
    std::function<unsigned(unsigned, unsigned)> func)
{
    std::vector<std::thread> threads;
    std::mutex mutex;
    std::condition_variable turns_holder;
    unsigned current_call = 0;
    unsigned current_tid = init_thread_id;

    threads.reserve(n_threads);

    mutex.lock();

    for (int tid = 0; tid < n_threads; ++tid) {
        threads.emplace_back([&, tid]() {
            while (current_call < n_calls) {
                std::unique_lock<std::mutex> lk(mutex);
                turns_holder.wait(lk, [&, tid] {
                    return current_tid == tid || current_call == n_calls;
                });
                if (current_call == n_calls) {
                    return;
                }
                current_tid = func(tid, current_call);
                ASSERT_LT(current_tid, n_threads) << "Incorrect thread id!";
                current_call++;
                lk.unlock();
                turns_holder.notify_all();
            }
        });
    }

    mutex.unlock();

    for (int i = 0; i < threads.size(); i++) {
        threads[i].join();
    }
}

/*
 * Create threads and measure the cost of maintaining allocations from threads.
 * Allocations order is sequential (otherwise the results might be very
 * nondeterministic).
 */
void run_test(memkind_t kind, size_t min_size, size_t max_size,
              unsigned n_threads, double malloc_probability = 1.0,
              unsigned n_calls = 1000)
{
    Worker worker(RandomSizesAllocator(kind, min_size, max_size, n_calls),
                  malloc_probability);

    MemoryFootprintStats mem_footprint_stats;

    auto func = [&](unsigned tid, unsigned id) -> unsigned {
        if (id == 0) {
            mem_footprint_stats.reset();
        }

        worker.work();
        mem_footprint_stats.sample(worker.get_requested_memory_sum_bytes());

        return (tid + 1) % n_threads; // next thread id
    };
    run_multithreaded_seq_exec(n_threads, 0, n_calls, func);
    mem_footprint_stats.log_data();
}

class MemoryFootprintTest: public ::testing::Test
{
};

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_DEFAULT_only_malloc_small_allocations_1_thread)
{
    run_test(MEMKIND_DEFAULT, 128, 15 * KB, 1);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_DEFAULT_only_malloc_small_allocations_10_thread)
{
    run_test(MEMKIND_DEFAULT, 128, 15 * KB, 10);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_DEFAULT_only_malloc_medium_allocations_1_thread)
{
    run_test(MEMKIND_DEFAULT, 16 * KB, 1 * MB, 1);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_DEFAULT_only_malloc_medium_allocations_10_thread)
{
    run_test(MEMKIND_DEFAULT, 16 * KB, 1 * MB, 10);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_DEFAULT_only_malloc_large_allocations_1_thread)
{
    run_test(MEMKIND_DEFAULT, 2 * MB, 10 * MB, 1, 1.0, 100);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_DEFAULT_only_malloc_large_allocations_10_thread)
{
    run_test(MEMKIND_DEFAULT, 2 * MB, 10 * MB, 10, 1.0, 100);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_DEFAULT_random_malloc80_free20_random_small_allocations_1_thread)
{
    run_test(MEMKIND_DEFAULT, 128, 15 * KB, 1, 0.8);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_DEFAULT_random_malloc80_free20_random_small_allocations_10_thread)
{
    run_test(MEMKIND_DEFAULT, 128, 15 * KB, 10, 0.8);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_DEFAULT_random_malloc80_free20_random_medium_allocations_1_thread)
{
    run_test(MEMKIND_DEFAULT, 16 * KB, 1 * MB, 1, 0.8);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_DEFAULT_random_malloc80_free20_random_medium_allocations_10_thread)
{
    run_test(MEMKIND_DEFAULT, 16 * KB, 1 * MB, 10, 0.8);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_DEFAULT_random_malloc80_free20_random_large_allocations_1_thread)
{
    run_test(MEMKIND_DEFAULT, 2 * MB, 10 * MB, 1, 0.8, 100);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_DEFAULT_random_malloc80_free20_random_large_allocations_10_thread)
{
    run_test(MEMKIND_DEFAULT, 2 * MB, 10 * MB, 10, 0.8, 100);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_HBW_only_malloc_small_allocations_1_thread)
{
    run_test(MEMKIND_HBW, 128, 15 * KB, 1);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_HBW_only_malloc_small_allocations_10_thread)
{
    run_test(MEMKIND_HBW, 128, 15 * KB, 10);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_HBW_only_malloc_medium_allocations_1_thread)
{
    run_test(MEMKIND_HBW, 16 * KB, 1 * MB, 1);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_HBW_only_malloc_medium_allocations_10_thread)
{
    run_test(MEMKIND_HBW, 16 * KB, 1 * MB, 10);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_HBW_only_malloc_large_allocations_1_thread)
{
    run_test(MEMKIND_HBW, 2 * MB, 10 * MB, 1, 100);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_HBW_only_malloc_large_allocations_10_thread)
{
    run_test(MEMKIND_HBW, 2 * MB, 10 * MB, 10, 100);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_HBW_random_malloc80_free20_random_small_allocations_1_thread)
{
    run_test(MEMKIND_HBW, 128, 15 * KB, 1, 0.8);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_HBW_random_malloc80_free20_random_small_allocations_10_thread)
{
    run_test(MEMKIND_HBW, 128, 15 * KB, 10, 0.8);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_HBW_random_malloc80_free20_random_medium_allocations_1_thread)
{
    run_test(MEMKIND_HBW, 16 * KB, 1 * MB, 1, 0.8);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_HBW_random_malloc80_free20_random_medium_allocations_10_thread)
{
    run_test(MEMKIND_HBW, 16 * KB, 1 * MB, 10, 0.8);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_HBW_random_malloc80_free20_random_large_allocations_1_thread)
{
    run_test(MEMKIND_HBW, 2 * MB, 10 * MB, 1, 0.8, 100);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_HBW_random_malloc80_free20_random_large_allocations_10_thread)
{
    run_test(MEMKIND_HBW, 2 * MB, 10 * MB, 10, 0.8, 100);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_DAX_KMEM_only_malloc_small_allocations_1_thread)
{
    run_test(MEMKIND_DAX_KMEM, 128, 15 * KB, 1);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_DAX_KMEM_only_malloc_small_allocations_10_thread)
{
    run_test(MEMKIND_DAX_KMEM, 128, 15 * KB, 10);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_DAX_KMEM_only_malloc_medium_allocations_1_thread)
{
    run_test(MEMKIND_DAX_KMEM, 16 * KB, 1 * MB, 1);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_DAX_KMEM_only_malloc_medium_allocations_10_thread)
{
    run_test(MEMKIND_DAX_KMEM, 16 * KB, 1 * MB, 10);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_DAX_KMEM_only_malloc_large_allocations_1_thread)
{
    run_test(MEMKIND_DAX_KMEM, 2 * MB, 10 * MB, 1, 100);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_DAX_KMEM_only_malloc_large_allocations_10_thread)
{
    run_test(MEMKIND_DAX_KMEM, 2 * MB, 10 * MB, 10, 100);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_DAX_KMEM_random_malloc80_free20_random_small_allocations_1_thread)
{
    run_test(MEMKIND_DAX_KMEM, 128, 15 * KB, 1, 0.8);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_DAX_KMEM_random_malloc80_free20_random_small_allocations_10_thread)
{
    run_test(MEMKIND_DAX_KMEM, 128, 15 * KB, 10, 0.8);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_DAX_KMEM_random_malloc80_free20_random_medium_allocations_1_thread)
{
    run_test(MEMKIND_DAX_KMEM, 16 * KB, 1 * MB, 1, 0.8);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_DAX_KMEM_random_malloc80_free20_random_medium_allocations_10_thread)
{
    run_test(MEMKIND_DAX_KMEM, 16 * KB, 1 * MB, 10, 0.8);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_DAX_KMEM_random_malloc80_free20_random_large_allocations_1_thread)
{
    run_test(MEMKIND_DAX_KMEM, 2 * MB, 10 * MB, 1, 0.8, 100);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_DAX_KMEM_random_malloc80_free20_random_large_allocations_10_thread)
{
    run_test(MEMKIND_DAX_KMEM, 2 * MB, 10 * MB, 10, 0.8, 100);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_MEMKIND_PMEM_only_malloc_small_allocations_1_thread)
{
    struct memkind *kind = nullptr;
    memkind_create_pmem(PMEM_DIR, 0, &kind);
    run_test(kind, 128, 15 * KB, 1);
    memkind_destroy_kind(kind);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_MEMKIND_PMEM_only_malloc_small_allocations_10_thread)
{
    struct memkind *kind = nullptr;
    memkind_create_pmem(PMEM_DIR, 0, &kind);
    run_test(kind, 128, 15 * KB, 10);
    memkind_destroy_kind(kind);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_MEMKIND_PMEM_only_malloc_medium_allocations_1_thread)
{
    struct memkind *kind = nullptr;
    memkind_create_pmem(PMEM_DIR, 0, &kind);
    run_test(kind, 16 * KB, 1 * MB, 1);
    memkind_destroy_kind(kind);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_MEMKIND_PMEM_only_malloc_medium_allocations_10_thread)
{
    struct memkind *kind = nullptr;
    memkind_create_pmem(PMEM_DIR, 0, &kind);
    run_test(kind, 16 * KB, 1 * MB, 10);
    memkind_destroy_kind(kind);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_MEMKIND_PMEM_only_malloc_large_allocations_1_thread)
{
    struct memkind *kind = nullptr;
    memkind_create_pmem(PMEM_DIR, 0, &kind);
    run_test(kind, 2 * MB, 10 * MB, 1, 100);
    memkind_destroy_kind(kind);
}

TEST_F(MemoryFootprintTest,
       test_TC_MEMKIND_MEMKIND_PMEM_only_malloc_large_allocations_10_thread)
{
    struct memkind *kind = nullptr;
    memkind_create_pmem(PMEM_DIR, 0, &kind);
    run_test(kind, 2 * MB, 10 * MB, 10, 100);
    memkind_destroy_kind(kind);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_MEMKIND_PMEM_random_malloc80_free20_random_small_allocations_1_thread)
{
    struct memkind *kind = nullptr;
    memkind_create_pmem(PMEM_DIR, 0, &kind);
    run_test(kind, 128, 15 * KB, 1, 0.8);
    memkind_destroy_kind(kind);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_MEMKIND_PMEM_random_malloc80_free20_random_small_allocations_10_thread)
{
    struct memkind *kind = nullptr;
    memkind_create_pmem(PMEM_DIR, 0, &kind);
    run_test(kind, 128, 15 * KB, 10, 0.8);
    memkind_destroy_kind(kind);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_MEMKIND_PMEM_random_malloc80_free20_random_medium_allocations_1_thread)
{
    struct memkind *kind = nullptr;
    memkind_create_pmem(PMEM_DIR, 0, &kind);
    run_test(kind, 16 * KB, 1 * MB, 1, 0.8);
    memkind_destroy_kind(kind);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_MEMKIND_PMEM_random_malloc80_free20_random_medium_allocations_10_thread)
{
    struct memkind *kind = nullptr;
    memkind_create_pmem(PMEM_DIR, 0, &kind);
    run_test(kind, 16 * KB, 1 * MB, 10, 0.8);
    memkind_destroy_kind(kind);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_MEMKIND_PMEM_random_malloc80_free20_random_large_allocations_1_thread)
{
    struct memkind *kind = nullptr;
    memkind_create_pmem(PMEM_DIR, 0, &kind);
    run_test(kind, 2 * MB, 10 * MB, 1, 0.8, 100);
    memkind_destroy_kind(kind);
}

TEST_F(
    MemoryFootprintTest,
    test_TC_MEMKIND_MEMKIND_PMEM_random_malloc80_free20_random_large_allocations_10_thread)
{
    struct memkind *kind = nullptr;
    memkind_create_pmem(PMEM_DIR, 0, &kind);
    run_test(kind, 2 * MB, 10 * MB, 10, 0.8, 100);
    memkind_destroy_kind(kind);
}
