pkmem: clear dbg and use mtx on delete

2 files changed, 66 insertions, 1 deletions

diff --git a/pkmem.h b/pkmem.h
index d4672a2..5adc9ea 100644
--- a/pkmem.h
+++ b/pkmem.h
@@ -270,7 +270,8 @@ pk_bucket_create(const char* description, int64_t sz, bool transient)
 void
 pk_bucket_destroy(struct pk_membucket* bkt)
 {
-	for (int64_t i = 0; i < pk_bucket_head; ++i) {
+	int64_t i;
+	for (i = 0; i < pk_bucket_head; ++i) {
 		if (&pk_buckets[i] == bkt) {
 			if (pk_bucket_head == i)
 				pk_bucket_head--;
@@ -287,11 +288,20 @@ pk_bucket_destroy(struct pk_membucket* bkt)
 	bkt->blocks = CAFE_BABE(struct pk_memblock);
 	bkt->ptr = CAFE_BABE(char);
 	bkt->transient = false;
+#ifdef PK_MEMORY_DEBUGGER
+	for (i = debug_alloc_head; i > -1; --i) {
+		if (debug_all_allocs[i].bkt == bkt) {
+			debug_all_allocs[i].blk.data = NULL;
+			debug_all_allocs[i].blk.size = 0u;
+		}
+	}
+#endif
 }
 
 void
 pk_bucket_reset(struct pk_membucket* bkt)
 {
+	int64_t i;
 	if (bkt->transient != true) {
 		PK_LOG_ERR("WARNING: pk_bucket_reset called on non-transient pk_membucket\n");
 	}
@@ -301,6 +311,14 @@ pk_bucket_reset(struct pk_membucket* bkt)
 	bkt->lastEmptyBlockIndex = 0;
 	bkt->blocks->data = bkt->ptr;
 	bkt->blocks->size = bkt->size - (sizeof(struct pk_memblock) * bkt->maxBlockCount);
+#ifdef PK_MEMORY_DEBUGGER
+	for (i = debug_alloc_head; i > -1; --i) {
+		if (debug_all_allocs[i].bkt == bkt) {
+			debug_all_allocs[i].blk.data = NULL;
+			debug_all_allocs[i].blk.size = 0u;
+		}
+	}
+#endif
 }
 
 void
@@ -466,6 +484,7 @@ pk_delete_bkt(const void* ptr, size_t sz, struct pk_membucket* bkt)
 #ifdef PK_MEMORY_FORCE_MALLOC
 	return std::free(const_cast<void*>(ptr));
 #endif
+	mtx_lock(&bkt->mtx);
 	assert(ptr >= bkt->raw && (char*)ptr < bkt->ptr + bkt->size && "pointer not in memory bucket range");
 	assert(sz > 0 && "attempted to free pointer of size 0");
 #ifdef PK_MEMORY_DEBUGGER
@@ -551,6 +570,7 @@ pk_delete_bkt(const void* ptr, size_t sz, struct pk_membucket* bkt)
 		assert(debug_tracked_alloc_size == debug_bucket_alloc_size && "allocation size mismatch!");
 	}
 #endif
+	mtx_unlock(&bkt->mtx);
 }
 
 void
diff --git a/test/pkmem.c b/test/pkmem.c
index b6826a2..2073459 100644
--- a/test/pkmem.c
+++ b/test/pkmem.c
@@ -1,6 +1,21 @@
 
 #include "../pkmem.h"
 
+struct mem_test {
+	struct pk_membucket *bkt1;
+	struct pk_membucket *bkt2;
+} mt;
+
+void spinup() {
+	mt.bkt1 = pk_bucket_create("bkt1", 1024, false);
+	mt.bkt2 = pk_bucket_create("bkt2", 1024, false);
+}
+
+void teardown() {
+	pk_bucket_destroy(mt.bkt1);
+	pk_bucket_destroy(mt.bkt2);
+}
+
 int main(int argc, char *argv[])
 {
 	(void)argc;
@@ -8,9 +23,39 @@ int main(int argc, char *argv[])
 
 	// pk_new_base
 	{
+		spinup();
 		char *some_dang_string = (char*)pk_new_base(64, alignof(char*));
 		fprintf(stdout, "some_dang_string: %p\n", some_dang_string);
 		pk_delete_base(some_dang_string, 64);
+		teardown();
+	}
+
+	/* template
+	{
+		spinup();
+		teardown();
+	}
+	*/
+
+	// zero-length alloc
+	{
+		spinup();
+		char *zero_length_string;
+		char *three_length_string;
+
+		zero_length_string = pk_new_bkt(0, 0, mt.bkt1);
+		three_length_string = pk_new_bkt(3, 0, mt.bkt1);
+		three_length_string[0] = '1';
+		three_length_string[1] = '2';
+
+		if (zero_length_string != NULL) exit(1);
+		if (strlen(three_length_string) != 2) exit(1);
+		if (mt.bkt1->allocs != 1) exit(1);
+		if (mt.bkt1->head != 3) exit(1);
+		if (mt.bkt1->lastEmptyBlockIndex != 0) exit(1);
+
+		PK_LOGV_INF("%s: %s\n", __FILE__, "handles zero-length alloc");
+		teardown();
 	}
 
 	return 0;