// SPDX-License-Identifier: GPL-2.0
#include <linux/bpf.h>
#include <bpf/bpf_tracing.h>
#include <stdbool.h>
#include "timerlat_bpf.h"

#define nosubprog __always_inline
#define MAX_ENTRIES_DEFAULT 4096

char LICENSE[] SEC("license") = "GPL";

struct trace_event_raw_timerlat_sample {
        unsigned long long timer_latency;
        int context;
} __attribute__((preserve_access_index));

struct {
        __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
        __uint(max_entries, MAX_ENTRIES_DEFAULT);
        __type(key, unsigned int);
        __type(value, unsigned long long);
} hist_irq SEC(".maps"), hist_thread SEC(".maps"), hist_user SEC(".maps");

struct {
        __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
        __uint(max_entries, SUMMARY_FIELD_N);
        __type(key, unsigned int);
        __type(value, unsigned long long);
} summary_irq SEC(".maps"), summary_thread SEC(".maps"), summary_user SEC(".maps");

struct {
        __uint(type, BPF_MAP_TYPE_ARRAY);
        __uint(max_entries, 1);
        __type(key, unsigned int);
        __type(value, unsigned long long);
} stop_tracing SEC(".maps");

struct {
        __uint(type, BPF_MAP_TYPE_RINGBUF);
        __uint(max_entries, 1);
} signal_stop_tracing SEC(".maps");

struct {
        __uint(type, BPF_MAP_TYPE_PROG_ARRAY);
        __uint(key_size, sizeof(unsigned int));
        __uint(max_entries, 1);
        __array(values, unsigned int (void *));
} bpf_action SEC(".maps") = {
        .values = {
                [0] = 0
        },
};

/* Params to be set by rtla */
const volatile int bucket_size = 1;
const volatile int output_divisor = 1000;
const volatile int entries = 256;
const volatile int irq_threshold;
const volatile int thread_threshold;
const volatile bool aa_only;

nosubprog unsigned long long map_get(void *map,
                                     unsigned int key)
{
        unsigned long long *value_ptr;

        value_ptr = bpf_map_lookup_elem(map, &key);

        return !value_ptr ? 0 : *value_ptr;
}

nosubprog void map_set(void *map,
                       unsigned int key,
                       unsigned long long value)
{
        bpf_map_update_elem(map, &key, &value, BPF_ANY);
}

nosubprog void map_increment(void *map,
                             unsigned int key)
{
        map_set(map, key, map_get(map, key) + 1);
}

nosubprog void update_main_hist(void *map,
                                int bucket)
{
        if (entries == 0)
                /* No histogram */
                return;

        if (bucket >= entries)
                /* Overflow */
                return;

        map_increment(map, bucket);
}

nosubprog void update_summary(void *map,
                              unsigned long long latency,
                              int bucket)
{
        if (aa_only)
                /* Auto-analysis only, nothing to be done here */
                return;

        map_set(map, SUMMARY_CURRENT, latency);

        if (bucket >= entries)
                /* Overflow */
                map_increment(map, SUMMARY_OVERFLOW);

        if (latency > map_get(map, SUMMARY_MAX))
                map_set(map, SUMMARY_MAX, latency);

        if (latency < map_get(map, SUMMARY_MIN) || map_get(map, SUMMARY_COUNT) == 0)
                map_set(map, SUMMARY_MIN, latency);

        map_increment(map, SUMMARY_COUNT);
        map_set(map, SUMMARY_SUM, map_get(map, SUMMARY_SUM) + latency);
}

nosubprog void set_stop_tracing(struct trace_event_raw_timerlat_sample *tp_args)
{
        int value = 0;

        /* Suppress further sample processing */
        map_set(&stop_tracing, 0, 1);

        /* Signal to userspace */
        bpf_ringbuf_output(&signal_stop_tracing, &value, sizeof(value), 0);

        /*
         * Call into BPF action program, if attached.
         * Otherwise, just silently fail.
         */
        bpf_tail_call(tp_args, &bpf_action, 0);
}

SEC("tp/osnoise/timerlat_sample")
int handle_timerlat_sample(struct trace_event_raw_timerlat_sample *tp_args)
{
        unsigned long long latency, latency_us;
        int bucket;

        if (map_get(&stop_tracing, 0))
                return 0;

        latency = tp_args->timer_latency / output_divisor;
        latency_us = tp_args->timer_latency / 1000;
        bucket = latency / bucket_size;

        if (tp_args->context == 0) {
                update_main_hist(&hist_irq, bucket);
                update_summary(&summary_irq, latency, bucket);

                if (irq_threshold != 0 && latency_us >= irq_threshold)
                        set_stop_tracing(tp_args);
        } else if (tp_args->context == 1) {
                update_main_hist(&hist_thread, bucket);
                update_summary(&summary_thread, latency, bucket);

                if (thread_threshold != 0 && latency_us >= thread_threshold)
                        set_stop_tracing(tp_args);
        } else {
                update_main_hist(&hist_user, bucket);
                update_summary(&summary_user, latency, bucket);

                if (thread_threshold != 0 && latency_us >= thread_threshold)
                        set_stop_tracing(tp_args);
        }

        return 0;
}