// SPDX-License-Identifier: GPL-2.0
/*
 * Test to verify that total_bw value remains consistent across all CPUs
 * in different BPF program states.
 *
 * Copyright (C) 2025 NVIDIA Corporation.
 */
#include <bpf/bpf.h>
#include <errno.h>
#include <pthread.h>
#include <scx/common.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/wait.h>
#include <unistd.h>
#include "minimal.bpf.skel.h"
#include "scx_test.h"

#define MAX_CPUS 512
#define STRESS_DURATION_SEC 5

struct total_bw_ctx {
        struct minimal *skel;
        long baseline_bw[MAX_CPUS];
        int nr_cpus;
};

static void *cpu_stress_thread(void *arg)
{
        volatile int i;
        time_t end_time = time(NULL) + STRESS_DURATION_SEC;

        while (time(NULL) < end_time)
                for (i = 0; i < 1000000; i++)
                        ;

        return NULL;
}

/*
 * The first enqueue on a CPU causes the DL server to start, for that
 * reason run stressor threads in the hopes it schedules on all CPUs.
 */
static int run_cpu_stress(int nr_cpus)
{
        pthread_t *threads;
        int i, ret = 0;

        threads = calloc(nr_cpus, sizeof(pthread_t));
        if (!threads)
                return -ENOMEM;

        /* Create threads to run on each CPU */
        for (i = 0; i < nr_cpus; i++) {
                if (pthread_create(&threads[i], NULL, cpu_stress_thread, NULL)) {
                        ret = -errno;
                        fprintf(stderr, "Failed to create thread %d: %s\n", i, strerror(-ret));
                        break;
                }
        }

        /* Wait for all threads to complete */
        for (i = 0; i < nr_cpus; i++) {
                if (threads[i])
                        pthread_join(threads[i], NULL);
        }

        free(threads);
        return ret;
}

static int read_total_bw_values(long *bw_values, int max_cpus)
{
        FILE *fp;
        char line[256];
        int cpu_count = 0;

        fp = fopen("/sys/kernel/debug/sched/debug", "r");
        if (!fp) {
                SCX_ERR("Failed to open debug file");
                return -1;
        }

        while (fgets(line, sizeof(line), fp)) {
                char *bw_str = strstr(line, "total_bw");

                if (bw_str) {
                        bw_str = strchr(bw_str, ':');
                        if (bw_str) {
                                /* Only store up to max_cpus values */
                                if (cpu_count < max_cpus)
                                        bw_values[cpu_count] = atol(bw_str + 1);
                                cpu_count++;
                        }
                }
        }

        fclose(fp);
        return cpu_count;
}

static bool verify_total_bw_consistency(long *bw_values, int count)
{
        int i;
        long first_value;

        if (count <= 0)
                return false;

        first_value = bw_values[0];

        for (i = 1; i < count; i++) {
                if (bw_values[i] != first_value) {
                        SCX_ERR("Inconsistent total_bw: CPU0=%ld, CPU%d=%ld",
                                first_value, i, bw_values[i]);
                        return false;
                }
        }

        return true;
}

static int fetch_verify_total_bw(long *bw_values, int nr_cpus)
{
        int attempts = 0;
        int max_attempts = 10;
        int count;

        /*
         * The first enqueue on a CPU causes the DL server to start, for that
         * reason run stressor threads in the hopes it schedules on all CPUs.
         */
        if (run_cpu_stress(nr_cpus) < 0) {
                SCX_ERR("Failed to run CPU stress");
                return -1;
        }

        /* Try multiple times to get stable values */
        while (attempts < max_attempts) {
                count = read_total_bw_values(bw_values, nr_cpus);
                fprintf(stderr, "Read %d total_bw values (testing %d CPUs)\n", count, nr_cpus);
                /* If system has more CPUs than we're testing, that's OK */
                if (count < nr_cpus) {
                        SCX_ERR("Expected at least %d CPUs, got %d", nr_cpus, count);
                        attempts++;
                        sleep(1);
                        continue;
                }

                /* Only verify the CPUs we're testing */
                if (verify_total_bw_consistency(bw_values, nr_cpus)) {
                        fprintf(stderr, "Values are consistent: %ld\n", bw_values[0]);
                        return 0;
                }

                attempts++;
                sleep(1);
        }

        return -1;
}

static enum scx_test_status setup(void **ctx)
{
        struct total_bw_ctx *test_ctx;

