Fix thread reference counting

sys/include/thread.h

@@ -31,6 +31,7 @@ typedef struct Thread {
     uintptr_t		kstack;
     uintptr_t		ustack;
     uint64_t		tid;
+    uint64_t		refCount;
     // Process
     struct Process	*proc;
     TAILQ_ENTRY(Thread)	threadList;
@@ -55,6 +56,7 @@ typedef struct Process {
     AS				*space;
     uintptr_t			ustackNext;	// Next user stack
     TAILQ_ENTRY(Process)	processList;
+    uint64_t			refCount;
     // Process
     Process			*parent;
     TAILQ_ENTRY(Process)	siblingList;

sys/kern/syscall.c

@@ -28,7 +28,11 @@ uint64_t
 Syscall_GetPID()
 {
     Thread *cur = Sched_Current();
-    return cur->tid;
+    uint64_t pid = cur->proc->pid;
+
+    Thread_Release(cur);
+
+    return pid;
 }
 
 void
@@ -37,6 +41,7 @@ Syscall_Exit(uint64_t status)
     Thread *cur = Sched_Current();
 
     Sched_SetZombie(cur);
+    Thread_Release(cur);
     Sched_Scheduler();
 
     // Should not return
@@ -121,6 +126,7 @@ Syscall_MMap(uint64_t addr, uint64_t len, uint64_t prot)
     bool status;
 
     status = PMap_AllocMap(cur->space, addr, pgs, PTE_W);
+    Thread_Release(cur);
     if (!status) {
 	// XXX: Need to unmap PMap_Unmap(cur->space, addr, pgs);
 	return 0;
@@ -141,6 +147,7 @@ Syscall_MUnmap(uint64_t addr, uint64_t len)
     }
 
     PMap_Unmap(cur->space, addr, len /= PGSIZE);
+    Thread_Release(cur);
 
     return 0;
 }
@@ -156,37 +163,55 @@ Syscall_MProtect(uint64_t addr, uint64_t len, uint64_t prot)
 uint64_t
 Syscall_Read(uint64_t fd, uint64_t addr, uint64_t off, uint64_t length)
 {
+    uint64_t status;
     Thread *cur = Sched_Current();
     Handle *handle = Handle_Lookup(cur->proc, fd);
 
-    if (handle == NULL)
+    if (handle == NULL) {
-	return -EBADF;
+	status = -EBADF;
+    } else {
+	status = (handle->read)(handle, (void *)addr, off, length);
+    }
 
-    return (handle->read)(handle, (void *)addr, off, length);
+    Thread_Release(cur);
+
+    return status;
 }
 
 uint64_t
 Syscall_Write(uint64_t fd, uint64_t addr, uint64_t off, uint64_t length)
 {
+    uint64_t status;
     Thread *cur = Sched_Current();
     Handle *handle = Handle_Lookup(cur->proc, fd);
 
-    if (handle == NULL)
+    if (handle == NULL) {
-	return -EBADF;
+	status = -EBADF;
+    } else {
+	status = (handle->write)(handle, (void *)addr, off, length);
+    }
 
-    return (handle->write)(handle, (void *)addr, off, length);;
+    Thread_Release(cur);
+
+    return status;
 }
 
 uint64_t
 Syscall_Flush(uint64_t fd)
 {
+    uint64_t status;
     Thread *cur = Sched_Current();
     Handle *handle = Handle_Lookup(cur->proc, fd);
 
-    if (handle == NULL)
+    if (handle == NULL) {
-	return -EBADF;
+	status = -EBADF;
+    } else {
+	status = (handle->flush)(handle);
+    }
 
-    return (handle->flush)(handle);
+    Thread_Release(cur);
+
+    return status;
 }
 
 // XXX: Cleanup
@@ -195,41 +220,57 @@ Handle *Console_OpenHandle();
 uint64_t
 Syscall_Open(uint64_t user_path, uint64_t flags)
 {
+    uint64_t handleNo;
     Thread *cur = Sched_Current();
     int status;
     char path[256];
 
     status = CopyStrIn(user_path, &path, sizeof(path));
-    if (status != 0)
+    if (status != 0) {
+	Thread_Release(cur);
 	return status;
+    }
 
     if (strncmp("/dev/", path, 5) == 0) {
 	if (strcmp("/dev/console", path) == 0) {
 	    Handle *handle = Console_OpenHandle();
-	    return Handle_Add(cur->proc, handle);
+	    handleNo = Handle_Add(cur->proc, handle);
+	    Thread_Release(cur);
+	    return handleNo;
 	}
 
+	Thread_Release(cur);
 	return -ENOENT;
     }
 
     Handle *handle;
     status = VFSUIO_Open(path, &handle);
-    if (status != 0)
+    if (status != 0) {
+	Thread_Release(cur);
 	return status;
+    }
 
-    return Handle_Add(cur->proc, handle);
+    handleNo = Handle_Add(cur->proc, handle);
+    Thread_Release(cur);
+    return handleNo;
 }
 
 uint64_t
 Syscall_Close(uint64_t fd)
 {
+    uint64_t status;
     Thread *cur = Sched_Current();
     Handle *handle = Handle_Lookup(cur->proc, fd);
 
-    if (handle == NULL)
+    if (handle == NULL) {
-	return -EBADF;
+	status = -EBADF;
+    } else {
+	status = (handle->close)(handle);
+    }
 
-    return (handle->close)(handle);
+    Thread_Release(cur);
+
+    return status;
 }
 
 uint64_t
@@ -267,23 +308,35 @@ Syscall_ReadDir(uint64_t fd, char *user_buf, size_t len, uintptr_t user_off)
     Handle *handle = Handle_Lookup(cur->proc, fd);
     uint64_t offset;
 
-    if (handle == NULL)
+    if (handle == NULL) {
+	Thread_Release(cur);
 	return -EBADF;
+    }
 
     status = CopyIn(user_off, &offset, sizeof(offset));
-    if (status != 0)
+    if (status != 0) {
+	Thread_Release(cur);
 	return status;
+    }
 
-    if (handle->type != HANDLE_TYPE_FILE)
+    if (handle->type != HANDLE_TYPE_FILE) {
+	Thread_Release(cur);
 	return -ENOTDIR;
+    }
 
     rstatus = VFS_ReadDir(handle->vnode, user_buf, len, &offset);
-    if (rstatus < 0)
+    if (rstatus < 0) {
+	Thread_Release(cur);
 	return rstatus;
+    }
 
     status = CopyOut(&offset, user_off, sizeof(offset));
-    if (status != 0)
+    if (status != 0) {
+	Thread_Release(cur);
 	return status;
+    }
+
+    Thread_Release(cur);
 
     return rstatus;
 }
