// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) 2025, Google LLC.
 */

#include <time.h>

#include "lru_gen_util.h"

/*
 * Tracks state while we parse memcg lru_gen stats. The file we're parsing is
 * structured like this (some extra whitespace elided):
 *
 * memcg (id) (path)
 * node (id)
 * (gen_nr) (age_in_ms) (nr_anon_pages) (nr_file_pages)
 */
struct memcg_stats_parse_context {
        bool consumed; /* Whether or not this line was consumed */
        /* Next parse handler to invoke */
        void (*next_handler)(struct memcg_stats *stats,
                             struct memcg_stats_parse_context *ctx,
                             char *line);
        int current_node_idx; /* Current index in nodes array */
        const char *name; /* The name of the memcg we're looking for */
};

static void memcg_stats_handle_searching(struct memcg_stats *stats,
                                         struct memcg_stats_parse_context *ctx,
                                         char *line);
static void memcg_stats_handle_in_memcg(struct memcg_stats *stats,
                                        struct memcg_stats_parse_context *ctx,
                                        char *line);
static void memcg_stats_handle_in_node(struct memcg_stats *stats,
                                       struct memcg_stats_parse_context *ctx,
                                       char *line);

struct split_iterator {
        char *str;
        char *save;
};

static char *split_next(struct split_iterator *it)
{
        char *ret = strtok_r(it->str, " \t\n\r", &it->save);

        it->str = NULL;
        return ret;
}

static void memcg_stats_handle_searching(struct memcg_stats *stats,
                                         struct memcg_stats_parse_context *ctx,
                                         char *line)
{
        struct split_iterator it = { .str = line };
        char *prefix = split_next(&it);
        char *memcg_id = split_next(&it);
        char *memcg_name = split_next(&it);
        char *end;

        ctx->consumed = true;

        if (!prefix || strcmp("memcg", prefix))
                return; /* Not a memcg line (maybe empty), skip */

        TEST_ASSERT(memcg_id && memcg_name,
                    "malformed memcg line; no memcg id or memcg_name");

        if (strcmp(memcg_name + 1, ctx->name))
                return; /* Wrong memcg, skip */

        /* Found it! */

        stats->memcg_id = strtoul(memcg_id, &end, 10);
        TEST_ASSERT(*end == '\0', "malformed memcg id '%s'", memcg_id);
        if (!stats->memcg_id)
                return; /* Removed memcg? */

        ctx->next_handler = memcg_stats_handle_in_memcg;
}

static void memcg_stats_handle_in_memcg(struct memcg_stats *stats,
                                        struct memcg_stats_parse_context *ctx,
                                        char *line)
{
        struct split_iterator it = { .str = line };
        char *prefix = split_next(&it);
        char *id = split_next(&it);
        long found_node_id;
        char *end;

        ctx->consumed = true;
        ctx->current_node_idx = -1;

        if (!prefix)
                return; /* Skip empty lines */

        if (!strcmp("memcg", prefix)) {
                /* Memcg done, found next one; stop. */
                ctx->next_handler = NULL;
                return;
        } else if (strcmp("node", prefix))
                TEST_ASSERT(false, "found malformed line after 'memcg ...',"
                                   "token: '%s'", prefix);

        /* At this point we know we have a node line. Parse the ID. */

        TEST_ASSERT(id, "malformed node line; no node id");

        found_node_id = strtol(id, &end, 10);
        TEST_ASSERT(*end == '\0', "malformed node id '%s'", id);

        ctx->current_node_idx = stats->nr_nodes++;
        TEST_ASSERT(ctx->current_node_idx < MAX_NR_NODES,
                    "memcg has stats for too many nodes, max is %d",
                    MAX_NR_NODES);
        stats->nodes[ctx->current_node_idx].node = found_node_id;

        ctx->next_handler = memcg_stats_handle_in_node;
}

