Add MallocCount::setAllocLimit() function

Sming/Components/malloc_count/include/malloc_count.h

@@ -34,30 +34,53 @@
 
 namespace MallocCount
 {
-/* returns the currently allocated amount of memory */
+/**
+ *  @brief Get the currently allocated amount of memory
+ */
 size_t getCurrent(void);
 
-/* returns the current peak memory allocation */
+/**
+ *  @brief Get the peak memory allocation
+ */
 size_t getPeak(void);
 
-/* resets the peak memory allocation to current */
+/**
+ * @brief Reset the peak memory allocation to current
+ */
 void resetPeak(void);
 
-/* returns the total number of allocations */
+/**
+ * @brief Get the total number of allocations
+ */
 size_t getAllocCount(void);
 
-/* typedef of callback function */
+/**
+ * @brief Set an allocation limit
+ * @param maxBytes Specify 0 for no limit
+ */
+void setAllocLimit(size_t maxBytes);
+
+/**
+ * @brief Callback function type
+ * @param current Current allocated bytes
+ */
 typedef std::function<void(size_t current)> MallocCountCallback;
 
-/* supply malloc_count with a callback function that is invoked on each change
- * of the current allocation. The callback function must not use malloc()/realloc()/free()
- * or it will go into an endless recursive loop! */
+/**
+ * @brief Set a callback function that is invoked on each change of the current allocation
+ * @note The callback function must not use malloc()/realloc()/free()
+ * or it will go into an endless recursive loop!
+ */
 void setCallback(MallocCountCallback callback);
 
-/* dynamically enable/disable logging */
+/**
+ * @brief Enable/disable logging
+ */
 void enableLogging(bool enable);
 
-/* allocations less than this threshold are never logged */
+/**
+ * @brief Set minimum allocation size for log output (when enabled)
+ */
 void setLogThreshold(size_t threshold);
 
 }; // namespace MallocCount

Sming/Components/malloc_count/malloc_count.cpp

@@ -114,6 +114,8 @@ static struct {
 	size_t count;   // Number of allocations called
 } stats;
 
+static size_t allocationLimit = 0;
+
 static MallocCountCallback userCallback = nullptr;
 
 /* add allocation to statistics */
@@ -153,19 +155,25 @@ size_t getPeak()
 }
 
 /* user function to reset the peak allocation to current */
-extern void resetPeak()
+void resetPeak()
 {
 	stats.peak = stats.current;
 }
 
 /* user function to return total number of allocations */
-extern size_t getAllocCount()
+size_t getAllocCount()
 {
 	return stats.count;
 }
 
+/* sets the maximum available memory */
+void setAllocLimit(size_t maxBytes)
+{
+	allocationLimit = maxBytes;
+}
+
 /* user function to supply a memory profile callback */
-extern void setCallback(MallocCountCallback callback)
+void setCallback(MallocCountCallback callback)
 {
 	userCallback = callback;
 }
@@ -181,6 +189,11 @@ extern "C" void* WRAP(F_MALLOC)(size_t size)
 		return nullptr;
 	}
 
+	if(allocationLimit != 0 && stats.current + size > allocationLimit) {
+		log("malloc(%u) -> exceeds maximum (current is %u bytes)", size, stats.current);
+		return nullptr;
+	}
+
 	/* call read malloc procedure in libc */
 	void* ret = REAL(F_MALLOC)(alignment + size);
 
@@ -246,8 +259,10 @@ extern "C" void* WRAP(F_CALLOC)(size_t nmemb, size_t size)
 		return nullptr;
 	}
 
-	void* ret = malloc(size);
-	memset(ret, 0, size);
+	void* ret = WRAP(F_MALLOC)(size);
+	if(ret != nullptr) {
+		memset(ret, 0, size);
+	}
 	return ret;
 }
 
@@ -271,6 +286,12 @@ extern "C" void* WRAP(F_REALLOC)(void* ptr, size_t size)
 		return REAL(F_REALLOC)(ptr, size);
 	}
 
+	// For testing purposes we'll assume the reallocation won't be in-place
+	if(allocationLimit != 0 && stats.current + size > allocationLimit) {
+		log("remalloc(%u) -> exceeds maximum (current is %u bytes)", size, stats.current);
+		return nullptr;
+	}
+
 	ptr = (char*)ptr - alignment;
 
 	size_t oldsize = *(size_t*)ptr;