From c70c6ccec5fc91f721fc15ce8968d69fb6ffa159 Mon Sep 17 00:00:00 2001
From: Matthew Kennedy <matthewkennedy@outlook.com>
Date: Sat, 4 Mar 2023 03:33:08 -0500
Subject: [PATCH] lua memory leak detection and prevention

(cherry picked from commit cfa9bfa44dd0964f3d9efacfc5a1efbdce26d3d5)
---
 firmware/controllers/lua/lua.cpp | 42 ++++++++++++++++++++++++++------
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/firmware/controllers/lua/lua.cpp b/firmware/controllers/lua/lua.cpp
index d69705b1e3..0b7ed0a984 100644
--- a/firmware/controllers/lua/lua.cpp
+++ b/firmware/controllers/lua/lua.cpp
@@ -29,6 +29,7 @@ public:
 
 	size_t m_memoryUsed = 0;
 	size_t m_size;
+	char* m_buffer;
 
 	void* alloc(size_t n) {
 		return chHeapAlloc(&m_heap, n);
@@ -41,15 +42,17 @@ public:
 public:
 	template<size_t TSize>
 	Heap(char (&buffer)[TSize])
-		: m_size(TSize)
 	{
-		chHeapObjectInit(&m_heap, buffer, TSize);
+		reinit(buffer, TSize);
 	}
 
-	void reinit(char *buffer, size_t m_size) {
+	void reinit(char *buffer, size_t size) {
 		efiAssertVoid(OBD_PCM_Processor_Fault, m_memoryUsed == 0, "Too late to reinit Lua heap");
-		chHeapObjectInit(&m_heap, buffer, m_size);
-		this->m_size = m_size;
+
+		m_size = size;
+		m_buffer = buffer;
+
+		reset();
 	}
 
 	void* realloc(void* ptr, size_t osize, size_t nsize) {
@@ -88,10 +91,23 @@ public:
 	size_t used() const {
 		return m_memoryUsed;
 	}
+
+	// Use only in case of emergency - obliterates all heap objects and starts over
+	void reset() {
+		chHeapObjectInit(&m_heap, m_buffer, m_size);
+		m_memoryUsed = 0;
+	}
 };
 
 static Heap userHeap(luaUserHeap);
 
+static void printLuaMemoryInfo() {
+	auto heapSize = userHeap.size();
+	auto memoryUsed = userHeap.used();
+	float pct = 100.0f * memoryUsed / heapSize;
+	efiPrintf("Lua memory heap usage: %d / %d bytes = %.1f%%", memoryUsed, heapSize, pct);
+}
+
 static void* myAlloc(void* /*ud*/, void* ptr, size_t osize, size_t nsize) {
 	if (engineConfiguration->debugMode == DBG_LUA) {
 		engine->outputChannels.debugIntField1 = userHeap.used();
@@ -195,7 +211,8 @@ static bool loadScript(LuaHandle& ls, const char* scriptStr) {
 	}
 
 	efiPrintf(TAG "script loaded successfully!");
-	printLuaMemory();
+
+	printLuaMemoryInfo();
 
 	return true;
 }
@@ -329,6 +346,15 @@ void LuaThread::ThreadTask() {
 	while (!chThdShouldTerminateX()) {
 		bool wasOk = runOneLua(myAlloc, config->luaScript);
 
+		auto usedAfterRun = userHeap.used();
+		if (usedAfterRun != 0) {
+			efiPrintf(TAG "MEMORY LEAK DETECTED: %d bytes used after teardown", usedAfterRun);
+
+			// Lua blew up in some terrible way that left memory allocated, reset the heap
+			// so that subsequent runs don't overflow the heap
+			userHeap.reset();
+		}
+
 		// Reset any lua adjustments the script made
 		engine->resetLua();
 
@@ -380,8 +406,8 @@ void startLua() {
 		needsReset = true;
 	});
 
-	addConsoleAction("luamemory", printLuaMemory);
-#endif // LUA_USER_HEAP
+	addConsoleAction("luamemory", printLuaMemoryInfo);
+#endif
 }
 
 #else // not EFI_UNIT_TEST