From dce77a7494f18a010afacd27a1147667c2c3fe29 Mon Sep 17 00:00:00 2001
From: "Fletcher T. Penney" <fletcher@fletcherpenney.net>
Date: Sat, 18 Feb 2017 16:15:25 -0500
Subject: [PATCH] ADDED: Improve token_pool memory handling

---
 src/main.c  | 15 ++++++++++++---
 src/token.c | 24 +++++++++++++++++++-----
 src/token.h |  7 +++++--
 3 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/src/main.c b/src/main.c
index 33ebeb8..fceebfe 100644
--- a/src/main.c
+++ b/src/main.c
@@ -322,7 +322,7 @@ int main(int argc, char** argv) {
 	FILE * output_stream;
 	char * output_filename;
 
-	// Prepare token pool
+	// Increment counter and prepare token pool
 #ifdef kUseObjectPool
 	token_pool_init();
 #endif
@@ -333,7 +333,6 @@ int main(int argc, char** argv) {
 		// Batch process 1 or more files
 		for (int i = 0; i < a_file->count; ++i)
 		{
-			token_pool_drain();
 
 			buffer = scan_file(a_file->filename[i]);
 
@@ -366,6 +365,11 @@ int main(int argc, char** argv) {
 	
 				// Don't free folder -- owned by dirname
 			}
+
+			// Increment counter and prepare token pool
+#ifdef kUseObjectPool
+			token_pool_init();
+#endif
 	
 			if (FORMAT_MMD == format) {
 				result = buffer->str;
@@ -387,6 +391,9 @@ int main(int argc, char** argv) {
 			if (FORMAT_MMD != format) {
 				free(result);
 			}
+
+			// Decrement counter and drain
+			token_pool_drain();
 		}
 	} else {
 		if (a_file->count) {
@@ -456,7 +463,9 @@ int main(int argc, char** argv) {
 
 exit:
 
-	// Clean up token pool
+	// Decrement counter and clean up token pool
+	token_pool_drain();
+	
 #ifdef kUseObjectPool
 	token_pool_free();
 #endif
diff --git a/src/token.c b/src/token.c
index ffb8750..c853141 100644
--- a/src/token.c
+++ b/src/token.c
@@ -69,7 +69,11 @@
 
 #include "object_pool.h"
 
-static pool * token_pool = NULL;
+static pool * token_pool = NULL;		//!< Pointer to our object pool
+
+/// Count number of uses of this pool to allow us know
+/// when it's safe to drain the pool
+static short token_pool_count = 0;
 
 /// Intialize object pool for token allocation
 void token_pool_init(void) {
@@ -77,19 +81,30 @@ void token_pool_init(void) {
 		// No pool exists
 		token_pool = pool_new(sizeof(token));
 	}
+
+	// Increment counter
+	token_pool_count++;
 }
 
 
 /// Drain token allocator pool to prepare for another parse
 void token_pool_drain(void) {
-	pool_drain(token_pool);
+	// Decrement counter
+	token_pool_count--;
+
+	if (token_pool_count == 0)
+		pool_drain(token_pool);
 }
 
 
 /// Free token allocator pool
 void token_pool_free(void) {
-	pool_free(token_pool);
-	token_pool = NULL;
+	if (token_pool_count == 0) {
+		pool_free(token_pool);
+		token_pool = NULL;
+	} else {
+		fprintf(stderr, "ERROR: Attempted to drain token pool while still in use.\n");
+	}
 }
 
 #endif
@@ -102,7 +117,6 @@ token * token_new(unsigned short type, size_t start, size_t len) {
 #ifdef kUseObjectPool
 	token * t = pool_allocate_object(token_pool);
 #else
-	//token * t = calloc(1, sizeof(token));
 	token * t = malloc(sizeof(token));
 #endif
 
diff --git a/src/token.h b/src/token.h
index 080f635..12de1bf 100644
--- a/src/token.h
+++ b/src/token.h
@@ -59,9 +59,12 @@
 
 
 #define kUseObjectPool 1		//!< Use an object pool to allocate tokens to improve
-										//!< performance in memory allocation. Frees all
-										//!< tokens at once, however, at end of parsing.
+								//!< performance in memory allocation. Frees all
+								//!< tokens at once, however, at end of parsing.
 
+/// Should call init() once per thread/use, and drain() once per thread/use.
+/// This allows us to know when the pool is no longer being used and it is safe
+/// to free.
 
 #ifdef kUseObjectPool
 void token_pool_init(void);				//!< Initialize object pool for allocating tokens
-- 
2.40.0