// SPDX-License-Identifier: GPL-2.0
/*
 * A scheduler that validates ops.dequeue() is called correctly:
 * - Tasks dispatched to terminal DSQs (local, global) bypass the BPF
 *   scheduler entirely: no ops.dequeue() should be called
 * - Tasks dispatched to user DSQs from ops.enqueue() enter BPF custody:
 *   ops.dequeue() must be called when they leave custody
 * - Every ops.enqueue() dispatch to non-terminal DSQs is followed by
 *   exactly one ops.dequeue() (validate 1:1 pairing and state machine)
 *
 * Copyright (c) 2026 NVIDIA Corporation.
 */

#include <scx/common.bpf.h>

#define SHARED_DSQ	0

/*
 * BPF internal queue.
 *
 * Tasks are stored here and consumed from ops.dispatch(), validating that
 * tasks on BPF internal structures still get ops.dequeue() when they
 * leave.
 */
struct {
	__uint(type, BPF_MAP_TYPE_QUEUE);
	__uint(max_entries, 32768);
	__type(value, s32);
} global_queue SEC(".maps");

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

UEI_DEFINE(uei);

/*
 * Counters to track the lifecycle of tasks:
 * - enqueue_cnt: Number of times ops.enqueue() was called
 * - dequeue_cnt: Number of times ops.dequeue() was called (any type)
 * - dispatch_dequeue_cnt: Number of regular dispatch dequeues (no flag)
 * - change_dequeue_cnt: Number of property change dequeues
 * - bpf_queue_full: Number of times the BPF internal queue was full
 */
u64 enqueue_cnt, dequeue_cnt, dispatch_dequeue_cnt, change_dequeue_cnt, bpf_queue_full;

/*
 * Test scenarios:
 * 0) Dispatch to local DSQ from ops.select_cpu() (terminal DSQ, bypasses BPF
 *    scheduler, no dequeue callbacks)
 * 1) Dispatch to global DSQ from ops.select_cpu() (terminal DSQ, bypasses BPF
 *    scheduler, no dequeue callbacks)
 * 2) Dispatch to shared user DSQ from ops.select_cpu() (enters BPF scheduler,
 *    dequeue callbacks expected)
 * 3) Dispatch to local DSQ from ops.enqueue() (terminal DSQ, bypasses BPF
 *    scheduler, no dequeue callbacks)
 * 4) Dispatch to global DSQ from ops.enqueue() (terminal DSQ, bypasses BPF
 *    scheduler, no dequeue callbacks)
 * 5) Dispatch to shared user DSQ from ops.enqueue() (enters BPF scheduler,
 *    dequeue callbacks expected)
 * 6) BPF internal queue from ops.enqueue(): store task PIDs in ops.enqueue(),
 *    consume in ops.dispatch() and dispatch to local DSQ (validates dequeue
 *    for tasks stored in internal BPF data structures)
 */
u32 test_scenario;

/*
 * Per-task state to track lifecycle and validate workflow semantics.
 * State transitions:
 *   NONE -> ENQUEUED (on enqueue)
 *   NONE -> DISPATCHED (on direct dispatch to terminal DSQ)
 *   ENQUEUED -> DISPATCHED (on dispatch dequeue)
 *   DISPATCHED -> NONE (on property change dequeue or re-enqueue)
 *   ENQUEUED -> NONE (on property change dequeue before dispatch)
 */
enum task_state {
	TASK_NONE = 0,
	TASK_ENQUEUED,
	TASK_DISPATCHED,
};

struct task_ctx {
	enum task_state state; /* Current state in the workflow */
	u64 enqueue_seq;       /* Sequence number for debugging */
};

struct {
	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
	__uint(map_flags, BPF_F_NO_PREALLOC);
	__type(key, int);
	__type(value, struct task_ctx);
} task_ctx_stor SEC(".maps");

static struct task_ctx *try_lookup_task_ctx(struct task_struct *p)
{
	return bpf_task_storage_get(&task_ctx_stor, p, 0, 0);
}

