// 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;
}