Linux Scheduler Class

my_scheduler.c

// SPDX-License-Identifier: GPL-2.0-only
/*
 * Generic entry points and implementation of Linux a scheduling class.
 */
#include "sched.h"
#include <linux/sched/myscheduler.h>
#include <linux/kernel.h>
#include <linux/jiffies.h>
#include <linux/sched/task.h>


void init_myscheduler_rq(struct myscheduler_rq *myscheduler_rq)
{
	INIT_LIST_HEAD(&myscheduler_rq->run_list);
	myscheduler_rq->myscheduler_nr_running = 0;
}

static inline int on_myscheduler_rq(struct sched_myscheduler_entity *myscheduler_se)
{
	return myscheduler_se->on_rq;
}

static inline struct task_struct *myscheduler_task_of(struct sched_myscheduler_entity *myscheduler_se)
{
	return container_of(myscheduler_se, struct task_struct, myscheduler_se);
}

void insert_task_by_weight(struct rq *rq, struct task_struct *p)
{
	struct list_head *pos;
	struct sched_myscheduler_entity *iter;

	list_for_each(pos, &rq->myscheduler.run_list) {
		iter = list_entry(pos, struct sched_myscheduler_entity, run_list);
		if (iter->weight > p->myscheduler_se.weight) {
			list_add_tail(&p->myscheduler_se.run_list, pos);
			return;
		}
	}
	list_add_tail(&p->myscheduler_se.run_list, &rq->myscheduler.run_list);
}

static void enqueue_task_myscheduler(struct rq *rq, struct task_struct *p, int flags)
{
	struct sched_myscheduler_entity *myscheduler_se = &p->myscheduler_se;

	if (p->policy != SCHED_MINE)
		return;

	if (!on_my_scheduler(my_scheduler)) {
		if (my_scheduler>weight < MAX_WEIGHT)
			insert_task_by_weight(rq, p);
		else
			list_add_tail(&my_scheduler>run_list, &rq->my_scheduler.run_list);
		my_scheduler>on_rq = 1;
		rq->my_scheduler.my_scheduler++;
	}
}

static void dequeue_task_my_scheduler(struct rq *rq, struct task_struct *p, int flags)
{
	struct sched_my_scheduler *my_scheduler = &p->my_scheduler;

	if (on_my_scheduler(my_scheduler)) {
		list_del_init(&my_scheduler>run_list);
		my_scheduler>on_rq = 0;
		rq->my_scheduler.my_scheduler;
	}
}

static void requeue_task_my_scheduler(struct rq *rq, struct task_struct *p)
{
	list_move_tail(&p->my_scheduler.run_list, &rq->my_scheduler.run_list);
}

static void yield_task_my_scheduler(struct rq *rq)
{
	WARN_ON_ONCE(1);
}

static struct task_struct *pick_next_task_my_scheduler(struct rq *rq)
{
	struct sched_my_scheduler *my_scheduler;
	struct task_struct *next;
	struct list_head *pos;

	// Initialize lowest weight to a high value
	int lowest_weight = 10000;
	struct sched_my_scheduler *lowest_weight_se = NULL;

	if (!rq->my_scheduler.my_scheduler)
		return NULL;

	// Iterate through the run list to find the task with the lowest weight
	list_for_each(pos, &rq->my_scheduler.run_list) {
		my_scheduler = list_entry(pos, struct sched_my_scheduler, run_list);
		if (my_scheduler>weight < lowest_weight) {
			lowest_weight = my_scheduler>weight;
			lowest_weight_se = my_scheduler;
		}
	}

	// If a task with the lowest weight is found, set it as next
	if (lowest_weight_se) {
		next = my_scheduler(lowest_weight_se);
	} else {
		// Otherwise, pick the first task in the list
		my_scheduler = list_first_entry(&rq->my_scheduler.run_list, struct sched_my_scheduler, run_list);
		next = my_scheduler(my_scheduler);
	}

	if (!next)
		return NULL;

	return next;
}

/* PREEMPTION */

/*
 *
 * Helper function to determine if the current task should be preempted
 */
static bool should_preempt_curr(struct rq *rq, struct task_struct *p)
{
	struct task_struct *next_task;

	/*
	 * Get the next task in the run queue (if any)
	 * Return if no other task to run
	 */
	next_task = pick_next_task_my_scheduler(rq);
	if (!next_task)
		return false;
	/*
	 * Compare priorities (or other criteria) to decide if preemption is needed
	 * For example, preempt if the next task has a higher priority
	 */
	return next_task->prio < p->prio;
}

static void check_preempt_curr_my_scheduler(struct rq *rq, struct task_struct *p, int flags)
{
	/*
	 * Check if there is a more urgent task to run
	 * and mark the current task for rescheduling
	 */
	if (should_preempt_curr(rq, p))
		resched_curr(rq);
}

static inline void set_next_task_my_scheduler(struct rq *rq, struct task_struct *p, bool first)
{
	WARN_ON_ONCE(1);
}