s32 BPF_STRUCT_OPS(dequeue_select_cpu, struct task_struct *p,
		   s32 prev_cpu, u64 wake_flags)
{
	struct task_ctx *tctx;

	tctx = try_lookup_task_ctx(p);
	if (!tctx)
		return prev_cpu;

	switch (test_scenario) {
	case 0:
		/*
		 * Direct dispatch to the local DSQ.
		 *
		 * Task bypasses BPF scheduler entirely: no enqueue
		 * tracking, no ops.dequeue() callbacks.
		 */
		scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, 0);
		tctx->state = TASK_DISPATCHED;
		break;
	case 1:
		/*
		 * Direct dispatch to the global DSQ.
		 *
		 * Task bypasses BPF scheduler entirely: no enqueue
		 * tracking, no ops.dequeue() callbacks.
		 */
		scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0);
		tctx->state = TASK_DISPATCHED;
		break;
	case 2:
		/*
		 * Dispatch to a shared user DSQ.
		 *
		 * Task enters BPF scheduler management: track
		 * enqueue/dequeue lifecycle and validate state
		 * transitions.
		 */
		if (tctx->state == TASK_ENQUEUED)
			scx_bpf_error("%d (%s): enqueue while in ENQUEUED state seq=%llu",
				      p->pid, p->comm, tctx->enqueue_seq);

		scx_bpf_dsq_insert(p, SHARED_DSQ, SCX_SLICE_DFL, 0);

		__sync_fetch_and_add(&enqueue_cnt, 1);

		tctx->state = TASK_ENQUEUED;
		tctx->enqueue_seq++;
		break;
	}

	return prev_cpu;
}

void BPF_STRUCT_OPS(dequeue_enqueue, struct task_struct *p, u64 enq_flags)
{
	struct task_ctx *tctx;
	s32 pid = p->pid;

	tctx = try_lookup_task_ctx(p);
	if (!tctx)
		return;

	switch (test_scenario) {
	case 3:
		/*
		 * Direct dispatch to the local DSQ.
		 *
		 * Task bypasses BPF scheduler entirely: no enqueue
		 * tracking, no ops.dequeue() callbacks.
		 */
		scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, enq_flags);
		tctx->state = TASK_DISPATCHED;
		break;
	case 4:
		/*
		 * Direct dispatch to the global DSQ.
		 *
		 * Task bypasses BPF scheduler entirely: no enqueue
		 * tracking, no ops.dequeue() callbacks.
		 */
		scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags);
		tctx->state = TASK_DISPATCHED;
		break;
	case 5:
		/*
		 * Dispatch to shared user DSQ.
		 *
		 * Task enters BPF scheduler management: track
		 * enqueue/dequeue lifecycle and validate state
		 * transitions.
		 */
		if (tctx->state == TASK_ENQUEUED)
			scx_bpf_error("%d (%s): enqueue while in ENQUEUED state seq=%llu",
				      p->pid, p->comm, tctx->enqueue_seq);

		scx_bpf_dsq_insert(p, SHARED_DSQ, SCX_SLICE_DFL, enq_flags);

		__sync_fetch_and_add(&enqueue_cnt, 1);

		tctx->state = TASK_ENQUEUED;
		tctx->enqueue_seq++;
		break;
	case 6:
		/*
		 * Store task in BPF internal queue.
		 *
		 * Task enters BPF scheduler management: track
		 * enqueue/dequeue lifecycle and validate state
		 * transitions.
		 */
		if (tctx->state == TASK_ENQUEUED)
			scx_bpf_error("%d (%s): enqueue while in ENQUEUED state seq=%llu",
				      p->pid, p->comm, tctx->enqueue_seq);

		if (bpf_map_push_elem(&global_queue, &pid, 0)) {
			scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags);
			__sync_fetch_and_add(&bpf_queue_full, 1);

			tctx->state = TASK_DISPATCHED;
		} else {
			__sync_fetch_and_add(&enqueue_cnt, 1);

			tctx->state = TASK_ENQUEUED;
			tctx->enqueue_seq++;
		}
		break;
	default:
		/* For all other scenarios, dispatch to the global DSQ */
		scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags);
		tctx->state = TASK_DISPATCHED;
		break;
	}

	scx_bpf_kick_cpu(scx_bpf_task_cpu(p), SCX_KICK_IDLE);
}

