From 6738fa0bd7c9cf39f3f68cbdd4d2d861cf7fdc03 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Wed, 20 Dec 2023 16:49:16 -0500
Subject: [PATCH] Added HNSW_MEMORY flag to show memory usage - #384 [skip ci]

---
 src/hnsw.h      |  2 ++
 src/hnswbuild.c | 12 ++++++++++++
 2 files changed, 14 insertions(+)

diff --git a/src/hnsw.h b/src/hnsw.h
index 3497b95..862d7b7 100644
--- a/src/hnsw.h
+++ b/src/hnsw.h
@@ -72,8 +72,10 @@
 
 #if PG_VERSION_NUM >= 150000
 #define RandomDouble() pg_prng_double(&pg_global_prng_state)
+#define SeedRandom(seed) pg_prng_seed(&pg_global_prng_state, seed)
 #else
 #define RandomDouble() (((double) random()) / MAX_RANDOM_VALUE)
+#define SeedRandom(seed) srandom(seed)
 #endif
 
 #if PG_VERSION_NUM < 130000
diff --git a/src/hnswbuild.c b/src/hnswbuild.c
index 2acb80f..afe0fb1 100644
--- a/src/hnswbuild.c
+++ b/src/hnswbuild.c
@@ -279,6 +279,14 @@ FreeElements(HnswBuildState * buildstate)
 static void
 FlushPages(HnswBuildState * buildstate)
 {
+#ifdef HNSW_MEMORY
+#if PG_VERSION_NUM >= 130000
+	elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(CurrentMemoryContext, false) / (1024 * 1024));
+#else
+	MemoryContextStats(CurrentMemoryContext);
+#endif
+#endif
+
 	CreateMetaPage(buildstate);
 	CreateElementPages(buildstate);
 	CreateNeighborPages(buildstate);
@@ -902,6 +910,10 @@ static void
 BuildIndex(Relation heap, Relation index, IndexInfo *indexInfo,
 		   HnswBuildState * buildstate, ForkNumber forkNum)
 {
+#ifdef HNSW_MEMORY
+	SeedRandom(42);
+#endif
+
 	InitBuildState(buildstate, heap, index, indexInfo, forkNum);
 
 	if (buildstate->heap != NULL)