From feb45381ba84b1aeba28698cf4435b85483aaf81 Mon Sep 17 00:00:00 2001
From: Ali Mashtizadeh <mashti@uwaterloo.ca>
Date: Sun, 3 Sep 2023 14:54:14 -0400
Subject: [PATCH] OSWait: Fix waiting on a specific PID

---
 sys/include/thread.h |  6 ++++++
 sys/kern/process.c   | 42 ++++++++++++++++++++++++++++++++++++------
 sys/kern/sched.c     |  5 +++++
 3 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/sys/include/thread.h b/sys/include/thread.h
index d5b2d36..b4f0fd9 100644
--- a/sys/include/thread.h
+++ b/sys/include/thread.h
@@ -58,6 +58,10 @@ typedef struct Thread {
 #define PROCESS_HANDLE_SLOTS	128
 #define PROCESS_TITLE_LENGTH	128
 
+#define PROC_STATE_NULL		0
+#define PROC_STATE_READY	1
+#define PROC_STATE_ZOMBIE	2
+
 typedef struct Process {
     uint64_t			pid;
     AS				*space;
@@ -68,6 +72,7 @@ typedef struct Process {
     TAILQ_ENTRY(Process)	processList;
     uint64_t			refCount;
     char			title[PROCESS_TITLE_LENGTH];
+    int				procState;
     uint64_t			exitCode;
     // Process
     Process			*parent;
@@ -76,6 +81,7 @@ typedef struct Process {
     ProcessQueue		zombieProc;
     Mutex			zombieProcLock;
     CV				zombieProcCV;
+    CV				zombieProcPCV;
     // Threads
     uint64_t			threads;
     ThreadQueue			threadList;
diff --git a/sys/kern/process.c b/sys/kern/process.c
index 0ace2a2..3449f1a 100644
--- a/sys/kern/process.c
+++ b/sys/kern/process.c
@@ -45,6 +45,7 @@ Process_Create(Process *parent, const char *title)
     proc->pid = nextProcessID++;
     proc->threads = 0;
     proc->refCount = 1;
+    proc->procState = PROC_STATE_NULL;
     TAILQ_INIT(&proc->threadList);
 
     if (title) {
@@ -77,6 +78,7 @@ Process_Create(Process *parent, const char *title)
     TAILQ_INIT(&proc->zombieProc);
     Mutex_Init(&proc->zombieProcLock, "Zombie Process Lock");
     CV_Init(&proc->zombieProcCV, "Zombie Process CV");
+    CV_Init(&proc->zombieProcPCV, "Zombie Process PCV");
 
     Spinlock_Lock(&procLock);
     TAILQ_INSERT_TAIL(&processList, proc, processList);
@@ -92,6 +94,7 @@ Process_Destroy(Process *proc)
 
     Spinlock_Destroy(&proc->lock);
     Semaphore_Destroy(&proc->zombieSemaphore);
+    CV_Destroy(&proc->zombieProcPCV);
     CV_Destroy(&proc->zombieProcCV);
     Mutex_Destroy(&proc->zombieProcLock);
     PMap_DestroyAS(proc->space);
@@ -146,6 +149,7 @@ Process_Release(Process *proc)
     }
 }
 
+// XXX: Cleanup return value
 uint64_t
 Process_Wait(Process *proc, uint64_t pid)
 {
@@ -157,14 +161,33 @@ Process_Wait(Process *proc, uint64_t pid)
     // XXX: Need to verify pid exists!
 
     Mutex_Lock(&proc->zombieProcLock);
-    while (1) {
-	p = TAILQ_FIRST(&proc->zombieProc);
-	if (p != NULL && (pid == 0 || p->pid == pid)) {
-	    TAILQ_REMOVE(&proc->zombieProc, p, siblingList);
-	    break;
+    if (pid == 0) {
+	while (1) {
+	    p = TAILQ_FIRST(&proc->zombieProc);
+	    if (p != NULL) {
+		TAILQ_REMOVE(&proc->zombieProc, p, siblingList);
+		break;
+	    }
+
+	    CV_Wait(&proc->zombieProcCV, &proc->zombieProcLock);
+	}
+    } else {
+	p = Process_Lookup(pid);
+	if (p == NULL) {
+	    Mutex_Unlock(&proc->zombieProcLock);
+	    return ENOENT;
 	}
 
-	CV_Wait(&proc->zombieProcCV, &proc->zombieProcLock);
+	while (p->procState != PROC_STATE_ZOMBIE) {
+	    CV_Wait(&p->zombieProcPCV, &proc->zombieProcLock);
+	}
+	TAILQ_REMOVE(&proc->zombieProc, p, siblingList);
+	/*
+	 * Release the reference from Process_Lookup but it's okay to continue 
+	 * using p because we're cleaning up the last reference as the parent 
+	 * process.
+	 */
+	Process_Release(p);
     }
     Mutex_Unlock(&proc->zombieProcLock);
 
@@ -188,8 +211,15 @@ Process_Wait(Process *proc, uint64_t pid)
 void
 Process_Dump(Process *proc)
 {
+    const char *stateStrings[] = {
+	"NULL",
+	"READY",
+	"ZOMBIE"
+    };
+
     kprintf("title      %s\n", proc->title);
     kprintf("pid        %llu\n", proc->pid);
+    kprintf("state	%s\n", stateStrings[proc->procState]);
     kprintf("space      %016llx\n", proc->space);
     kprintf("threads    %llu\n", proc->threads);
     kprintf("refCount   %d\n", proc->refCount);
diff --git a/sys/kern/sched.c b/sys/kern/sched.c
index 6218f7b..fb3e819 100644
--- a/sys/kern/sched.c
+++ b/sys/kern/sched.c
@@ -46,6 +46,9 @@ Sched_SetRunnable(Thread *thr)
 {
     Spinlock_Lock(&schedLock);
 
+    if (thr->proc->procState == PROC_STATE_NULL)
+	thr->proc->procState = PROC_STATE_READY;
+
     if (thr->schedState == SCHED_STATE_WAITING) {
 	thr->waitTime += KTime_GetEpochNS() - thr->waitStart;
 	thr->waitStart = 0;
@@ -86,9 +89,11 @@ Sched_SetZombie(Thread *thr)
 	ASSERT(proc->parent != NULL);
 	Mutex_Lock(&proc->parent->zombieProcLock);
 	Spinlock_Lock(&proc->parent->lock); // Guards child list
+	proc->procState = PROC_STATE_ZOMBIE;
 	TAILQ_REMOVE(&proc->parent->childrenList, proc, siblingList);
 	TAILQ_INSERT_TAIL(&proc->parent->zombieProc, proc, siblingList);
 	Spinlock_Unlock(&proc->parent->lock);
+	CV_Signal(&proc->zombieProcPCV);
 	CV_Signal(&proc->parent->zombieProcCV);
 	Mutex_Unlock(&proc->parent->zombieProcLock);
     }