void BPF_STRUCT_OPS(dequeue_dequeue, struct task_struct *p, u64 deq_flags)
{
	struct task_ctx *tctx;

	__sync_fetch_and_add(&dequeue_cnt, 1);

	tctx = try_lookup_task_ctx(p);
	if (!tctx)
		return;

	/*
	 * For scenarios 0, 1, 3, and 4 (terminal DSQs: local and global),
	 * ops.dequeue() should never be called because tasks bypass the
	 * BPF scheduler entirely. If we get here, it's a kernel bug.
	 */
	if (test_scenario == 0 || test_scenario == 3) {
		scx_bpf_error("%d (%s): dequeue called for local DSQ scenario",
			      p->pid, p->comm);
		return;
	}

	if (test_scenario == 1 || test_scenario == 4) {
		scx_bpf_error("%d (%s): dequeue called for global DSQ scenario",
			      p->pid, p->comm);
		return;
	}

	if (deq_flags & SCX_DEQ_SCHED_CHANGE) {
		/*
		 * Property change interrupting the workflow. Valid from
		 * both ENQUEUED and DISPATCHED states. Transitions task
		 * back to NONE state.
		 */
		__sync_fetch_and_add(&change_dequeue_cnt, 1);

		/* Validate state transition */
		if (tctx->state != TASK_ENQUEUED && tctx->state != TASK_DISPATCHED)
			scx_bpf_error("%d (%s): invalid property change dequeue state=%d seq=%llu",
				      p->pid, p->comm, tctx->state, tctx->enqueue_seq);

		/*
		 * Transition back to NONE: task outside scheduler control.
		 *
		 * Scenario 6: dispatch() checks tctx->state after popping a
		 * PID, if the task is in state NONE, it was dequeued by
		 * property change and must not be dispatched (this
		 * prevents "target CPU not allowed").
		 */
		tctx->state = TASK_NONE;
	} else {
		/*
		 * Regular dispatch dequeue: kernel is moving the task from
		 * BPF custody to a terminal DSQ. Normally we come from
		 * ENQUEUED state. We can also see TASK_NONE if the task
		 * was dequeued by property change (SCX_DEQ_SCHED_CHANGE)
		 * while it was already on a DSQ (dispatched but not yet
		 * consumed); in that case we just leave state as NONE.
		 */
		__sync_fetch_and_add(&dispatch_dequeue_cnt, 1);

		/*
		 * Must be ENQUEUED (normal path) or NONE (already dequeued
		 * by property change while on a DSQ).
		 */
		if (tctx->state != TASK_ENQUEUED && tctx->state != TASK_NONE)
			scx_bpf_error("%d (%s): dispatch dequeue from state %d seq=%llu",
				      p->pid, p->comm, tctx->state, tctx->enqueue_seq);

		if (tctx->state == TASK_ENQUEUED)
			tctx->state = TASK_DISPATCHED;

		/* NONE: leave as-is, task was already property-change dequeued */
	}
}

void BPF_STRUCT_OPS(dequeue_dispatch, s32 cpu, struct task_struct *prev)
{
	if (test_scenario == 6) {
		struct task_ctx *tctx;
		struct task_struct *p;
		s32 pid;

		if (bpf_map_pop_elem(&global_queue, &pid))
			return;

		p = bpf_task_from_pid(pid);
		if (!p)
			return;

		/*
		 * If the task was dequeued by property change
		 * (ops.dequeue() set tctx->state = TASK_NONE), skip
		 * dispatch.
		 */
		tctx = try_lookup_task_ctx(p);
		if (!tctx || tctx->state == TASK_NONE) {
			bpf_task_release(p);
			return;
		}

		/*
		 * Dispatch to this CPU's local DSQ if allowed, otherwise
		 * fallback to the global DSQ.
		 */
		if (bpf_cpumask_test_cpu(cpu, p->cpus_ptr))
			scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | cpu, SCX_SLICE_DFL, 0);
		else
			scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0);

		bpf_task_release(p);
	} else {
		scx_bpf_dsq_move_to_local(SHARED_DSQ, 0);
	}
}

s32 BPF_STRUCT_OPS(dequeue_init_task, struct task_struct *p,
		   struct scx_init_task_args *args)
{
	struct task_ctx *tctx;

	tctx = bpf_task_storage_get(&task_ctx_stor, p, 0,
				   BPF_LOCAL_STORAGE_GET_F_CREATE);
	if (!tctx)
		return -ENOMEM;

	return 0;
}

s32 BPF_STRUCT_OPS_SLEEPABLE(dequeue_init)
{
	s32 ret;

	ret = scx_bpf_create_dsq(SHARED_DSQ, -1);
	if (ret)
		return ret;

	return 0;
}

void BPF_STRUCT_OPS(dequeue_exit, struct scx_exit_info *ei)
{
	UEI_RECORD(uei, ei);
}

SEC(".struct_ops.link")
struct sched_ext_ops dequeue_ops = {
	.select_cpu		= (void *)dequeue_select_cpu,
	.enqueue		= (void *)dequeue_enqueue,
	.dequeue		= (void *)dequeue_dequeue,
	.dispatch		= (void *)dequeue_dispatch,
	.init_task		= (void *)dequeue_init_task,
	.init			= (void *)dequeue_init,
	.exit			= (void *)dequeue_exit,
	.flags			= SCX_OPS_ENQ_LAST,
	.name			= "dequeue_test",
};