From 388e6388a9aee45ad6bc75c9ffc5c9d0b83c6812 Mon Sep 17 00:00:00 2001
From: Scott Haseley <shaseley@buffalo.edu>
Date: Thu, 7 Apr 2016 14:35:09 -0400
Subject: [PATCH] Added thread_wait_for_count() which allows the menu to wait
 for all (new) threads to exit before continuing. We need this for khu testing
 because we have to make sure cleanup is done before getting the used byte
 count.

---
 kern/include/thread.h |  2 ++
 kern/main/menu.c      |  7 +++++++
 kern/thread/thread.c  | 37 +++++++++++++++++++++++++++++++++++++
 3 files changed, 46 insertions(+)

diff --git a/kern/include/thread.h b/kern/include/thread.h
index 42d56b4..04a46d3 100644
--- a/kern/include/thread.h
+++ b/kern/include/thread.h
@@ -178,5 +178,7 @@ void schedule(void);
  */
 void thread_consider_migration(void);
 
+extern unsigned thread_count;
+void thread_wait_for_count(unsigned);
 
 #endif /* _THREAD_H_ */
diff --git a/kern/main/menu.c b/kern/main/menu.c
index 057438a..2451b13 100644
--- a/kern/main/menu.c
+++ b/kern/main/menu.c
@@ -117,6 +117,7 @@ common_prog(int nargs, char **args)
 {
 	struct proc *proc;
 	int result;
+	unsigned tc;
 
 	/* Create a process for the new program to run in. */
 	proc = proc_create_runprogram(args[0] /* name */);
@@ -124,6 +125,8 @@ common_prog(int nargs, char **args)
 		return ENOMEM;
 	}
 
+	tc = thread_count;
+
 	result = thread_fork(args[0] /* thread name */,
 			proc /* new process */,
 			cmd_progthread /* thread function */,
@@ -139,6 +142,10 @@ common_prog(int nargs, char **args)
 	 * once you write the code for handling that.
 	 */
 
+	// Wait for all threads to finish cleanup, otherwise khu be a bit behind,
+	// especially once swapping is enabled.
+	thread_wait_for_count(tc);
+
 	return 0;
 }
 
diff --git a/kern/thread/thread.c b/kern/thread/thread.c
index 7a06d07..8ceeaf5 100644
--- a/kern/thread/thread.c
+++ b/kern/thread/thread.c
@@ -70,6 +70,11 @@ unsigned num_cpus;
 /* Used to wait for secondary CPUs to come online. */
 static struct semaphore *cpu_startup_sem;
 
+/* Used to synchronize exit cleanup. */
+unsigned thread_count = 0;
+static struct spinlock thread_count_lock = SPINLOCK_INITIALIZER;
+static struct wchan *thread_count_wchan;
+
 ////////////////////////////////////////////////////////////
 
 /*
@@ -425,6 +430,7 @@ thread_start_cpus(void)
 	kprintf("cpu0: %s\n", buf);
 
 	cpu_startup_sem = sem_create("cpu_hatch", 0);
+	thread_count_wchan = wchan_create("thread_count");
 	mainbus_start_cpus();
 
 	num_cpus = cpuarray_num(&allcpus);
@@ -438,6 +444,12 @@ thread_start_cpus(void)
 		kprintf("%d CPUs online\n", i + 1);
 	}
 	cpu_startup_sem = NULL;
+
+	// Gross hack to deal with os/161 "idle" threads. Hardcode the thread count
+	// to 1 so the inc/dec properly works in thread_[fork/exit]. The one thread
+	// is the cpu0 boot thread (menu), which is the only thread that hasn't
+	// exited yet.
+	thread_count = 1;
 }
 
 /*
@@ -536,6 +548,11 @@ thread_fork(const char *name,
 	 */
 	newthread->t_iplhigh_count++;
 
+	spinlock_acquire(&thread_count_lock);
+	++thread_count;
+	wchan_wakeall(thread_count_wchan, &thread_count_lock);
+	spinlock_release(&thread_count_lock);
+
 	/* Set up the switchframe so entrypoint() gets called */
 	switchframe_init(newthread, entrypoint, data1, data2);
 
@@ -799,6 +816,14 @@ thread_exit(void)
 	/* Check the stack guard band. */
 	thread_checkstack(cur);
 
+	// Decrement the thread count and notify anyone interested.
+	if (thread_count) {
+		spinlock_acquire(&thread_count_lock);
+		--thread_count;
+		wchan_wakeall(thread_count_wchan, &thread_count_lock);
+		spinlock_release(&thread_count_lock);
+	}
+
 	/* Interrupts off on this processor */
 	splhigh();
 	thread_switch(S_ZOMBIE, NULL, NULL);
@@ -1196,3 +1221,15 @@ interprocessor_interrupt(void)
 	curcpu->c_ipi_pending = 0;
 	spinlock_release(&curcpu->c_ipi_lock);
 }
+
+/*
+ * Wait for the thread count to equal tc.
+ */
+void thread_wait_for_count(unsigned tc)
+{
+	spinlock_acquire(&thread_count_lock);
+	while (thread_count != tc) {
+		wchan_sleep(thread_count_wchan, &thread_count_lock);
+	}
+	spinlock_release(&thread_count_lock);
+}