pkmem: clear dbg and use mtx on delete

1 files changed, 21 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