static void update_curr_my_scheduler(struct rq *rq)
{
	struct task_struct *curr = rq->curr;
	u64 now, delta_exec;

	if (curr->sched_class != &my_scheduler)
		return;
	curr->sched_info.pcount++;
	now = jiffies;
	delta_exec = now - curr->sched_info.last_arrival;
	curr->sched_info.last_arrival = now;
}

static bool task_has_high_priority(struct task_struct *p)
{
	/* use my_scheduler weight field */
	struct sched_my_scheduler *my_scheduler = &p->my_scheduler;

	return my_scheduler>weight < MAX_WEIGHT;
}

/*
 * Preempt the current task with a newly woken task if needed:
 * We have to consider the priority.
 */
static void prio_changed_my_scheduler(struct rq *rq, struct task_struct *p, int oldprio)
{
	struct sched_my_scheduler *my_scheduler = &p->my_scheduler;

	//if the priority is higher than the current task, we reschedule
	if (p->policy != SCHED_MINE)
		return;
	if  (my_scheduler>weight < rq->curr->my_scheduler.weight)
		resched_curr(rq);
}

/*
 * Helper function to determine if a task should be running
 */
static bool should_be_running(struct task_struct *p)
{
	/*
	 * Implement logic to determine if the task should be running
	 * This could be based on various factors like task state, priority, etc.
	 * For example:
	 */
	return p->__state == TASK_RUNNING && task_has_high_priority(p);
}


static void switched_to_my_scheduler(struct rq *rq, struct task_struct *p)
{
	struct sched_my_scheduler *my_scheduler = &p->my_scheduler;

	// Check if the task is already on the my_scheduler run queue
	if (on_my_scheduler(my_scheduler)) {
		// If the task is already running but should not be, dequeue it
		if (!should_be_running(p))
			dequeue_task_my_scheduler(rq, p, 0);
	} else {
		// If the task is not running but should be, enqueue it
		if (should_be_running(p))
			enqueue_task_my_scheduler(rq, p, 0);
	}

	// Handle priority change for tasks already in the my_scheduler scheduling class
	if (p->policy == SCHED_MINE)
		check_preempt_curr_my_scheduler(rq, p, 0);
}

static void put_prev_task_my_scheduler(struct rq *rq, struct task_struct *prev)
{
	if (on_my_scheduler_rq(&prev->my_scheduler_se))
		update_curr_my_scheduler(rq);
}

#ifdef CONFIG_SMP
static int select_task_rq_my_scheduler(struct task_struct *p, int cpu, int sd_flag)
{
	return task_cpu(p);
}

static int balance_my_scheduler(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
{
	return WARN_ON_ONCE(1);
}

static struct rq *find_lock_lowest_my_scheduler(struct task_struct *task, struct rq *rq)
{
	return NULL;
}
#endif /* CONFIG_SMP */

static void task_tick_my_scheduler(struct rq *rq, struct task_struct *p, int queued)
{
	struct sched_my_scheduler_entity *my_scheduler_se = &p->my_scheduler_se;

	if (p->policy != SCHED_MINE)
		return;

	/* my_scheduler_se->weight is not 10000 we don't swap */
	if (my_scheduler_se->weight != MAX_WEIGHT)
		return;

	//every tick we make a swap
	my_scheduler_se->time_slice = sched_rr_timeslice;

	if (my_scheduler_se->run_list.prev == my_scheduler_se->run_list.next)
		return;

	requeue_task_my_scheduler(rq, p);
	resched_curr(rq);
}

static unsigned int get_rr_interval_myscheduler(struct rq *rq, struct task_struct *task)
{
	return my_scheduler;
}


/*
 * Simple, special scheduling class for the per-CPU tasks:
 */
DEFINE_SCHED_CLASS(my_scheduler) = {

	.enqueue_task		= enqueue_task_my_scheduler,
	.dequeue_task		= dequeue_task_my_scheduler,
	.yield_task		= yield_task_my_scheduler,

	.check_preempt_curr	= check_preempt_curr_my_scheduler,

	.pick_next_task		= pick_next_task_my_scheduler,
	.put_prev_task		= put_prev_task_my_scheduler,
	.set_next_task          = set_next_task_my_scheduler,

#ifdef CONFIG_SMP
	.balance		= balance_my_scheduler,
	.select_task_rq		= select_task_rq_my_scheduler,
	.set_cpus_allowed       = set_cpus_allowed_common,
	.rq_online              = rq_online_my_scheduler,
	.rq_offline             = rq_offline_my_scheduler,
	.task_woken		= task_woken_my_scheduler,
	.switched_from		= switched_from_my_scheduler,
	.find_lock_rq		= find_lock_lowest_my_scheduler,
#endif

	.task_tick		= task_tick_my_scheduler,

	.get_rr_interval	= get_rr_interval_my_scheduler,

	.prio_changed		= prio_changed_my_scheduler,
	.switched_to		= switched_to_my_scheduler,

	.update_curr		= update_curr_my_scheduler,
};