return task back to cfs after plane is done

kernel/sched/pb.c

@@ -157,30 +157,8 @@ EXPORT_SYMBOL(init_pb_rq);
 static void enqueue_task_pb(struct rq *rq, struct task_struct *p, int flags)
 {
 	struct pb_rq *pb = &(rq->pb);
-	bool premature_finish = false;
 
 	pb->waiting_on_io = 0;
-
-	pb->c_entry++;
-
-	/**
-		* Don't schedule a task that is dead. (e.g. plan was incorrect and program finished quicker)
-		* TODO: if we have multiple tasks structs try the next plan entry
-		*/
-	if (pb->c_entry < pb->size && pb->plan[pb->c_entry].task_struct->state == TASK_DEAD) {
-		premature_finish = true;
-		pb->is_initialized = 0;
-	}
-
-	if (pb->c_entry >= pb->size || premature_finish) {
-		if (premature_finish) {
-			printk(KERN_WARNING "PLAN TERMINATED PREMATURELY \n");
-		}
-		else {
-			printk(KERN_WARNING "PLAN DONE \n");
-		}
-	}
-	printk("DONE");
 }
 
 // task started IO and thus it is finished
@@ -191,6 +169,8 @@ static void dequeue_task_pb(struct rq *rq, struct task_struct *p, int flags)
 	u64 perf_counter;
 	u64 counter_diff;
 	u64 read_error;
+	bool premature_finish = false;
+
 
 	if (pb->waiting_on_io) {
 		return;
@@ -221,6 +201,31 @@ static void dequeue_task_pb(struct rq *rq, struct task_struct *p, int flags)
 
 		printk(KERN_WARNING "PB TASK %llu RAN %llu INSTRUCTIONS TOO LONG\n", pb->plan[pb->c_entry].task_id, over_time);
 	}
+
+
+	pb->c_entry++;
+
+	/**
+		* Don't schedule a task that is dead. (e.g. plan was incorrect and program finished quicker)
+		* TODO: if we have multiple tasks structs try the next plan entry
+		*/
+	if (pb->c_entry < pb->size && pb->plan[pb->c_entry].task_struct->state == TASK_DEAD) {
+		premature_finish = true;
+	}
+
+	if (pb->c_entry >= pb->size || premature_finish) {
+		if (premature_finish) {
+			printk(KERN_WARNING "PLAN TERMINATED PREMATURELY \n");
+		}
+		else {
+			printk(KERN_WARNING "PLAN DONE \n");
+		}
+
+		// set back to cfs for completion of task
+		pb->is_initialized = 0;
+		pb->plan[0].task_struct->sched_class = &fair_sched_class;
+		resched_curr(rq);
+	}
 }
 
 static void yield_task_pb(struct rq *rq)