// SPDX-License-Identifier: GPL-2.0
#include <linux/ftrace.h>
#include <linux/tracepoint.h>
#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/init.h>
#include <linux/rv.h>
#include <rv/instrumentation.h>

#define MODULE_NAME "stall"

#include <trace/events/sched.h>
#include <rv_trace.h>

#define RV_MON_TYPE RV_MON_PER_TASK
#define HA_TIMER_TYPE HA_TIMER_WHEEL
#include "stall.h"
#include <rv/ha_monitor.h>

static u64 threshold_jiffies = 1000;
module_param(threshold_jiffies, ullong, 0644);

static u64 ha_get_env(struct ha_monitor *ha_mon, enum envs_stall env, u64 time_ns)
{
        if (env == clk_stall)
                return ha_get_clk_jiffy(ha_mon, env);
        return ENV_INVALID_VALUE;
}

static void ha_reset_env(struct ha_monitor *ha_mon, enum envs_stall env, u64 time_ns)
{
        if (env == clk_stall)
                ha_reset_clk_jiffy(ha_mon, env);
}

static inline bool ha_verify_invariants(struct ha_monitor *ha_mon,
                                        enum states curr_state, enum events event,
                                        enum states next_state, u64 time_ns)
{
        if (curr_state == enqueued_stall)
                return ha_check_invariant_jiffy(ha_mon, clk_stall, time_ns);
        return true;
}

static inline bool ha_verify_guards(struct ha_monitor *ha_mon,
                                    enum states curr_state, enum events event,
                                    enum states next_state, u64 time_ns)
{
        bool res = true;

        if (curr_state == dequeued_stall && event == sched_wakeup_stall)
                ha_reset_env(ha_mon, clk_stall, time_ns);
        else if (curr_state == running_stall && event == sched_switch_preempt_stall)
                ha_reset_env(ha_mon, clk_stall, time_ns);
        return res;
}

static inline void ha_setup_invariants(struct ha_monitor *ha_mon,
                                       enum states curr_state, enum events event,
                                       enum states next_state, u64 time_ns)
{
        if (next_state == curr_state)
                return;
        if (next_state == enqueued_stall)
                ha_start_timer_jiffy(ha_mon, clk_stall, threshold_jiffies, time_ns);
        else if (curr_state == enqueued_stall)
                ha_cancel_timer(ha_mon);
}

static bool ha_verify_constraint(struct ha_monitor *ha_mon,
                                 enum states curr_state, enum events event,
                                 enum states next_state, u64 time_ns)
{
        if (!ha_verify_invariants(ha_mon, curr_state, event, next_state, time_ns))
                return false;

        if (!ha_verify_guards(ha_mon, curr_state, event, next_state, time_ns))
                return false;

        ha_setup_invariants(ha_mon, curr_state, event, next_state, time_ns);

        return true;
}

static void handle_sched_switch(void *data, bool preempt,
                                struct task_struct *prev,
                                struct task_struct *next,
                                unsigned int prev_state)
{
        if (!preempt && prev_state != TASK_RUNNING)
                da_handle_start_event(prev, sched_switch_wait_stall);
        else
                da_handle_event(prev, sched_switch_preempt_stall);
        da_handle_event(next, sched_switch_in_stall);
}

static void handle_sched_wakeup(void *data, struct task_struct *p)
{
        da_handle_event(p, sched_wakeup_stall);
}

static int enable_stall(void)
{
        int retval;

        retval = da_monitor_init();
        if (retval)
                return retval;

        rv_attach_trace_probe("stall", sched_switch, handle_sched_switch);
        rv_attach_trace_probe("stall", sched_wakeup, handle_sched_wakeup);

        return 0;
}

static void disable_stall(void)
{
        rv_this.enabled = 0;

        rv_detach_trace_probe("stall", sched_switch, handle_sched_switch);
        rv_detach_trace_probe("stall", sched_wakeup, handle_sched_wakeup);

        da_monitor_destroy();
}

static struct rv_monitor rv_this = {
        .name = "stall",
        .description = "identify tasks stalled for longer than a threshold.",
        .enable = enable_stall,
        .disable = disable_stall,
        .reset = da_monitor_reset_all,
        .enabled = 0,
};

static int __init register_stall(void)
{
        return rv_register_monitor(&rv_this, NULL);
}

static void __exit unregister_stall(void)
{
        rv_unregister_monitor(&rv_this);
}

module_init(register_stall);
module_exit(unregister_stall);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Gabriele Monaco <gmonaco@redhat.com>");
MODULE_DESCRIPTION("stall: identify tasks stalled for longer than a threshold.");