From 53cb80ebf74b059d1a473e1afba260dcdeb33f85 Mon Sep 17 00:00:00 2001
From: Cory McWilliams <cory@unprompted.com>
Date: Sat, 18 Feb 2023 19:14:06 +0000
Subject: [PATCH] Replace the sqlite allocator, and use our own tracking for
 stats.  Want to use this to collect callstacks for all allocations.

git-svn-id: https://www.unprompted.com/svn/projects/tildefriends/trunk@4185 ed5197a5-7fde-0310-b194-c3ffbd925b24
---
 src/main.c |  1 +
 src/mem.c  | 62 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 src/mem.h  |  3 +++
 src/task.c |  2 +-
 4 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/src/main.c b/src/main.c
index e94be9db..e75ad98e 100644
--- a/src/main.c
+++ b/src/main.c
@@ -685,6 +685,7 @@ int main(int argc, char* argv[])
 #endif
 	tf_mem_replace_uv_allocator();
 	tf_mem_replace_tls_allocator();
+	tf_mem_replace_sqlite_allocator();
 	uv_setup_args(argc, argv);
 	tf_taskstub_startup();
 
diff --git a/src/mem.c b/src/mem.c
index 753820c9..7968be88 100644
--- a/src/mem.c
+++ b/src/mem.c
@@ -5,6 +5,7 @@
 #include <openssl/crypto.h>
 
 #include <quickjs.h>
+#include <sqlite3.h>
 
 #include <string.h>
 
@@ -12,6 +13,7 @@ static int64_t s_tf_malloc_size;
 static int64_t s_uv_malloc_size;
 static int64_t s_tls_malloc_size;
 static int64_t s_js_malloc_size;
+static int64_t s_sqlite_malloc_size;
 
 static void* _tf_alloc(int64_t* total, size_t size)
 {
@@ -251,3 +253,63 @@ size_t tf_mem_get_js_malloc_size()
 {
 	return s_js_malloc_size;
 }
+
+static void* _tf_sqlite_malloc(int size)
+{
+	return _tf_alloc(&s_sqlite_malloc_size, size);
+}
+
+static void _tf_sqlite_free(void* ptr)
+{
+	_tf_free(&s_sqlite_malloc_size, ptr);
+}
+
+static void* _tf_sqlite_realloc(void* ptr, int size)
+{
+	return _tf_realloc(&s_sqlite_malloc_size, ptr, size);
+}
+
+static int _tf_sqlite_size(void* ptr)
+{
+	void* old_ptr = ptr ? (void*)((intptr_t)ptr - sizeof(size_t)) : NULL;
+	size_t old_size = 0;
+	if (old_ptr)
+	{
+		memcpy(&old_size, old_ptr, sizeof(size_t));
+	}
+	return (int)old_size;
+}
+
+static int _tf_sqlite_roundup(int size)
+{
+	return (size + 7) & ~7;
+}
+
+static int _tf_sqlite_init(void* user_data)
+{
+	return SQLITE_OK;
+}
+
+static void _tf_sqlite_shutdown(void* user_data)
+{
+}
+
+void tf_mem_replace_sqlite_allocator()
+{
+	sqlite3_mem_methods methods =
+	{
+		.xMalloc = _tf_sqlite_malloc,
+		.xFree = _tf_sqlite_free,
+		.xRealloc = _tf_sqlite_realloc,
+		.xSize = _tf_sqlite_size,
+		.xRoundup = _tf_sqlite_roundup,
+		.xInit = _tf_sqlite_init,
+		.xShutdown = _tf_sqlite_shutdown,
+	};
+	sqlite3_config(SQLITE_CONFIG_MALLOC, &methods);
+}
+
+size_t tf_mem_get_sqlite_malloc_size()
+{
+	return s_sqlite_malloc_size;
+}
diff --git a/src/mem.h b/src/mem.h
index df3ff6be..ab1a96c6 100644
--- a/src/mem.h
+++ b/src/mem.h
@@ -10,6 +10,9 @@ size_t tf_mem_get_uv_malloc_size();
 void tf_mem_replace_tls_allocator();
 size_t tf_mem_get_tls_malloc_size();
 
+void tf_mem_replace_sqlite_allocator();
+size_t tf_mem_get_sqlite_malloc_size();
+
 size_t tf_mem_get_tf_malloc_size();
 
 void* tf_malloc(size_t size);
diff --git a/src/task.c b/src/task.c
index c1ecad11..e65b1c42 100644
--- a/src/task.c
+++ b/src/task.c
@@ -724,7 +724,7 @@ static JSValue _tf_task_getStats(JSContext* context, JSValueConst this_val, int
 		JS_SetPropertyStr(context, result, "memory_percent", JS_NewFloat64(context, 100.0f * rss / total_memory));
 	}
 
-	JS_SetPropertyStr(context, result, "sqlite3_memory_percent", JS_NewFloat64(context, 100.0f * sqlite3_memory_used() / total_memory));
+	JS_SetPropertyStr(context, result, "sqlite3_memory_percent", JS_NewFloat64(context, 100.0f * tf_mem_get_sqlite_malloc_size() / total_memory));
 	JS_SetPropertyStr(context, result, "js_malloc_percent", JS_NewFloat64(context, 100.0f * tf_mem_get_js_malloc_size() / total_memory));
 	JS_SetPropertyStr(context, result, "uv_malloc_percent", JS_NewFloat64(context, 100.0f * tf_mem_get_uv_malloc_size() / total_memory));
 	JS_SetPropertyStr(context, result, "tls_malloc_percent", JS_NewFloat64(context, 100.0f * tf_mem_get_tls_malloc_size() / total_memory));