static void memcg_stats_handle_in_node(struct memcg_stats *stats,
                                       struct memcg_stats_parse_context *ctx,
                                       char *line)
{
        char *my_line = strdup(line);
        struct split_iterator it = { .str = my_line };
        char *gen, *age, *nr_anon, *nr_file;
        struct node_stats *node_stats;
        struct generation_stats *gen_stats;
        char *end;

        TEST_ASSERT(it.str, "failed to copy input line");

        gen = split_next(&it);

        if (!gen)
                goto out_consume; /* Skip empty lines */

        if (!strcmp("memcg", gen) || !strcmp("node", gen)) {
                /*
                 * Reached next memcg or node section. Don't consume, let the
                 * other handler deal with this.
                 */
                ctx->next_handler = memcg_stats_handle_in_memcg;
                goto out;
        }

        node_stats = &stats->nodes[ctx->current_node_idx];
        TEST_ASSERT(node_stats->nr_gens < MAX_NR_GENS,
                    "found too many generation lines; max is %d",
                    MAX_NR_GENS);
        gen_stats = &node_stats->gens[node_stats->nr_gens++];

        age = split_next(&it);
        nr_anon = split_next(&it);
        nr_file = split_next(&it);

        TEST_ASSERT(age && nr_anon && nr_file,
                    "malformed generation line; not enough tokens");

        gen_stats->gen = (int)strtol(gen, &end, 10);
        TEST_ASSERT(*end == '\0', "malformed generation number '%s'", gen);

        gen_stats->age_ms = strtol(age, &end, 10);
        TEST_ASSERT(*end == '\0', "malformed generation age '%s'", age);

        gen_stats->nr_anon = strtol(nr_anon, &end, 10);
        TEST_ASSERT(*end == '\0', "malformed anonymous page count '%s'",
                    nr_anon);

        gen_stats->nr_file = strtol(nr_file, &end, 10);
        TEST_ASSERT(*end == '\0', "malformed file page count '%s'", nr_file);

out_consume:
        ctx->consumed = true;
out:
        free(my_line);
}

static void print_memcg_stats(const struct memcg_stats *stats, const char *name)
{
        int node, gen;

        pr_debug("stats for memcg %s (id %lu):\n", name, stats->memcg_id);
        for (node = 0; node < stats->nr_nodes; ++node) {
                pr_debug("\tnode %d\n", stats->nodes[node].node);
                for (gen = 0; gen < stats->nodes[node].nr_gens; ++gen) {
                        const struct generation_stats *gstats =
                                &stats->nodes[node].gens[gen];

                        pr_debug("\t\tgen %d\tage_ms %ld"
                                 "\tnr_anon %ld\tnr_file %ld\n",
                                 gstats->gen, gstats->age_ms, gstats->nr_anon,
                                 gstats->nr_file);
                }
        }
}

/* Re-read lru_gen debugfs information for @memcg into @stats. */
void lru_gen_read_memcg_stats(struct memcg_stats *stats, const char *memcg)
{
        FILE *f;
        ssize_t read = 0;
        char *line = NULL;
        size_t bufsz;
        struct memcg_stats_parse_context ctx = {
                .next_handler = memcg_stats_handle_searching,
                .name = memcg,
        };

        memset(stats, 0, sizeof(struct memcg_stats));

        f = fopen(LRU_GEN_DEBUGFS, "r");
        TEST_ASSERT(f, "fopen(%s) failed", LRU_GEN_DEBUGFS);

        while (ctx.next_handler && (read = getline(&line, &bufsz, f)) > 0) {
                ctx.consumed = false;

                do {
                        ctx.next_handler(stats, &ctx, line);
                        if (!ctx.next_handler)
                                break;
                } while (!ctx.consumed);
        }

        if (read < 0 && !feof(f))
                TEST_ASSERT(false, "getline(%s) failed", LRU_GEN_DEBUGFS);

        TEST_ASSERT(stats->memcg_id > 0, "Couldn't find memcg: %s\n"
                    "Did the memcg get created in the proper mount?",
                    memcg);
        if (line)
                free(line);
        TEST_ASSERT(!fclose(f), "fclose(%s) failed", LRU_GEN_DEBUGFS);

        print_memcg_stats(stats, memcg);
}

/*
 * Find all pages tracked by lru_gen for this memcg in generation @target_gen.
 *
 * If @target_gen is negative, look for all generations.
 */
long lru_gen_sum_memcg_stats_for_gen(int target_gen,
                                     const struct memcg_stats *stats)
{
        int node, gen;
        long total_nr = 0;

        for (node = 0; node < stats->nr_nodes; ++node) {
                const struct node_stats *node_stats = &stats->nodes[node];

                for (gen = 0; gen < node_stats->nr_gens; ++gen) {
                        const struct generation_stats *gen_stats =
                                &node_stats->gens[gen];

                        if (target_gen >= 0 && gen_stats->gen != target_gen)
                                continue;

                        total_nr += gen_stats->nr_anon + gen_stats->nr_file;
                }
        }

        return total_nr;
}

