// SPDX-License-Identifier: GPL-2.0
#include <linux/ceph/ceph_debug.h>

#include <linux/math64.h>
#include <linux/slab.h>
#include <linux/seq_file.h>

#include "subvolume_metrics.h"
#include "mds_client.h"
#include "super.h"

/**
 * struct ceph_subvol_metric_rb_entry - Per-subvolume I/O metrics node
 * @node: Red-black tree linkage for tracker->tree
 * @subvolume_id: Subvolume identifier (key for rb-tree lookup)
 * @read_ops: Accumulated read operation count since last snapshot
 * @write_ops: Accumulated write operation count since last snapshot
 * @read_bytes: Accumulated bytes read since last snapshot
 * @write_bytes: Accumulated bytes written since last snapshot
 * @read_latency_us: Sum of read latencies in microseconds
 * @write_latency_us: Sum of write latencies in microseconds
 */
struct ceph_subvol_metric_rb_entry {
        struct rb_node node;
        u64 subvolume_id;
        u64 read_ops;
        u64 write_ops;
        u64 read_bytes;
        u64 write_bytes;
        u64 read_latency_us;
        u64 write_latency_us;
};

static struct kmem_cache *ceph_subvol_metric_entry_cachep;

void ceph_subvolume_metrics_init(struct ceph_subvolume_metrics_tracker *tracker)
{
        spin_lock_init(&tracker->lock);
        tracker->tree = RB_ROOT_CACHED;
        tracker->nr_entries = 0;
        tracker->enabled = false;
        atomic64_set(&tracker->snapshot_attempts, 0);
        atomic64_set(&tracker->snapshot_empty, 0);
        atomic64_set(&tracker->snapshot_failures, 0);
        atomic64_set(&tracker->record_calls, 0);
        atomic64_set(&tracker->record_disabled, 0);
        atomic64_set(&tracker->record_no_subvol, 0);
        atomic64_set(&tracker->total_read_ops, 0);
        atomic64_set(&tracker->total_read_bytes, 0);
        atomic64_set(&tracker->total_write_ops, 0);
        atomic64_set(&tracker->total_write_bytes, 0);
}

static struct ceph_subvol_metric_rb_entry *
__lookup_entry(struct ceph_subvolume_metrics_tracker *tracker, u64 subvol_id)
{
        struct rb_node *node;

        node = tracker->tree.rb_root.rb_node;
        while (node) {
                struct ceph_subvol_metric_rb_entry *entry =
                        rb_entry(node, struct ceph_subvol_metric_rb_entry, node);

                if (subvol_id < entry->subvolume_id)
                        node = node->rb_left;
                else if (subvol_id > entry->subvolume_id)
                        node = node->rb_right;
                else
                        return entry;
        }

        return NULL;
}

static struct ceph_subvol_metric_rb_entry *
__insert_entry(struct ceph_subvolume_metrics_tracker *tracker,
               struct ceph_subvol_metric_rb_entry *entry)
{
        struct rb_node **link = &tracker->tree.rb_root.rb_node;
        struct rb_node *parent = NULL;
        bool leftmost = true;

        while (*link) {
                struct ceph_subvol_metric_rb_entry *cur =
                        rb_entry(*link, struct ceph_subvol_metric_rb_entry, node);

                parent = *link;
                if (entry->subvolume_id < cur->subvolume_id)
                        link = &(*link)->rb_left;
                else if (entry->subvolume_id > cur->subvolume_id) {
                        link = &(*link)->rb_right;
                        leftmost = false;
                } else
                        return cur;
        }

        rb_link_node(&entry->node, parent, link);
        rb_insert_color_cached(&entry->node, &tracker->tree, leftmost);
        tracker->nr_entries++;
        return entry;
}

static void ceph_subvolume_metrics_clear_locked(
                struct ceph_subvolume_metrics_tracker *tracker)
{
        struct rb_node *node = rb_first_cached(&tracker->tree);

        while (node) {
                struct ceph_subvol_metric_rb_entry *entry =
                        rb_entry(node, struct ceph_subvol_metric_rb_entry, node);
                struct rb_node *next = rb_next(node);

                rb_erase_cached(&entry->node, &tracker->tree);
                tracker->nr_entries--;
                kmem_cache_free(ceph_subvol_metric_entry_cachep, entry);
                node = next;
        }

        tracker->tree = RB_ROOT_CACHED;
}

void ceph_subvolume_metrics_destroy(struct ceph_subvolume_metrics_tracker *tracker)
{
        spin_lock(&tracker->lock);
        ceph_subvolume_metrics_clear_locked(tracker);
        tracker->enabled = false;
        spin_unlock(&tracker->lock);
}

void ceph_subvolume_metrics_enable(struct ceph_subvolume_metrics_tracker *tracker,
                                   bool enable)
{
        spin_lock(&tracker->lock);
        if (enable) {
                tracker->enabled = true;
        } else {
                tracker->enabled = false;
                ceph_subvolume_metrics_clear_locked(tracker);
        }
        spin_unlock(&tracker->lock);
}

