scheduler: load balancing initial implementation

kernel-src/include/kernel/scheduler.h

@@ -57,6 +57,7 @@ void sched_thread_wakeup_callback(thread_t *thread);
 thread_t *sched_select_next_thread(void);
 void sched_insert_in_cpu_queue(struct cpu_t *cpu, thread_t *thread);
 void sched_preempt_cpu(struct cpu_t *cpu);
+thread_t *sched_steal_work_from_cpu(struct cpu_t *cpu);
 
 static inline int sched_thread_run_queue_index(int interactivity) {
 	return min(interactivity * SCHED_RUN_QUEUE_SIZE / SCHED_MAX_INTERACTIVITY, SCHED_RUN_QUEUE_SIZE - 1);

kernel-src/include/x86-64/arch/cpu.h

@@ -56,6 +56,7 @@ typedef struct cpu_t {
 	int last_interactivity;
 	int last_queue;
 	size_t thread_count;
+	size_t stealable_thread_count;
 	spinlock_t sched_lock;
 
 	timerentry_t calendar_tick_timer_entry;

kernel-src/sys/scheduler/queue.c

@@ -72,7 +72,6 @@ static thread_t *pop_from_run_queue(sched_run_queue_t *run_queue) {
 	if (idx == -1)
 		return NULL;
 
-
 	thread_t *thread = run_queue->queues[idx].run;
 	
 	run_queue->queues[idx].run = thread->next;
@@ -87,6 +86,72 @@ static thread_t *pop_from_run_queue(sched_run_queue_t *run_queue) {
 	return thread;
 }
 
+static thread_t *steal_from_calendar_queue(sched_calendar_queue_t *calendar_queue) {
+	if (calendar_queue->thread_count == 0)
+		return NULL;
+
+	int run = calendar_queue->run;
+	
+	thread_t *thread = NULL;
+	do {
+		thread = calendar_queue->queues[run];
+		thread_t *prev = NULL;
+		while (thread) {
+			if (thread->cputarget) {
+				prev = thread;
+				thread = thread->next;
+				continue;
+			}
+
+			if (prev)
+				prev->next = thread->next;
+			else
+				calendar_queue->queues[run] = thread->next;
+
+			--calendar_queue->thread_count;
+			return thread;
+		}
+
+		run = (run + 1) % SCHED_RUN_QUEUE_SIZE;
+	} while (run != calendar_queue->run);
+
+	// no thread found to steal
+	return NULL;
+}
+
+static thread_t *steal_from_run_queue(sched_run_queue_t *run_queue) {
+	long idx = bitmap_find_first_set(&run_queue->thread_bitmap);
+	if (idx == -1)
+		return NULL;
+
+	for (; idx < SCHED_RUN_QUEUE_SIZE; ++idx) {
+		for (thread_t *search = run_queue->queues[idx].run; search; search = search->next) {
+			if (search->cputarget)
+				continue;
+
+			if (search->next) {
+				search->next->prev = search->prev;
+			} else {
+				run_queue->queues[idx].ins = search->prev;
+			}
+
+			if (search->prev) {
+				search->prev->next = search->next;
+			} else {
+				run_queue->queues[idx].run = search->next;
+			}
+
+			if (run_queue->queues[idx].run == NULL)
+				bitmap_set(&run_queue->thread_bitmap, idx, 0);
+
+			return search;
+		}
+	}
+
+	// no thread found in this runqueue in which we can steal (only cpu targetting ones were found)
+	return NULL;
+}
+
 static void insert_in_calendar_queue(sched_calendar_queue_t *calendar_queue, thread_t *thread) {
 	int idx = (calendar_queue->ins + thread->metrics.interactivity_score - SCHED_INTERACTIVITY_LIMIT) % SCHED_RUN_QUEUE_SIZE;
 
@@ -116,6 +181,24 @@ static thread_t *pop_from_calendar_queue(sched_calendar_queue_t *calendar_queue)
 	return thread;
 }
 
+// requires cpu queue to be locked
+thread_t *sched_steal_work_from_cpu(cpu_t *cpu) {
+	thread_t *thread = steal_from_run_queue(&cpu->rt_queue);
+	if (thread)
+		goto found;
+
+	thread = steal_from_calendar_queue(&cpu->ts_queue);
+	if (thread)
+		goto found;
+
+	thread = steal_from_run_queue(&cpu->idle_queue);
+	found:
+
+	--cpu->thread_count;
+	--cpu->stealable_thread_count;
+	return thread;
+}
+
 // requires cpu queue to be locked
 thread_t *sched_select_next_thread(void) {
 	thread_t *thread = pop_from_run_queue(&current_cpu()->rt_queue);
@@ -139,6 +222,9 @@ found:
 
 	current_cpu()->last_interactivity = thread->metrics.interactivity_score;
 	thread->last_cpu = current_cpu();
+	--current_cpu()->thread_count;
+	if (thread->cputarget == NULL)
+		--current_cpu()->stealable_thread_count;
 
 	return thread;
 }
@@ -164,6 +250,10 @@ void sched_insert_in_cpu_queue(cpu_t *cpu, thread_t *thread) {
 			break;
 	}
 
+	++cpu->thread_count;
+	if (thread->cputarget == NULL)
+		++cpu->stealable_thread_count;
+
 	spinlock_acquire(&sched_idle_cpu_bitmap_lock);
 
 	bitmap_set(&sched_idle_cpu_bitmap, cpu->internal_id, 0);
@@ -191,9 +281,63 @@ void sched_queue(thread_t *thread) {
 
 // ran every 10 ms per core
 void sched_calendar_tick(context_t *, dpcarg_t) {
+	interrupt_set(false);
 	spinlock_acquire(&current_cpu()->sched_lock);
 
 	current_cpu()->ts_queue.ins = (current_cpu()->ts_queue.ins + 1) % SCHED_RUN_QUEUE_SIZE;
 
 	spinlock_release(&current_cpu()->sched_lock);
 }
+
+// ran every second TODO randomize
+// takes a single thread from the most loaded cpu and gives it to the least loaded cpu
+// note that there has to be a migrateable thread for this to be possible. 
+// (so if a thread targets the cpu it is ignored in the count)
+void sched_load_balancer(context_t *, dpcarg_t) {
+	interrupt_set(false);
+	cpu_t *most_loaded = NULL;
+	cpu_t *least_loaded = NULL;
+
+	// find the least loaded and most loaded cpus
+	for (int i = 0; i < arch_smp_cpusawake; ++i) {
+		cpu_t *old_most_loaded = NULL;
+		cpu_t *old_least_loaded = NULL;
+
+		spinlock_acquire(&smp_cpus[i]->sched_lock);
+		if ((most_loaded == NULL || most_loaded->thread_count < smp_cpus[i]->thread_count) && smp_cpus[i]->stealable_thread_count) {
+			old_most_loaded = most_loaded;
+			most_loaded = smp_cpus[i];
+		}
+
+		if (least_loaded == NULL || least_loaded->thread_count > smp_cpus[i]->thread_count) {
+			old_least_loaded = least_loaded;
+			least_loaded = smp_cpus[i];
+		}
+
+		if (old_most_loaded)
+			spinlock_release(&old_most_loaded->sched_lock);
+
+		if (old_least_loaded && old_least_loaded != most_loaded && old_least_loaded != old_most_loaded)
+			spinlock_release(&old_least_loaded->sched_lock);
+
+		if (old_least_loaded == NULL && old_most_loaded == NULL)
+			spinlock_release(&smp_cpus[i]->sched_lock);
+	}
+
+	if (most_loaded == NULL || most_loaded->thread_count <= 1 || most_loaded->thread_count == least_loaded->thread_count)
+		goto leave;
+
+	thread_t *thread = sched_steal_work_from_cpu(most_loaded);
+	__assert(thread);
+
+	sched_insert_in_cpu_queue(least_loaded, thread);
+
+	if (sched_thread_can_run_in_cpu(thread, least_loaded->last_queue, least_loaded->last_interactivity))
+		sched_preempt_cpu(least_loaded);
+
+	leave:
+	if (most_loaded)
+		spinlock_release(&most_loaded->sched_lock);
+	if (most_loaded != least_loaded)
+		spinlock_release(&least_loaded->sched_lock);
+}

kernel-src/sys/scheduler/scheduler.c

@@ -295,9 +295,11 @@ static void reschedule_yield(context_t *context, void *_cpu) {
 
 	spinlock_acquire(&current_cpu()->sched_lock);
 	thread_t *next = sched_select_next_thread();
-	spinlock_release(&current_cpu()->sched_lock);
 
-	arch_smp_send_ipi(cpu, cpu->reschedule_isr, ARCH_SMP_IPI_TARGET, false);
+	if (sched_thread_can_run_in_cpu(thread, cpu->last_queue, cpu->last_interactivity))
+		arch_smp_send_ipi(cpu, cpu->reschedule_isr, ARCH_SMP_IPI_TARGET, false);
+
+	spinlock_release(&current_cpu()->sched_lock);
 	switch_thread(next);
 }
 
@@ -341,6 +343,7 @@ static void set_up_bitmaps(void) {
 }
 
 void sched_calendar_tick(context_t *, dpcarg_t);
+void sched_load_balancer(context_t *, dpcarg_t);
 
 void sched_ap_entry() {
 	SPINLOCK_INIT(current_cpu()->sched_lock);
@@ -367,6 +370,7 @@ void sched_ap_entry() {
 }
 
 void sched_init() {
+	static timerentry_t sched_load_balancer_entry;
 	proc_init();
 
 	SPINLOCK_INIT(current_cpu()->sched_lock);
@@ -399,6 +403,8 @@ void sched_init() {
 
 	timer_insert(current_cpu()->timer, &current_cpu()->schedtimerentry, reschedule_timer_dpc, NULL, QUANTUM_US, true);
 	timer_insert(current_cpu()->timer, &current_cpu()->calendar_tick_timer_entry, sched_calendar_tick, NULL, 10000, true);
+	// TODO randomize interval
+	timer_insert(current_cpu()->timer, &sched_load_balancer_entry, sched_load_balancer, NULL, 1000000, true);
 	// XXX move this resume to a more appropriate place
 	timer_resume(current_cpu()->timer);
 }