@@ -294,6 +347,8 @@ Syscall_ThreadCreate(uint64_t rip, uint64_t arg)
     uint64_t threadId;
     Thread *curThread = Sched_Current();
     Thread *newThread = Thread_UThreadCreate(curThread, rip, arg);
+
+    Thread_Release(curThread);
     if (newThread == NULL) {
 	return -1;
     }
@@ -314,6 +369,7 @@ Syscall_ThreadExit(uint64_t status)
 
     Sched_SetZombie(cur);
     Semaphore_Release(&cur->proc->zombieSemaphore);
+    Thread_Release(cur);
     Sched_Scheduler();
 
     // Should not return
@@ -328,6 +384,7 @@ ThreadWakeupHelper(void *arg)
     Sched_SetRunnable(thr);
     KTimer_Release(thr->timerEvt);
     thr->timerEvt = NULL;
+    Thread_Release(thr);
 }
 
 uint64_t
@@ -345,6 +402,8 @@ Syscall_ThreadSleep(uint64_t time)
     }
     Sched_Scheduler();
 
+    Thread_Release(cur);
+
     return 0;
 }
 
@@ -366,8 +425,10 @@ Syscall_ThreadWait(uint64_t tid)
     while (1) {
 	Semaphore_Acquire(&cur->proc->zombieSemaphore);
 	status = Thread_Wait(cur, tid);
-	if (status != 0)
+	if (status != 0) {
+	    Thread_Release(cur);
 	    return status;
+	}
 	Semaphore_Release(&cur->proc->zombieSemaphore);
 	Sched_Scheduler();
     }

sys/kern/thread.c

@@ -73,6 +73,7 @@ Process_Create()
 
     proc->pid = nextProcessID++;
     proc->threads = 0;
+    proc->refCount = 1;
     TAILQ_INIT(&proc->threadList);
 
     proc->space = PMap_NewAS();