void ceph_subvolume_metrics_record(struct ceph_subvolume_metrics_tracker *tracker,
                                   u64 subvol_id, bool is_write,
                                   size_t size, u64 latency_us)
{
        struct ceph_subvol_metric_rb_entry *entry, *new_entry = NULL;
        bool retry = false;

        /* CEPH_SUBVOLUME_ID_NONE (0) means unknown/unset subvolume */
        if (!READ_ONCE(tracker->enabled) ||
            subvol_id == CEPH_SUBVOLUME_ID_NONE || !size || !latency_us)
                return;

        /*
         * Retry loop for lock-free allocation pattern:
         * 1. First iteration: lookup under lock, if miss -> drop lock, alloc, retry
         * 2. Second iteration: lookup again (may have been inserted), insert if still missing
         * 3. On race (another thread inserted same key): free our alloc, retry
         * All successful paths exit via return, so retry flag doesn't need reset.
         */
        do {
                spin_lock(&tracker->lock);
                if (!tracker->enabled) {
                        spin_unlock(&tracker->lock);
                        if (new_entry)
                                kmem_cache_free(ceph_subvol_metric_entry_cachep, new_entry);
                        return;
                }

                entry = __lookup_entry(tracker, subvol_id);
                if (!entry) {
                        if (!new_entry) {
                                spin_unlock(&tracker->lock);
                                new_entry = kmem_cache_zalloc(ceph_subvol_metric_entry_cachep,
                                                      GFP_NOFS);
                                if (!new_entry)
                                        return;
                                new_entry->subvolume_id = subvol_id;
                                retry = true;
                                continue;
                        }
                        entry = __insert_entry(tracker, new_entry);
                        if (entry != new_entry) {
                                /* raced with another insert */
                                spin_unlock(&tracker->lock);
                                kmem_cache_free(ceph_subvol_metric_entry_cachep, new_entry);
                                new_entry = NULL;
                                retry = true;
                                continue;
                        }
                        new_entry = NULL;
                }

                if (is_write) {
                        entry->write_ops++;
                        entry->write_bytes += size;
                        entry->write_latency_us += latency_us;
                        atomic64_inc(&tracker->total_write_ops);
                        atomic64_add(size, &tracker->total_write_bytes);
                } else {
                        entry->read_ops++;
                        entry->read_bytes += size;
                        entry->read_latency_us += latency_us;
                        atomic64_inc(&tracker->total_read_ops);
                        atomic64_add(size, &tracker->total_read_bytes);
                }
                spin_unlock(&tracker->lock);
                if (new_entry)
                        kmem_cache_free(ceph_subvol_metric_entry_cachep, new_entry);
                return;
        } while (retry);
}

int ceph_subvolume_metrics_snapshot(struct ceph_subvolume_metrics_tracker *tracker,
                                    struct ceph_subvol_metric_snapshot **out,
                                    u32 *nr, bool consume)
{
        struct ceph_subvol_metric_snapshot *snap = NULL;
        struct rb_node *node;
        u32 count = 0, idx = 0;
        int ret = 0;

        *out = NULL;
        *nr = 0;

        if (!READ_ONCE(tracker->enabled))
                return 0;

        atomic64_inc(&tracker->snapshot_attempts);

        spin_lock(&tracker->lock);
        for (node = rb_first_cached(&tracker->tree); node; node = rb_next(node)) {
                struct ceph_subvol_metric_rb_entry *entry =
                        rb_entry(node, struct ceph_subvol_metric_rb_entry, node);

                /* Include entries with ANY I/O activity (read OR write) */
                if (entry->read_ops || entry->write_ops)
                        count++;
        }
        spin_unlock(&tracker->lock);

        if (!count) {
                atomic64_inc(&tracker->snapshot_empty);
                return 0;
        }

        snap = kcalloc(count, sizeof(*snap), GFP_NOFS);
        if (!snap) {
                atomic64_inc(&tracker->snapshot_failures);
                return -ENOMEM;
        }

        spin_lock(&tracker->lock);
        node = rb_first_cached(&tracker->tree);
        while (node) {
                struct ceph_subvol_metric_rb_entry *entry =
                        rb_entry(node, struct ceph_subvol_metric_rb_entry, node);
                struct rb_node *next = rb_next(node);

                /* Skip entries with NO I/O activity at all */
                if (!entry->read_ops && !entry->write_ops) {
                        rb_erase_cached(&entry->node, &tracker->tree);
                        tracker->nr_entries--;
                        kmem_cache_free(ceph_subvol_metric_entry_cachep, entry);
                        node = next;
                        continue;
                }

                if (idx >= count) {
                        pr_warn("ceph: subvol metrics snapshot race (idx=%u count=%u)\n",
                                idx, count);
                        break;
                }

                snap[idx].subvolume_id = entry->subvolume_id;
                snap[idx].read_ops = entry->read_ops;
                snap[idx].write_ops = entry->write_ops;
                snap[idx].read_bytes = entry->read_bytes;
                snap[idx].write_bytes = entry->write_bytes;
                snap[idx].read_latency_us = entry->read_latency_us;
                snap[idx].write_latency_us = entry->write_latency_us;
                idx++;

                if (consume) {
                        entry->read_ops = 0;
                        entry->write_ops = 0;
                        entry->read_bytes = 0;
                        entry->write_bytes = 0;
                        entry->read_latency_us = 0;
                        entry->write_latency_us = 0;
                        rb_erase_cached(&entry->node, &tracker->tree);
                        tracker->nr_entries--;
                        kmem_cache_free(ceph_subvol_metric_entry_cachep, entry);
                }
                node = next;
        }
        spin_unlock(&tracker->lock);

        if (!idx) {
                kfree(snap);
                snap = NULL;
                ret = 0;
        } else {
                *nr = idx;
                *out = snap;
        }

        return ret;
}

