diff --git a/main.c b/main.c
index 505ca6c..a49a4e2 100644
--- a/main.c
+++ b/main.c
@@ -274,8 +274,16 @@
 		ret = hf_vcpu_run(vcpu->vm->id, vcpu->vcpu_index);
 
 		switch (ret.code) {
-		/* Yield (forcibly or voluntarily). */
+		/* Preempted. */
+		case HF_VCPU_RUN_PREEMPTED:
+			if (need_resched())
+				schedule();
+			break;
+
+		/* Yield. */
 		case HF_VCPU_RUN_YIELD:
+			if (!kthread_should_stop())
+				schedule();
 			break;
 
 		 /* WFI. */