// SPDX-License-Identifier: GPL-2.0
/*
 * KVM dirty logging page splitting test
 *
 * Based on dirty_log_perf.c
 *
 * Copyright (C) 2018, Red Hat, Inc.
 * Copyright (C) 2023, Google, Inc.
 */

#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <linux/bitmap.h>

#include "kvm_util.h"
#include "test_util.h"
#include "memstress.h"
#include "guest_modes.h"
#include "ucall_common.h"

#define VCPUS           2
#define SLOTS           2
#define ITERATIONS      2

static u64 guest_percpu_mem_size = DEFAULT_PER_VCPU_MEM_SIZE;

static enum vm_mem_backing_src_type backing_src = VM_MEM_SRC_ANONYMOUS_HUGETLB;

static u64 dirty_log_manual_caps;
static bool host_quit;
static int iteration;
static int vcpu_last_completed_iteration[KVM_MAX_VCPUS];

struct kvm_page_stats {
        u64 pages_4k;
        u64 pages_2m;
        u64 pages_1g;
        u64 hugepages;
};

static void get_page_stats(struct kvm_vm *vm, struct kvm_page_stats *stats, const char *stage)
{
        stats->pages_4k = vm_get_stat(vm, pages_4k);
        stats->pages_2m = vm_get_stat(vm, pages_2m);
        stats->pages_1g = vm_get_stat(vm, pages_1g);
        stats->hugepages = stats->pages_2m + stats->pages_1g;

        pr_debug("\nPage stats after %s: 4K: %ld 2M: %ld 1G: %ld huge: %ld\n",
                 stage, stats->pages_4k, stats->pages_2m, stats->pages_1g,
                 stats->hugepages);
}

static void run_vcpu_iteration(struct kvm_vm *vm)
{
        int i;

        iteration++;
        for (i = 0; i < VCPUS; i++) {
                while (READ_ONCE(vcpu_last_completed_iteration[i]) !=
                       iteration)
                        ;
        }
}

static void vcpu_worker(struct memstress_vcpu_args *vcpu_args)
{
        struct kvm_vcpu *vcpu = vcpu_args->vcpu;
        int vcpu_idx = vcpu_args->vcpu_idx;

        while (!READ_ONCE(host_quit)) {
                int current_iteration = READ_ONCE(iteration);

                vcpu_run(vcpu);

                TEST_ASSERT_EQ(get_ucall(vcpu, NULL), UCALL_SYNC);

                vcpu_last_completed_iteration[vcpu_idx] = current_iteration;

                /* Wait for the start of the next iteration to be signaled. */
                while (current_iteration == READ_ONCE(iteration) &&
                       READ_ONCE(iteration) >= 0 &&
                       !READ_ONCE(host_quit))
                        ;
        }
}