void ceph_subvolume_metrics_free_snapshot(struct ceph_subvol_metric_snapshot *snapshot)
{
        kfree(snapshot);
}

/*
 * Dump subvolume metrics to a seq_file for debugfs.
 *
 * Iterates the rb-tree directly under spinlock to avoid allocation.
 * The lock hold time is minimal since we're only doing seq_printf calls.
 */
void ceph_subvolume_metrics_dump(struct ceph_subvolume_metrics_tracker *tracker,
                                 struct seq_file *s)
{
        struct rb_node *node;
        bool found = false;

        spin_lock(&tracker->lock);
        if (!tracker->enabled) {
                spin_unlock(&tracker->lock);
                seq_puts(s, "subvolume metrics disabled\n");
                return;
        }

        for (node = rb_first_cached(&tracker->tree); node; node = rb_next(node)) {
                struct ceph_subvol_metric_rb_entry *entry =
                        rb_entry(node, struct ceph_subvol_metric_rb_entry, node);
                u64 avg_rd_lat, avg_wr_lat;

                if (!entry->read_ops && !entry->write_ops)
                        continue;

                if (!found) {
                        seq_puts(s, "subvol_id       rd_ops    rd_bytes    rd_avg_lat_us  wr_ops    wr_bytes    wr_avg_lat_us\n");
                        seq_puts(s, "------------------------------------------------------------------------------------------------\n");
                        found = true;
                }

                avg_rd_lat = entry->read_ops ?
                        div64_u64(entry->read_latency_us, entry->read_ops) : 0;
                avg_wr_lat = entry->write_ops ?
                        div64_u64(entry->write_latency_us, entry->write_ops) : 0;

                seq_printf(s, "%-15llu%-10llu%-12llu%-16llu%-10llu%-12llu%-16llu\n",
                           entry->subvolume_id,
                           entry->read_ops,
                           entry->read_bytes,
                           avg_rd_lat,
                           entry->write_ops,
                           entry->write_bytes,
                           avg_wr_lat);
        }
        spin_unlock(&tracker->lock);

        if (!found)
                seq_puts(s, "(no subvolume metrics collected)\n");
}

void ceph_subvolume_metrics_record_io(struct ceph_mds_client *mdsc,
                                      struct ceph_inode_info *ci,
                                      bool is_write, size_t bytes,
                                      ktime_t start, ktime_t end)
{
        struct ceph_subvolume_metrics_tracker *tracker;
        u64 subvol_id;
        s64 delta_us;

        if (!mdsc || !ci || !bytes)
                return;

        tracker = &mdsc->subvol_metrics;
        atomic64_inc(&tracker->record_calls);

        if (!ceph_subvolume_metrics_enabled(tracker)) {
                atomic64_inc(&tracker->record_disabled);
                return;
        }

        subvol_id = READ_ONCE(ci->i_subvolume_id);
        if (subvol_id == CEPH_SUBVOLUME_ID_NONE) {
                atomic64_inc(&tracker->record_no_subvol);
                return;
        }

        delta_us = ktime_to_us(ktime_sub(end, start));
        if (delta_us <= 0)
                delta_us = 1;

        ceph_subvolume_metrics_record(tracker, subvol_id, is_write,
                                      bytes, (u64)delta_us);
}

int __init ceph_subvolume_metrics_cache_init(void)
{
        ceph_subvol_metric_entry_cachep = KMEM_CACHE(ceph_subvol_metric_rb_entry,
                                                    SLAB_RECLAIM_ACCOUNT);
        if (!ceph_subvol_metric_entry_cachep)
                return -ENOMEM;
        return 0;
}

void ceph_subvolume_metrics_cache_destroy(void)
{
        kmem_cache_destroy(ceph_subvol_metric_entry_cachep);
}