#ifndef __BPF_RATELIMIT_H__
#define __BPF_RATELIMIT_H__

#include <bpf/bpf_helpers.h>

struct bpf_ratelimit {
	uint64_t interval; // unit: second
	uint64_t begin;
	uint64_t burst;	    // max events/interval
	uint64_t max_burst; // max burst
	uint64_t events;    // current events/interval
	uint64_t nmissed;   // missed events/interval

	uint64_t total_events;	 // total events
	uint64_t total_nmissed;	 // total missed events
	uint64_t total_interval; // total interval
};

#define BPF_RATELIMIT(name, interval, burst)                                   \
	struct bpf_ratelimit name = {interval, 0, burst, 0, 0, 0, 0, 0, 0}

// bpf_ratelimited: whether the threshold is exceeded
//
// @rate: struct bpf_ratelimit *
// @return:
//   true: the threshold is exceeded
//   false: the threshold is not exceeded
static __always_inline bool bpf_ratelimited(struct bpf_ratelimit *rate)
{
	// validate
	if (rate == NULL || rate->interval == 0)
		return false;

	u64 curr = bpf_ktime_get_ns() / 1000000000;

	if (rate->begin == 0)
		rate->begin = curr;

	if (curr > rate->begin + rate->interval) {
		__sync_fetch_and_add(&rate->total_interval, curr - rate->begin);
		rate->begin  = curr;
		rate->events = rate->nmissed = 0;
	}

	if (rate->burst && rate->burst > rate->events) {
		__sync_fetch_and_add(&rate->events, 1);
		__sync_fetch_and_add(&rate->total_events, 1);
		return false;
	}

	__sync_fetch_and_add(&rate->nmissed, 1);
	__sync_fetch_and_add(&rate->total_nmissed, 1);
	return true;
}

#define BPF_RATELIMIT_IN_MAP(name, interval, burst, max_burst)                 \
	struct {                                                               \
		__uint(type, BPF_MAP_TYPE_ARRAY);                              \
		__uint(key_size, sizeof(u32));                                 \
		__uint(value_size, sizeof(struct bpf_ratelimit));              \
		__uint(max_entries, 1);                                        \
	} bpf_rlimit_##name SEC(".maps");                                      \
	struct {                                                               \
		__uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);                   \
		__uint(key_size, sizeof(int));                                 \
		__uint(value_size, sizeof(u32));                               \
	} event_bpf_rlimit_##name SEC(".maps");                                \
	volatile const struct bpf_ratelimit ___bpf_rlimit_cfg_##name = {       \
		interval, 0, burst, max_burst, 0, 0, 0, 0, 0}

// bpf_ratelimited_in_map: whether the threshold is exceeded
//
// @rate: struct bpf_ratelimit *
// @return:
//   true: the threshold is exceeded
//   false: the threshold is not exceeded
#define bpf_ratelimited_in_map(ctx, rate)                                      \
	bpf_ratelimited_core_in_map(ctx, &bpf_rlimit_##rate,                   \
				    &event_bpf_rlimit_##rate,                  \
				    &___bpf_rlimit_cfg_##rate)

static __always_inline bool
bpf_ratelimited_core_in_map(void *ctx, void *map, void *perf_map,
			    const volatile struct bpf_ratelimit *cfg)
{
	u32 key			   = 0;
	struct bpf_ratelimit *rate = NULL;

	rate = bpf_map_lookup_elem(map, &key);
	if (rate == NULL)
		return false;

	// init from cfg
	if (rate->interval == 0) {
		rate->interval	= cfg->interval;
		rate->burst	= cfg->burst;
		rate->max_burst = cfg->max_burst;
	}

	// the threshold is not exceeded, return false
	u64 old_nmissed = rate->nmissed;
	if (!bpf_ratelimited(rate))
		return false;

	// the threshold/max_burst is exceeded, notify once in a cycle
	if (old_nmissed == 0 || (rate->max_burst > 0 &&
				 rate->nmissed > rate->max_burst - rate->burst))
		bpf_perf_event_output(ctx, perf_map, BPF_F_CURRENT_CPU, rate,
				      sizeof(struct bpf_ratelimit));
	return true;
}

#endif
