sched: use spinlock for atomic context

arch/at91sam3x8e/sched.c

@@ -11,13 +11,15 @@
 #include <stdbool.h>
 #include <toolchain.h>
 
+#include <ardix/spinlock.h>
+
 void irq_sys_tick(void)
 {
 	/*
 	 * fire a PendSV interrupt and do the actual context switching there
 	 * because it is faster that way (according to the docs, at least)
 	 */
-	if (!_is_atomic_context)
+	if (!spinlock_count(&_in_atomic_context))
 		arch_irq_invoke(IRQNO_PEND_SV);
 }
 

kernel/sched.c

@@ -51,7 +51,7 @@ static inline bool sched_proc_should_run(const struct process *proc)
 {
 	enum proc_state state = proc->state;
 
-	if (state == PROC_QUEUE || state == PROC_READY)
+	if (state == PROC_QUEUE || state == PROC_READY || state == PROC_IOWAIT)
 		return true;
 
 	return false;
@@ -68,14 +68,13 @@ void *sched_process_switch(void *curr_sp)
 	while (1) {
 		nextpid++;
 		nextpid %= CONFIG_SCHED_MAXPROC;
-		if (proc_table[nextpid] != NULL && proc_table[nextpid]->state == PROC_QUEUE) {
+		if (proc_table[nextpid] != NULL && sched_proc_should_run(proc_table[nextpid])) {
 			_current_process = proc_table[nextpid];
 			break;
 		}
 		/* TODO: Add idle thread */
 	}
 
-	_current_process = proc_table[nextpid];
 	_current_process->state = PROC_READY;
 	return _current_process->sp;
 }