        if (access("/sys/kernel/debug/sched/debug", R_OK) != 0) {
                fprintf(stderr, "Skipping test: debugfs sched/debug not accessible\n");
                return SCX_TEST_SKIP;
        }

        test_ctx = calloc(1, sizeof(*test_ctx));
        if (!test_ctx)
                return SCX_TEST_FAIL;

        test_ctx->nr_cpus = sysconf(_SC_NPROCESSORS_ONLN);
        if (test_ctx->nr_cpus <= 0) {
                free(test_ctx);
                return SCX_TEST_FAIL;
        }

        /* If system has more CPUs than MAX_CPUS, just test the first MAX_CPUS */
        if (test_ctx->nr_cpus > MAX_CPUS)
                test_ctx->nr_cpus = MAX_CPUS;

        /* Test scenario 1: BPF program not loaded */
        /* Read and verify baseline total_bw before loading BPF program */
        fprintf(stderr, "BPF prog initially not loaded, reading total_bw values\n");
        if (fetch_verify_total_bw(test_ctx->baseline_bw, test_ctx->nr_cpus) < 0) {
                SCX_ERR("Failed to get stable baseline values");
                free(test_ctx);
                return SCX_TEST_FAIL;
        }

        /* Load the BPF skeleton */
        test_ctx->skel = minimal__open();
        if (!test_ctx->skel) {
                free(test_ctx);
                return SCX_TEST_FAIL;
        }

        SCX_ENUM_INIT(test_ctx->skel);
        if (minimal__load(test_ctx->skel)) {
                minimal__destroy(test_ctx->skel);
                free(test_ctx);
                return SCX_TEST_FAIL;
        }

        *ctx = test_ctx;
        return SCX_TEST_PASS;
}

static enum scx_test_status run(void *ctx)
{
        struct total_bw_ctx *test_ctx = ctx;
        struct bpf_link *link;
        long loaded_bw[MAX_CPUS];
        long unloaded_bw[MAX_CPUS];
        int i;

        /* Test scenario 2: BPF program loaded */
        link = bpf_map__attach_struct_ops(test_ctx->skel->maps.minimal_ops);
        if (!link) {
                SCX_ERR("Failed to attach scheduler");
                return SCX_TEST_FAIL;
        }

        fprintf(stderr, "BPF program loaded, reading total_bw values\n");
        if (fetch_verify_total_bw(loaded_bw, test_ctx->nr_cpus) < 0) {
                SCX_ERR("Failed to get stable values with BPF loaded");
                bpf_link__destroy(link);
                return SCX_TEST_FAIL;
        }
        bpf_link__destroy(link);

        /* Test scenario 3: BPF program unloaded */
        fprintf(stderr, "BPF program unloaded, reading total_bw values\n");
        if (fetch_verify_total_bw(unloaded_bw, test_ctx->nr_cpus) < 0) {
                SCX_ERR("Failed to get stable values after BPF unload");
                return SCX_TEST_FAIL;
        }

        /* Verify all three scenarios have the same total_bw values */
        for (i = 0; i < test_ctx->nr_cpus; i++) {
                if (test_ctx->baseline_bw[i] != loaded_bw[i]) {
                        SCX_ERR("CPU%d: baseline_bw=%ld != loaded_bw=%ld",
                                i, test_ctx->baseline_bw[i], loaded_bw[i]);
                        return SCX_TEST_FAIL;
                }

                if (test_ctx->baseline_bw[i] != unloaded_bw[i]) {
                        SCX_ERR("CPU%d: baseline_bw=%ld != unloaded_bw=%ld",
                                i, test_ctx->baseline_bw[i], unloaded_bw[i]);
                        return SCX_TEST_FAIL;
                }
        }

        fprintf(stderr, "All total_bw values are consistent across all scenarios\n");
        return SCX_TEST_PASS;
}

static void cleanup(void *ctx)
{
        struct total_bw_ctx *test_ctx = ctx;

        if (test_ctx) {
                if (test_ctx->skel)
                        minimal__destroy(test_ctx->skel);
                free(test_ctx);
        }
}

struct scx_test total_bw = {
        .name = "total_bw",
        .description = "Verify total_bw consistency across BPF program states",
        .setup = setup,
        .run = run,
        .cleanup = cleanup,
};
REGISTER_SCX_TEST(&total_bw)