#include "vmlinux.h"
#include "bpf_common.h"
#include <bpf/bpf_tracing.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_core_read.h>
#include "bpf_ratelimit.h"

char __license[] SEC("license") = "Dual MIT/GPL";

#define CPU_NUM 128
BPF_RATELIMIT_IN_MAP(rate, 1, CPU_NUM * 10000, 0);

struct {
    __uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
    __uint(key_size, sizeof(int));
    __uint(value_size, sizeof(u32));
} oom_perf_events SEC(".maps");

struct oom_info {
    char trigger_comm[COMPAT_TASK_COMM_LEN];
    char victim_comm[COMPAT_TASK_COMM_LEN];
    u32 trigger_pid;
    u32 victim_pid;
    u64 trigger_memcg_css;
    u64 victim_memcg_css;
};

SEC("kprobe/oom_kill_process")
int kprobe_oom_kill_process(struct pt_regs *ctx)
{
    struct oom_control *oc;
    struct oom_info info = {};
    struct task_struct *trigger_task, *victim_task;

    if (bpf_ratelimited_in_map(ctx, rate))
        return 0;

    oc = (void *)ctx->di;

    if (!oc)
        return 0;
    trigger_task = (struct task_struct *)bpf_get_current_task();
    victim_task = BPF_CORE_READ(oc, chosen);
    info.trigger_pid = BPF_CORE_READ(trigger_task, pid);
    info.victim_pid = BPF_CORE_READ(victim_task, pid);
    BPF_CORE_READ_STR_INTO(&info.trigger_comm, trigger_task, comm);
    BPF_CORE_READ_STR_INTO(&info.victim_comm, victim_task, comm);
    
    info.victim_memcg_css = (u64)BPF_CORE_READ(victim_task, cgroups, subsys[4]);
    info.trigger_memcg_css = (u64)BPF_CORE_READ(trigger_task, cgroups, subsys[4]);
    
    bpf_perf_event_output(ctx, &oom_perf_events, COMPAT_BPF_F_CURRENT_CPU, &info, sizeof(info));
    return 0;
}