/* Find all pages tracked by lru_gen for this memcg. */
long lru_gen_sum_memcg_stats(const struct memcg_stats *stats)
{
        return lru_gen_sum_memcg_stats_for_gen(-1, stats);
}

/*
 * If lru_gen aging should force page table scanning.
 *
 * If you want to set this to false, you will need to do eviction
 * before doing extra aging passes.
 */
static const bool force_scan = true;

static void run_aging_impl(unsigned long memcg_id, int node_id, int max_gen)
{
        FILE *f = fopen(LRU_GEN_DEBUGFS, "w");
        char *command;
        size_t sz;

        TEST_ASSERT(f, "fopen(%s) failed", LRU_GEN_DEBUGFS);
        sz = asprintf(&command, "+ %lu %d %d 1 %d\n",
                      memcg_id, node_id, max_gen, force_scan);
        TEST_ASSERT(sz > 0, "creating aging command failed");

        pr_debug("Running aging command: %s", command);
        if (fwrite(command, sizeof(char), sz, f) < sz) {
                TEST_ASSERT(false, "writing aging command %s to %s failed",
                            command, LRU_GEN_DEBUGFS);
        }

        TEST_ASSERT(!fclose(f), "fclose(%s) failed", LRU_GEN_DEBUGFS);
}

void lru_gen_do_aging(struct memcg_stats *stats, const char *memcg)
{
        int node, gen;

        pr_debug("lru_gen: invoking aging...\n");

        /* Must read memcg stats to construct the proper aging command. */
        lru_gen_read_memcg_stats(stats, memcg);

        for (node = 0; node < stats->nr_nodes; ++node) {
                int max_gen = 0;

                for (gen = 0; gen < stats->nodes[node].nr_gens; ++gen) {
                        int this_gen = stats->nodes[node].gens[gen].gen;

                        max_gen = max_gen > this_gen ? max_gen : this_gen;
                }

                run_aging_impl(stats->memcg_id, stats->nodes[node].node,
                               max_gen);
        }

        /* Re-read so callers get updated information */
        lru_gen_read_memcg_stats(stats, memcg);
}

/*
 * Find which generation contains at least @pages pages, assuming that
 * such a generation exists.
 */
int lru_gen_find_generation(const struct memcg_stats *stats,
                            unsigned long pages)
{
        int node, gen, gen_idx, min_gen = INT_MAX, max_gen = -1;

        for (node = 0; node < stats->nr_nodes; ++node)
                for (gen_idx = 0; gen_idx < stats->nodes[node].nr_gens;
                     ++gen_idx) {
                        gen = stats->nodes[node].gens[gen_idx].gen;
                        max_gen = gen > max_gen ? gen : max_gen;
                        min_gen = gen < min_gen ? gen : min_gen;
                }

        for (gen = min_gen; gen <= max_gen; ++gen)
                /* See if this generation has enough pages. */
                if (lru_gen_sum_memcg_stats_for_gen(gen, stats) > pages)
                        return gen;

        return -1;
}

bool lru_gen_usable(void)
{
        long required_features = LRU_GEN_ENABLED | LRU_GEN_MM_WALK;
        int lru_gen_fd, lru_gen_debug_fd;
        char mglru_feature_str[8] = {};
        long mglru_features;

        lru_gen_fd = open(LRU_GEN_ENABLED_PATH, O_RDONLY);
        if (lru_gen_fd < 0) {
                puts("lru_gen: Could not open " LRU_GEN_ENABLED_PATH);
                return false;
        }
        if (read(lru_gen_fd, &mglru_feature_str, 7) < 7) {
                puts("lru_gen: Could not read from " LRU_GEN_ENABLED_PATH);
                close(lru_gen_fd);
                return false;
        }
        close(lru_gen_fd);

        mglru_features = strtol(mglru_feature_str, NULL, 16);
        if ((mglru_features & required_features) != required_features) {
                printf("lru_gen: missing features, got: 0x%lx, expected: 0x%lx\n",
                       mglru_features, required_features);
                printf("lru_gen: Try 'echo 0x%lx > /sys/kernel/mm/lru_gen/enabled'\n",
                       required_features);
                return false;
        }

        lru_gen_debug_fd = open(LRU_GEN_DEBUGFS, O_RDWR);
        __TEST_REQUIRE(lru_gen_debug_fd >= 0,
                       "lru_gen: Could not open " LRU_GEN_DEBUGFS ", "
                       "but lru_gen is enabled, so cannot use page_idle.");
        close(lru_gen_debug_fd);
        return true;
}