@@ -105,11 +106,17 @@ Process_Lookup(uint64_t pid)
 void
 Process_Retain(Process *proc)
 {
+    ASSERT(proc->refCount != 0);
+    __sync_fetch_and_add(&proc->refCount, 1);
 }
 
 void
 Process_Release(Process *proc)
 {
+    ASSERT(proc->refCount != 0);
+    if (__sync_fetch_and_sub(&proc->refCount, 1) == 1) {
+	Process_Destroy(proc);
+    }
 }
 
 uint64_t
@@ -147,6 +154,7 @@ Thread_Create(Process *proc)
 
 	thr->ustack = proc->ustackNext;
 	proc->ustackNext += MEM_USERSPACE_STKLEN;
+	Process_Retain(proc);
     }
 
     if ((proc != NULL) && (proc->space != NULL)) {
@@ -163,6 +171,7 @@ Thread_Create(Process *proc)
 
     thr->schedState = SCHED_STATE_NULL;
     thr->timerEvt = NULL;
+    thr->refCount = 1;
 
     Thread_InitArch(thr);
     // Initialize queue
@@ -202,6 +211,7 @@ Thread_UThreadCreate(Thread *oldThr, uint64_t rip, uint64_t arg)
 
     thr->space = oldThr->space;
     thr->schedState = SCHED_STATE_NULL;
+    thr->refCount = 1;
 
     thr->ustack = proc->ustackNext;
     proc->ustackNext += MEM_USERSPACE_STKLEN;
@@ -242,11 +252,17 @@ Thread_Lookup(Process *proc, uint64_t tid)
 void
 Thread_Retain(Thread *thr)
 {
+    ASSERT(thr->refCount != 0);
+    __sync_fetch_and_add(&thr->refCount, 1);
 }
 
 void
 Thread_Release(Thread *thr)
 {
+    ASSERT(thr->refCount != 0);
+    if (__sync_fetch_and_sub(&thr->refCount, 1) == 1) {
+	Thread_Destroy(thr);
+    }
 }
 
 uint64_t
@@ -265,7 +281,7 @@ Thread_Wait(Thread *thr, uint64_t tid)
 
 	TAILQ_REMOVE(&thr->proc->zombieQueue, t, schedQueue);
 	status = t->exitValue;
-	Thread_Destroy(t);
+	Thread_Release(t);
 	return status;
     }
 
@@ -273,7 +289,7 @@ Thread_Wait(Thread *thr, uint64_t tid)
 	if (t->tid == tid) {
 	    TAILQ_REMOVE(&thr->proc->zombieQueue, t, schedQueue);
 	    status = t->exitValue;
-	    Thread_Destroy(t);
+	    Thread_Release(t);
 	    return status;
 	}
     }
@@ -288,7 +304,14 @@ Thread_Wait(Thread *thr, uint64_t tid)
 Thread *
 Sched_Current()
 {
-    return curProc;
+    Spinlock_Lock(&schedLock);
+
+    Thread *thr = curProc;
+    Thread_Retain(thr);
+
+    Spinlock_Unlock(&schedLock);
+
+    return thr;
 }
 
 void
@@ -404,6 +427,7 @@ Process_Dump(Process *proc)
     kprintf("pid        %llu\n", proc->pid);
     kprintf("space      %016llx\n", proc->space);
     kprintf("threads    %llu\n", proc->threads);
+    kprintf("refCount   %d\n", proc->refCount);
     kprintf("nextFD     %llu\n", proc->nextFD);
 }
 
@@ -414,6 +438,7 @@ Thread_Dump(Thread *thr)
     kprintf("space      %016llx\n", thr->space);
     kprintf("kstack     %016llx\n", thr->kstack);
     kprintf("tid        %llu\n", thr->tid);
+    kprintf("refCount   %d\n", thr->refCount);
     kprintf("state      %d\n", thr->schedState);
     kprintf("ctxswtch   %llu\n", thr->ctxSwitches);
     kprintf("utime      %llu\n", thr->userTime);
@@ -459,6 +484,7 @@ Debug_Processes(int argc, const char *argv[])
     TAILQ_FOREACH(proc, &processList, processList)
     {
 	kprintf("Process: %d(%016llx)\n", proc->pid, proc);
+	Process_Dump(proc);
     }
 
     //Spinlock_Unlock(&threadLock);