Improve reliability of memory stats

The allocation count is managed atomically and where it actually should
change (for instance, not counting an allocation before its success has
been checked).

Bonus: Improve readability of the pre-pad checks.

2 files changed, 27 insertions, 51 deletions

diff --git a/core/os/memory.cpp b/core/os/memory.cpp
index 069ee48fae..a4a664c9ef 100644
--- a/core/os/memory.cpp
+++ b/core/os/memory.cpp
@@ -29,6 +29,7 @@
 /*************************************************************************/
 #include "memory.h"
 #include "copymem.h"
+#include "core/safe_refcount.h"
 #include "error_macros.h"
 #include <stdio.h>
 #include <stdlib.h>
@@ -43,41 +44,36 @@ void *operator new(size_t p_size, void *(*p_allocfunc)(size_t p_size)) {
 	return p_allocfunc(p_size);
 }
 
-#include <stdio.h>
-
 #ifdef DEBUG_ENABLED
 size_t Memory::mem_usage = 0;
 size_t Memory::max_usage = 0;
+#define PREPAD true
+#else
+#define PREPAD p_pad_align
 #endif
 
-size_t Memory::alloc_count = 0;
+uint64_t Memory::alloc_count = 0;
 
 void *Memory::alloc_static(size_t p_bytes, bool p_pad_align) {
 
-#ifdef DEBUG_ENABLED
-	bool prepad = true;
-#else
-	bool prepad = p_pad_align;
-#endif
-
-	void *mem = malloc(p_bytes + (prepad ? PAD_ALIGN : 0));
-
-	alloc_count++;
+	void *mem = malloc(p_bytes + (PREPAD ? PAD_ALIGN : 0));
 
 	ERR_FAIL_COND_V(!mem, NULL);
 
-	if (prepad) {
-		uint64_t *s = (uint64_t *)mem;
-		*s = p_bytes;
-
-		uint8_t *s8 = (uint8_t *)mem;
+	atomic_increment(&alloc_count);
 
+	if (PREPAD) {
 #ifdef DEBUG_ENABLED
 		mem_usage += p_bytes;
 		if (mem_usage > max_usage) {
 			max_usage = mem_usage;
 		}
 #endif
+
+		uint64_t *s = (uint64_t *)mem;
+		*s = p_bytes;
+
+		uint8_t *s8 = (uint8_t *)mem;
 		return s8 + PAD_ALIGN;
 	} else {
 		return mem;
@@ -92,38 +88,30 @@ void *Memory::realloc_static(void *p_memory, size_t p_bytes, bool p_pad_align) {
 
 	uint8_t *mem = (uint8_t *)p_memory;
 
-#ifdef DEBUG_ENABLED
-	bool prepad = true;
-#else
-	bool prepad = p_pad_align;
-#endif
-
-	if (prepad) {
+	if (PREPAD) {
 		mem -= PAD_ALIGN;
 		uint64_t *s = (uint64_t *)mem;
 
+		if (p_bytes == 0) {
 #ifdef DEBUG_ENABLED
-		mem_usage -= *s;
-		mem_usage += p_bytes;
+			mem_usage -= *s;
 #endif
-
-		if (p_bytes == 0) {
 			free(mem);
 			return NULL;
 		} else {
-			*s = p_bytes;
-
 			mem = (uint8_t *)realloc(mem, p_bytes + PAD_ALIGN);
 			ERR_FAIL_COND_V(!mem, NULL);
 
 			s = (uint64_t *)mem;
-
+#ifdef DEBUG_ENABLED
+			mem_usage -= *s;
+			mem_usage += p_bytes;
+#endif
 			*s = p_bytes;
 
 			return mem + PAD_ALIGN;
 		}
 	} else {
-
 		mem = (uint8_t *)realloc(mem, p_bytes);
 
 		ERR_FAIL_COND_V(mem == NULL && p_bytes > 0, NULL);
@@ -138,27 +126,16 @@ void Memory::free_static(void *p_ptr, bool p_pad_align) {
 
 	uint8_t *mem = (uint8_t *)p_ptr;
 
-#ifdef DEBUG_ENABLED
-	bool prepad = true;
-#else
-	bool prepad = p_pad_align;
-#endif
-
-	alloc_count--;
-
-	if (prepad) {
+	if (PREPAD) {
 		mem -= PAD_ALIGN;
-		uint64_t *s = (uint64_t *)mem;
-
 #ifdef DEBUG_ENABLED
-		mem_usage -= *s;
+		const uint64_t s = *((uint64_t *)mem);
+		mem_usage -= s;
 #endif
-
-		free(mem);
-	} else {
-
-		free(mem);
 	}
+	atomic_decrement(&alloc_count);
+
+	free(mem);
 }
 
 size_t Memory::get_mem_available() {
diff --git a/core/os/memory.h b/core/os/memory.h
index b3eb599955..8c3e26219a 100644
--- a/core/os/memory.h
+++ b/core/os/memory.h
@@ -48,8 +48,7 @@ class Memory {
 	static size_t mem_usage;
 	static size_t max_usage;
 #endif
-
-	static size_t alloc_count;
+	static uint64_t alloc_count;
 
 public:
 	static void *alloc_static(size_t p_bytes, bool p_pad_align = false);