From 14705bb277cc3f7f16ee27ceac040a90151a87de Mon Sep 17 00:00:00 2001
From: Nick Krichevsky <nick@ollien.com>
Date: Thu, 26 Sep 2019 12:23:24 -0400
Subject: [PATCH] Fix memory leak when nb sends fail

---
 adder.c |  5 ++++-
 main.c  | 22 ++++++++++++++--------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/adder.c b/adder.c
index 9114590..eab14c1 100644
--- a/adder.c
+++ b/adder.c
@@ -13,10 +13,10 @@
  */
 void *adder(void *raw_mailbox_id) {
 	int mailbox_id = *(int*)(raw_mailbox_id);
-	int *ret_val = malloc(sizeof(int));
 	// If we're thread 0, we shouldn't be running adder.
 	// Any thread id below 0 is invalid.
 	if (mailbox_id <= 0) {
+		int *ret_val = malloc(sizeof(int));
 		*ret_val = -1;
 		return ret_val;
 	}
@@ -28,6 +28,7 @@ void *adder(void *raw_mailbox_id) {
 	while (true) {
 		int recv_result = RecvMsg(mailbox_id, &message);
 		if (recv_result != 0) {
+			int *ret_val = malloc(sizeof(int));
 			*ret_val = -1;
 			return ret_val;
 		}
@@ -52,10 +53,12 @@ void *adder(void *raw_mailbox_id) {
 
 	int send_result = SendMsg(0, &res_mesage);
 	if (send_result < 0) {
+		int *ret_val = malloc(sizeof(int));
 		*ret_val = -1;
 		return ret_val;
 	}
 
+	int *ret_val = malloc(sizeof(int));
 	*ret_val = 0;
 	return ret_val;
 }
\ No newline at end of file
diff --git a/main.c b/main.c
index f494b0c..6cf99d4 100644
--- a/main.c
+++ b/main.c
@@ -79,10 +79,15 @@ static int get_input(int num_threads, struct input **res) {
  * Kills the given thread. Should only be used in exceptional circumstances.
  *
  * @param thread The thread to kill.
+ * @param free_ret If true, the memory associated with the return value will be freed.
  */
-static void kill_thread(pthread_t thread) {
+static void kill_thread(pthread_t thread, bool free_ret) {
 	pthread_cancel(thread);
-	pthread_join(thread, NULL);
+	void *ret;
+	pthread_join(thread, free_ret ? &ret : NULL);
+	if (free_ret && ret != PTHREAD_CANCELED) {
+		free(ret);
+	}
 }
 
 /**
@@ -90,10 +95,11 @@ static void kill_thread(pthread_t thread) {
  *
  * @param threads An array of threads to kill.
  * @param num_threads The number of threads to kill.
+ * @param free_ret If true, the memory associated with the return value will be freed.
  */
-static void kill_threads(pthread_t *threads, int num_threads) {
+static void kill_threads(pthread_t *threads, int num_threads, bool free_ret) {
 	for (int i = 0; i < num_threads; i++) {
-		kill_thread(threads[i]);
+		kill_thread(threads[i], free_ret);
 	}
 }
 
@@ -142,7 +148,7 @@ static int start_adder_threads(int *thread_ids, int num_threads, pthread_t *thre
 		int create_result = pthread_create(&threads[i], NULL, adder, &thread_ids[i]);
 		if (create_result != 0) {
 			// Kill the i threads that we have spawned.
-			kill_threads(threads, i);
+			kill_threads(threads, i, true);
 
 			return -1;
 		}
@@ -426,7 +432,7 @@ int main(int argc, char *argv[]) {
 	if (spawn_result != 0) {
 		printf("Failed to spawn collection thread.\n");
 		free(inputs);
-		kill_threads(adder_threads, num_adder_threads);
+		kill_threads(adder_threads, num_adder_threads, true);
 		free_mailboxes();
 		return 1;
 	}
@@ -440,8 +446,8 @@ int main(int argc, char *argv[]) {
 
 	free(inputs);
 	if (send_result != 0) {
-		kill_threads(adder_threads, num_adder_threads);
-		kill_thread(collect_thread);
+		kill_threads(adder_threads, num_adder_threads, true);
+		kill_thread(collect_thread, false);
 		free_mailboxes();
 		return 1;
 	}