sched: enforce ref-counting on current task/thread pointers

sched/core.c

@@ -215,18 +215,19 @@ void schedule_thread_on_cpu(struct thread *thr)
 
 void start_charge_period(void)
 {
-	struct thread *self = current_thread();
+	struct thread *self = get_current_thread();
 	if (!self) {
 		return;
 	}
 
 	self->tr_charge_period_start = get_cycles();
+	put_current_thread(self);
 }
 
 void end_charge_period(void)
 {
 	preempt_disable();
-	struct thread *self = current_thread();
+	struct thread *self = get_current_thread();
 	if (!self) {
 		return;
 	}
@@ -244,6 +245,7 @@ void end_charge_period(void)
 	}
 
 	self->tr_charge_period_start = 0;
+	put_current_thread(self);
 
 	// printk("%llu cycles charged to %s/%u", charge,
 	// self->tr_parent->t_name, self->tr_parent->t_id);

sched/task.c

@@ -282,7 +282,7 @@ struct task *task_from_tid(tid_t id)
 
 void task_exit(int status)
 {
-	struct task *self = current_task();
+	struct task *self = get_current_task();
 	unsigned long flags;
 	task_lock_irqsave(self, &flags);
 	struct task *parent = self->t_parent;
@@ -295,7 +295,7 @@ void task_exit(int status)
 		task_unlock(parent);
 	}
 
-	struct thread *cur_thread = current_thread();
+	struct thread *cur_thread = get_current_thread();
 
 	self->t_state = TASK_STOPPED;
 	cur_thread->tr_state = THREAD_STOPPED;
@@ -340,6 +340,9 @@ void task_exit(int status)
 	       self->t_base.ob_refcount);
 	spin_unlock_irqrestore(handles_lock, flags);
 
+	put_current_thread(cur_thread);
+	put_current_task(self);
+
 	while (1) {
 		schedule(SCHED_NORMAL);
 	}
@@ -410,8 +413,19 @@ struct thread *task_create_thread(struct task *parent)
 	return thread;
 }
 
-struct task *current_task(void)
+struct task *get_current_task(void)
 {
-	struct thread *thr = current_thread();
-	return thr ? thr->tr_parent : NULL;
+	struct thread *thr = get_current_thread();
+	if (!thr) {
+		return NULL;
+	}
+
+	struct task *out = task_ref(thr->tr_parent);
+	put_current_thread(thr);
+	return out;
+}
+
+void put_current_task(struct task *task)
+{
+	task_unref(task);
 }

sched/thread.c

@@ -105,7 +105,7 @@ void thread_free(struct thread *thr)
 {
 }
 
-struct thread *current_thread(void)
+struct thread *get_current_thread(void)
 {
 	struct cpu_data *cpu = get_this_cpu();
 	if (!cpu) {
@@ -113,13 +113,24 @@ struct thread *current_thread(void)
 	}
 
 	struct thread *out = cpu->c_rq.rq_cur;
+	object_ref(&out->tr_base);
+
 	put_cpu(cpu);
 	return out;
 }
 
+void put_current_thread(struct thread *thr)
+{
+	object_unref(&thr->tr_base);
+}
+
 bool need_resched(void)
 {
-	return (current_thread()->tr_flags & THREAD_F_NEED_RESCHED) != 0;
+	struct thread *thr = get_current_thread();
+	bool result = (thr->tr_flags & THREAD_F_NEED_RESCHED) != 0;
+	put_current_thread(thr);
+
+	return result;
 }
 
 int thread_priority(struct thread *thr)
@@ -143,7 +154,7 @@ void thread_awaken(struct thread *thr)
 
 void thread_exit(void)
 {
-	struct thread *self = current_thread();
+	struct thread *self = get_current_thread();
 	unsigned long flags;
 	thread_lock_irqsave(self, &flags);
 	self->tr_state = THREAD_STOPPED;
@@ -153,6 +164,7 @@ void thread_exit(void)
 	       self->tr_parent->t_id,
 	       self->tr_id);
 	thread_unlock_irqrestore(self, flags);
+	put_current_thread(self);
 
 	while (1) {
 		schedule(SCHED_NORMAL);

sched/timer.c

@@ -44,7 +44,7 @@ unsigned long schedule_timeout(unsigned long ticks)
 {
 	struct timer timer;
 
-	struct thread *self = current_thread();
+	struct thread *self = get_current_thread();
 
 	timer.t_entry = QUEUE_ENTRY_INIT;
 	timer.t_expiry = clock_ticks + ticks;
@@ -58,6 +58,7 @@ unsigned long schedule_timeout(unsigned long ticks)
 	schedule(SCHED_NORMAL);
 
 	remove_timer(&timer);
+	put_current_thread(self);
 
 	return 0;
 }

sched/wait.c

@@ -110,7 +110,7 @@ void wakeup_one(struct waitqueue *q)
 
 void sleep_forever(void)
 {
-	struct thread *thr = current_thread();
+	struct thread *thr = get_current_thread();
 	struct runqueue *rq = thr->tr_rq;
 
 	unsigned long flags;
@@ -121,7 +121,17 @@ void sleep_forever(void)
 
 	rq_unlock(rq, flags);
 
-	while (thr->tr_state == THREAD_SLEEPING) {
+	put_current_thread(thr);
+
+	while (1) {
+		thr = get_current_thread();
+		bool sleep = (thr->tr_state == THREAD_SLEEPING);
+		put_current_thread(thr);
+
+		if (!sleep) {
+			break;
+		}
+
 		schedule(SCHED_NORMAL);
 	}
 }