static void run_test(enum vm_guest_mode mode, void *unused)
{
        struct kvm_vm *vm;
        unsigned long **bitmaps;
        u64 guest_num_pages;
        u64 host_num_pages;
        u64 pages_per_slot;
        int i;
        struct kvm_page_stats stats_populated;
        struct kvm_page_stats stats_dirty_logging_enabled;
        struct kvm_page_stats stats_dirty_pass[ITERATIONS];
        struct kvm_page_stats stats_clear_pass[ITERATIONS];
        struct kvm_page_stats stats_dirty_logging_disabled;
        struct kvm_page_stats stats_repopulated;

        vm = memstress_create_vm(mode, VCPUS, guest_percpu_mem_size,
                                 SLOTS, backing_src, false);

        guest_num_pages = (VCPUS * guest_percpu_mem_size) >> vm->page_shift;
        guest_num_pages = vm_adjust_num_guest_pages(mode, guest_num_pages);
        host_num_pages = vm_num_host_pages(mode, guest_num_pages);
        pages_per_slot = host_num_pages / SLOTS;
        TEST_ASSERT_EQ(host_num_pages, pages_per_slot * SLOTS);
        TEST_ASSERT(!(host_num_pages % 512),
                    "Number of pages, '%lu' not a multiple of 2MiB", host_num_pages);

        bitmaps = memstress_alloc_bitmaps(SLOTS, pages_per_slot);

        if (dirty_log_manual_caps)
                vm_enable_cap(vm, KVM_CAP_MANUAL_DIRTY_LOG_PROTECT2,
                              dirty_log_manual_caps);

        /* Start the iterations */
        iteration = -1;
        host_quit = false;

        for (i = 0; i < VCPUS; i++)
                vcpu_last_completed_iteration[i] = -1;

        memstress_start_vcpu_threads(VCPUS, vcpu_worker);

        run_vcpu_iteration(vm);
        get_page_stats(vm, &stats_populated, "populating memory");

        /* Enable dirty logging */
        memstress_enable_dirty_logging(vm, SLOTS);

        get_page_stats(vm, &stats_dirty_logging_enabled, "enabling dirty logging");

        while (iteration < ITERATIONS) {
                run_vcpu_iteration(vm);
                get_page_stats(vm, &stats_dirty_pass[iteration - 1],
                               "dirtying memory");

                memstress_get_dirty_log(vm, bitmaps, SLOTS);

                if (dirty_log_manual_caps) {
                        memstress_clear_dirty_log(vm, bitmaps, SLOTS, pages_per_slot);

                        get_page_stats(vm, &stats_clear_pass[iteration - 1], "clearing dirty log");
                }
        }

        /* Disable dirty logging */
        memstress_disable_dirty_logging(vm, SLOTS);

        get_page_stats(vm, &stats_dirty_logging_disabled, "disabling dirty logging");

        /* Run vCPUs again to fault pages back in. */
        run_vcpu_iteration(vm);
        get_page_stats(vm, &stats_repopulated, "repopulating memory");

        /*
         * Tell the vCPU threads to quit.  No need to manually check that vCPUs
         * have stopped running after disabling dirty logging, the join will
         * wait for them to exit.
         */
        host_quit = true;
        memstress_join_vcpu_threads(VCPUS);

        memstress_free_bitmaps(bitmaps, SLOTS);
        memstress_destroy_vm(vm);

        TEST_ASSERT_EQ((stats_populated.pages_2m * 512 +
                        stats_populated.pages_1g * 512 * 512), host_num_pages);

        /*
         * Check that all huge pages were split. Since large pages can only
         * exist in the data slot, and the vCPUs should have dirtied all pages
         * in the data slot, there should be no huge pages left after splitting.
         * Splitting happens at dirty log enable time without
         * KVM_CAP_MANUAL_DIRTY_LOG_PROTECT2 and after the first clear pass
         * with that capability.
         */
        if (dirty_log_manual_caps) {
                TEST_ASSERT_EQ(stats_clear_pass[0].hugepages, 0);
                TEST_ASSERT(stats_clear_pass[0].pages_4k >= host_num_pages,
                            "Expected at least '%lu' 4KiB pages, found only '%lu'",
                            host_num_pages, stats_clear_pass[0].pages_4k);
                TEST_ASSERT_EQ(stats_dirty_logging_enabled.hugepages, stats_populated.hugepages);
        } else {
                TEST_ASSERT_EQ(stats_dirty_logging_enabled.hugepages, 0);
                TEST_ASSERT(stats_dirty_logging_enabled.pages_4k >= host_num_pages,
                            "Expected at least '%lu' 4KiB pages, found only '%lu'",
                            host_num_pages, stats_dirty_logging_enabled.pages_4k);
        }

        /*
         * Once dirty logging is disabled and the vCPUs have touched all their
         * memory again, the hugepage counts should be the same as they were
         * right after initial population of memory.
         */
        TEST_ASSERT_EQ(stats_populated.pages_2m, stats_repopulated.pages_2m);
        TEST_ASSERT_EQ(stats_populated.pages_1g, stats_repopulated.pages_1g);
}

static void help(char *name)
{
        puts("");
        printf("usage: %s [-h] [-b vcpu bytes] [-s mem type]\n",
               name);
        puts("");
        printf(" -b: specify the size of the memory region which should be\n"
               "     dirtied by each vCPU. e.g. 10M or 3G.\n"
               "     (default: 1G)\n");
        backing_src_help("-s");
        puts("");
}

int main(int argc, char *argv[])
{
        int opt;

        TEST_REQUIRE(get_kvm_param_bool("eager_page_split"));
        TEST_REQUIRE(get_kvm_param_bool("tdp_mmu"));

        while ((opt = getopt(argc, argv, "b:hs:")) != -1) {
                switch (opt) {
                case 'b':
                        guest_percpu_mem_size = parse_size(optarg);
                        break;
                case 'h':
                        help(argv[0]);
                        exit(0);
                case 's':
                        backing_src = parse_backing_src_type(optarg);
                        break;
                default:
                        help(argv[0]);
                        exit(1);
                }
        }

        if (!is_backing_src_hugetlb(backing_src)) {
                pr_info("This test will only work reliably with HugeTLB memory. "
                        "It can work with THP, but that is best effort.\n");
        }

        guest_modes_append_default();

        dirty_log_manual_caps = 0;
        for_each_guest_mode(run_test, NULL);

        dirty_log_manual_caps =
                kvm_check_cap(KVM_CAP_MANUAL_DIRTY_LOG_PROTECT2);

        if (dirty_log_manual_caps) {
                dirty_log_manual_caps &= (KVM_DIRTY_LOG_MANUAL_PROTECT_ENABLE |
                                          KVM_DIRTY_LOG_INITIALLY_SET);
                for_each_guest_mode(run_test, NULL);
        } else {
                pr_info("Skipping testing with MANUAL_PROTECT as it is not supported");
        }

        return 0;
}