From efd8d4431ea49612cfa8174e60db8ea2e8ea6b7f Mon Sep 17 00:00:00 2001
From: Jayesh D <jayesh.dp@zohocorp.com>
Date: Wed, 13 Nov 2024 20:04:53 +0530
Subject: [PATCH] Custom malloc context implementation

---
 configure                                     |  32 +++
 src/backend/utils/misc/guc_tables.c           |  14 +-
 src/backend/utils/misc/postgresql.conf.sample |   2 +-
 src/backend/utils/mmgr/Makefile               |   1 +
 src/backend/utils/mmgr/aset.c                 |   4 +
 src/backend/utils/mmgr/generation.c           |   4 +
 src/backend/utils/mmgr/malloccxt.c            | 214 ++++++++++++++++++
 src/backend/utils/mmgr/mcxt.c                 |  14 +-
 src/backend/utils/mmgr/meson.build            |   1 +
 src/backend/utils/mmgr/slab.c                 |   4 +
 src/include/pg_config.h.in                    |   3 +
 src/include/utils/memutils.h                  |   4 +
 src/include/utils/memutils_internal.h         |  18 +-
 src/test/regress/expected/sysviews.out        |   5 +-
 src/test/regress/sql/sysviews.sql             |   2 +
 15 files changed, 317 insertions(+), 5 deletions(-)
 create mode 100644 src/backend/utils/mmgr/malloccxt.c

diff --git a/configure b/configure
index 4b01b682b1..2f871ac210 100755
--- a/configure
+++ b/configure
@@ -848,6 +848,7 @@ with_wal_blocksize
 with_llvm
 enable_depend
 enable_cassert
+enable_asan
 with_icu
 with_tcl
 with_tclconfig
@@ -1536,6 +1537,7 @@ Optional Features:
                           enable injection points (for testing)
   --enable-depend         turn on automatic dependency tracking
   --enable-cassert        enable assertion checks (for debugging)
+  --enable-asan           build with AddressSanitizer support
   --disable-largefile     omit support for large files
 
 Optional Packages:
@@ -7936,6 +7938,36 @@ fi
 
 
 
+#
+# Enable AddressSanitizer Support
+#
+
+
+# Check whether --enable-asan was given.
+if test "${enable_asan+set}" = set; then :
+  enableval=$enable_asan;
+  case $enableval in
+    yes)
+
+$as_echo "#define ENABLE_ASAN 1" >>confdefs.h
+CFLAGS="-DRANDOMIZE_ALLOCATED_MEMORY -O0 -fno-omit-frame-pointer -fsanitize=address"
+
+      ;;
+    no)
+      :
+      ;;
+    *)
+      as_fn_error $? "no argument expected for --enable-asan option" "$LINENO" 5
+      ;;
+  esac
+
+else
+  enable_asan=no
+
+fi
+
+
+
 #
 # Include directories
 #
diff --git a/src/backend/utils/misc/guc_tables.c b/src/backend/utils/misc/guc_tables.c
index 8a67f01200..9ef82c3592 100644
--- a/src/backend/utils/misc/guc_tables.c
+++ b/src/backend/utils/misc/guc_tables.c
@@ -499,6 +499,7 @@ bool		Debug_write_read_parse_plan_trees;
 bool		Debug_raw_expression_coverage_test;
 #endif
 
+bool        enable_malloc_context = false;
 bool		log_parser_stats = false;
 bool		log_planner_stats = false;
 bool		log_executor_stats = false;
@@ -2075,7 +2076,18 @@ struct config_bool ConfigureNamesBool[] =
 		false,
 		NULL, NULL, NULL
 	},
-
+	{
+		{"enable_malloc_context", PGC_USERSET, CLIENT_CONN_STATEMENT,
+			gettext_noop("Enable malloc context usage. Enabled by default if compiled with --enable-asan."),
+		},
+		&enable_malloc_context,
+#ifdef ENABLE_ASAN
+		true,
+#else
+		false,
+#endif
+		NULL, NULL, NULL
+	},
 	/* End-of-list marker */
 	{
 		{NULL, 0, 0, NULL, NULL}, NULL, false, NULL, NULL, NULL
diff --git a/src/backend/utils/misc/postgresql.conf.sample b/src/backend/utils/misc/postgresql.conf.sample
index 39a3ac2312..ca7c88a882 100644
--- a/src/backend/utils/misc/postgresql.conf.sample
+++ b/src/backend/utils/misc/postgresql.conf.sample
@@ -126,7 +126,7 @@
 #------------------------------------------------------------------------------
 
 # - Memory -
-
+#enable_malloc_context = off (default off in release build and default on in debug build)
 #shared_buffers = 128MB			# min 128kB
 					# (change requires restart)
 #huge_pages = try			# on, off, or try
diff --git a/src/backend/utils/mmgr/Makefile b/src/backend/utils/mmgr/Makefile
index 01a1fb8527..3240f5774d 100644
--- a/src/backend/utils/mmgr/Makefile
+++ b/src/backend/utils/mmgr/Makefile
@@ -19,6 +19,7 @@ OBJS = \
 	dsa.o \
 	freepage.o \
 	generation.o \
+	malloccxt.o \
 	mcxt.o \
 	memdebug.o \
 	portalmem.o \
diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index dede30dd86..6b594e9350 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -385,6 +385,10 @@ AllocSetContextCreateInternal(MemoryContext parent,
 			minContextSize <= maxBlockSize));
 	Assert(maxBlockSize <= MEMORYCHUNK_MAX_BLOCKOFFSET);
 
+	if(enable_malloc_context) {
+		return MallocContextCreate(parent, name);
+	}
+
 	/*
 	 * Check whether the parameters match either available freelist.  We do
 	 * not need to demand a match of maxBlockSize.
diff --git a/src/backend/utils/mmgr/generation.c b/src/backend/utils/mmgr/generation.c
index 0238c111d2..5fa5b8cebd 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -194,6 +194,10 @@ GenerationContextCreate(MemoryContext parent,
 			minContextSize <= maxBlockSize));
 	Assert(maxBlockSize <= MEMORYCHUNK_MAX_BLOCKOFFSET);
 
+	if(enable_malloc_context) {
+		return MallocContextCreate(parent, name);
+	}
+
 	/* Determine size of initial block */
 	allocSize = MAXALIGN(sizeof(GenerationContext)) +
 		Generation_BLOCKHDRSZ + Generation_CHUNKHDRSZ;
diff --git a/src/backend/utils/mmgr/malloccxt.c b/src/backend/utils/mmgr/malloccxt.c
new file mode 100644
index 0000000000..50d19da892
--- /dev/null
+++ b/src/backend/utils/mmgr/malloccxt.c
@@ -0,0 +1,214 @@
+/* ------------------------------------------------------------------------------
+ *
+ * malloccxt.c
+ * 		Custom malloc context implementation.
+ *		To be used only for debugging purpose.
+ *
+ * ------------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+#include "utils/memutils_memorychunk.h"
+#include "lib/ilist.h"
+
+/* Designed to use malloc/free to enable ADDRESS Sanitizer to catch errors */
+typedef struct MallocContext
+{
+	MemoryContextData header;
+	dlist_head chunks; /* head of list of chunks in this context */
+} MallocContext;
+
+typedef struct AllocChunk
+{
+	dlist_node chunk_list_node; /* Used to chain allocated chunks */
+	Size size;
+	MallocContext *mcxt;
+	MemoryChunk mchunk;
+} AllocChunk;
+
+#define ALLOC_CHUNKHDRSZ MAXALIGN(sizeof(AllocChunk))
+#define AllocPointerGetChunk(ptr)                                                \
+	((AllocChunk *) (((char *) (ptr)) - ALLOC_CHUNKHDRSZ))
+#define AllocChunkGetPointer(chk) ((void *) (((char *) (chk)) + ALLOC_CHUNKHDRSZ))
+
+#ifdef CLOBBER_FREED_MEMORY
+
+/* Wipe freed memory for debugging purposes */
+static void
+malloc_wipe_mem(void *ptr, size_t size)
+{
+	memset(ptr, 0x7F, size);
+}
+#else
+#define malloc_wipe_mem(p, s)
+#endif /* CLOBBER_FREED_MEMORY */
+
+static inline void
+free_chunk(AllocChunk *chunk)
+{
+	dlist_delete(&chunk->chunk_list_node);
+	malloc_wipe_mem(AllocChunkGetPointer(chunk), chunk->size);
+	free(chunk);
+}
+
+void
+MallocContextReset(MemoryContext context)
+{
+	if (!context->isReset)
+	{
+		MallocContext *mcxt = (MallocContext *) (context);
+		dlist_mutable_iter iter;
+
+		dlist_foreach_modify(iter, &mcxt->chunks)
+		{
+			free_chunk(dlist_container(AllocChunk, chunk_list_node, iter.cur));
+		}
+
+		mcxt->header.mem_allocated = 0;
+	}
+}
+
+void
+MallocContextDelete(MemoryContext context)
+{
+	MallocContextReset(context);
+	free(context);
+}
+
+MemoryContext
+MallocContextGetChunkContext(void *pointer)
+{
+	AllocChunk *chunk = AllocPointerGetChunk(pointer);
+	return &chunk->mcxt->header;
+}
+
+void *
+MallocContextAlloc(MemoryContext context, Size size)
+{
+	MallocContext *mcxt;
+	AllocChunk *chunk;
+
+	/* Should be called in single thread */
+	mcxt = (MallocContext *) (context);
+	chunk = (AllocChunk *) (malloc(ALLOC_CHUNKHDRSZ + size));
+
+	if (chunk)
+	{
+		chunk->mcxt = mcxt;
+		chunk->size = size;
+		dlist_push_tail(&mcxt->chunks, &chunk->chunk_list_node);
+		mcxt->header.mem_allocated += size + ALLOC_CHUNKHDRSZ;
+		MemoryChunkSetHdrMaskExternal(&chunk->mchunk, MCTX_MALLOC_ID);
+	}
+
+	return chunk ? AllocChunkGetPointer(chunk) : NULL;
+}
+
+void
+MallocContextFree(void *pointer)
+{
+	MallocContext *mcxt;
+	AllocChunk *chunk;
+
+	chunk = AllocPointerGetChunk(pointer);
+	mcxt = chunk->mcxt;
+	mcxt->header.mem_allocated -= chunk->size + ALLOC_CHUNKHDRSZ;
+
+	free_chunk(chunk);
+}
+
+void *
+MallocContextRealloc(void *pointer, Size size)
+{
+	MallocContext *mcxt;
+	AllocChunk *chunk;
+
+	/* Should be called in single thread */
+	chunk = AllocPointerGetChunk(pointer);
+	mcxt = chunk->mcxt;
+
+	dlist_delete(&chunk->chunk_list_node);
+	chunk = (AllocChunk *) (realloc(chunk, size + ALLOC_CHUNKHDRSZ));
+
+	if (chunk)
+	{
+		mcxt->header.mem_allocated += size - chunk->size;
+		chunk->size = size;
+		dlist_push_tail(&mcxt->chunks, &chunk->chunk_list_node);
+
+		return AllocChunkGetPointer(chunk);
+	}
+	else
+	{
+		chunk = AllocPointerGetChunk(pointer);
+		dlist_push_tail(&mcxt->chunks, &chunk->chunk_list_node);
+		return NULL;
+	}
+}
+
+Size
+MallocContextGetChunkSpace(void *pointer)
+{
+	AllocChunk *chunk = AllocPointerGetChunk(pointer);
+
+	return chunk->size + ALLOC_CHUNKHDRSZ;
+}
+
+bool
+MallocContextIsEmpty(MemoryContext context)
+{
+	return context->isReset;
+}
+
+void
+MallocContextStats(MemoryContext context, MemoryStatsPrintFunc func,
+				   void *passthru, MemoryContextCounters *totals,
+				   bool print_to_stderr)
+{
+	MallocContext *mcxt;
+	Size totalspace;
+	mcxt = (MallocContext *) (context);
+	totalspace = mcxt->header.mem_allocated;
+	if (func)
+	{
+		int level = *(int *) passthru;
+		for (int i = 0; i < level; i++)
+		{
+			fprintf(stderr, "  ");
+		}
+
+		fprintf(stderr, "%s: %zu used\n", mcxt->header.name, totalspace);
+	}
+
+	if (totals)
+	{
+		totals->totalspace += totalspace;
+		totals->freespace = 0;
+	}
+}
+
+#ifdef MEMORY_CONTEXT_CHECKING
+void
+MallocContextCheck(MemoryContext context)
+{
+}
+#endif
+
+
+MemoryContext
+MallocContextCreate(MemoryContext parent, const char *name)
+{
+	MallocContext *mmcxt;
+	/*
+	 * We need to set context node type as AllocSetContext, as there is no way to
+	 * add new `node` info pg system
+	 */
+	mmcxt = (MallocContext *) (malloc(sizeof(MallocContext)));
+
+	mmcxt->header.mem_allocated = 0;
+	dlist_init(&mmcxt->chunks);
+	MemoryContextCreate(&mmcxt->header, T_AllocSetContext, MCTX_MALLOC_ID, parent,
+						name);
+
+	return &mmcxt->header;
+}
\ No newline at end of file
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index bde54326c6..3d4469cde5 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -114,6 +114,19 @@ static const MemoryContextMethods mcxt_methods[] = {
 	[MCTX_BUMP_ID].check = BumpCheck,
 #endif
 
+	/* malloccxt.c */
+	[MCTX_MALLOC_ID].alloc = MallocContextAlloc,
+	[MCTX_MALLOC_ID].free_p = MallocContextFree,
+	[MCTX_MALLOC_ID].realloc = MallocContextRealloc,
+	[MCTX_MALLOC_ID].reset = MallocContextReset,
+	[MCTX_MALLOC_ID].delete_context = MallocContextDelete,
+	[MCTX_MALLOC_ID].get_chunk_context = MallocContextGetChunkContext,
+	[MCTX_MALLOC_ID].get_chunk_space = MallocContextGetChunkSpace,
+	[MCTX_MALLOC_ID].is_empty = MallocContextIsEmpty,
+	[MCTX_MALLOC_ID].stats = MallocContextStats,
+#ifdef MEMORY_CONTEXT_CHECKING
+	[MCTX_MALLOC_ID].check = MallocContextCheck,
+#endif
 
 	/*
 	 * Reserved and unused IDs should have dummy entries here.  This allows us
@@ -123,7 +136,6 @@ static const MemoryContextMethods mcxt_methods[] = {
 	 */
 	BOGUS_MCTX(MCTX_1_RESERVED_GLIBC_ID),
 	BOGUS_MCTX(MCTX_2_RESERVED_GLIBC_ID),
-	BOGUS_MCTX(MCTX_8_UNUSED_ID),
 	BOGUS_MCTX(MCTX_9_UNUSED_ID),
 	BOGUS_MCTX(MCTX_10_UNUSED_ID),
 	BOGUS_MCTX(MCTX_11_UNUSED_ID),
diff --git a/src/backend/utils/mmgr/meson.build b/src/backend/utils/mmgr/meson.build
index dd43a6844c..ea2271cd7d 100644
--- a/src/backend/utils/mmgr/meson.build
+++ b/src/backend/utils/mmgr/meson.build
@@ -7,6 +7,7 @@ backend_sources += files(
   'dsa.c',
   'freepage.c',
   'generation.c',
+  'malloccxt.c',
   'mcxt.c',
   'memdebug.c',
   'portalmem.c',
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index 3e15d59683..abde24d149 100644
--- a/src/backend/utils/mmgr/slab.c
+++ b/src/backend/utils/mmgr/slab.c
@@ -334,6 +334,10 @@ SlabContextCreate(MemoryContext parent,
 					 "sizeof(MemoryChunk) is not maxaligned");
 	Assert(blockSize <= MEMORYCHUNK_MAX_BLOCKOFFSET);
 
+	if(enable_malloc_context) {
+		return MallocContextCreate(parent, name);
+	}
+
 	/*
 	 * Ensure there's enough space to store the pointer to the next free chunk
 	 * in the memory of the (otherwise) unused allocation.
diff --git a/src/include/pg_config.h.in b/src/include/pg_config.h.in
index cdd9a6e935..2fba2bde67 100644
--- a/src/include/pg_config.h.in
+++ b/src/include/pg_config.h.in
@@ -663,6 +663,9 @@
 /* Define to 1 to build with assertion checks. (--enable-cassert) */
 #undef USE_ASSERT_CHECKING
 
+/* Define to 1 to build with AddressSanitizer support. (--enable-asan) */
+#undef ENABLE_ASAN
+
 /* Define to 1 to use AVX-512 popcount instructions with a runtime check. */
 #undef USE_AVX512_POPCNT_WITH_RUNTIME_CHECK
 
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index 3590c8bad9..10631f5ca1 100644
--- a/src/include/utils/memutils.h
+++ b/src/include/utils/memutils.h
@@ -62,6 +62,7 @@ extern PGDLLIMPORT MemoryContext CacheMemoryContext;
 extern PGDLLIMPORT MemoryContext MessageContext;
 extern PGDLLIMPORT MemoryContext TopTransactionContext;
 extern PGDLLIMPORT MemoryContext CurTransactionContext;
+extern PGDLLIMPORT bool enable_malloc_context;
 
 /* This is a transient link to the active portal's memory context: */
 extern PGDLLIMPORT MemoryContext PortalContext;
@@ -150,6 +151,9 @@ extern MemoryContext BumpContextCreate(MemoryContext parent,
 									   Size initBlockSize,
 									   Size maxBlockSize);
 
+/* malloccxt.c */
+extern MemoryContext MallocContextCreate(MemoryContext parent, const char *name);
+
 /*
  * Recommended default alloc parameters, suitable for "ordinary" contexts
  * that might hold quite a lot of data.
diff --git a/src/include/utils/memutils_internal.h b/src/include/utils/memutils_internal.h
index 5ce2d9b587..9a8a008b58 100644
--- a/src/include/utils/memutils_internal.h
+++ b/src/include/utils/memutils_internal.h
@@ -95,6 +95,22 @@ extern void BumpStats(MemoryContext context, MemoryStatsPrintFunc printfunc,
 extern void BumpCheck(MemoryContext context);
 #endif
 
+/* Malloc memory context */
+extern void MallocContextReset(MemoryContext context);
+extern void MallocContextDelete(MemoryContext context);
+extern MemoryContext MallocContextGetChunkContext(void *pointer);
+extern void *MallocContextAlloc(MemoryContext context, Size size);
+extern void MallocContextFree(void *pointer);
+extern void *MallocContextRealloc(void *pointer, Size size);
+extern Size MallocContextGetChunkSpace(void *pointer);
+extern bool MallocContextIsEmpty(MemoryContext context);
+extern void MallocContextStats(MemoryContext context, MemoryStatsPrintFunc func,
+							   void *passthru, MemoryContextCounters *totals,
+							   bool /*print_to_stderr*/);
+#ifdef MEMORY_CONTEXT_CHECKING
+extern void MallocContextCheck(MemoryContext /* context */);
+#endif
+
 /*
  * How many extra bytes do we need to request in order to ensure that we can
  * align a pointer to 'alignto'.  Since palloc'd pointers are already aligned
@@ -128,7 +144,7 @@ typedef enum MemoryContextMethodID
 	MCTX_SLAB_ID,
 	MCTX_ALIGNED_REDIRECT_ID,
 	MCTX_BUMP_ID,
-	MCTX_8_UNUSED_ID,
+	MCTX_MALLOC_ID,   /* Custom malloc context */
 	MCTX_9_UNUSED_ID,
 	MCTX_10_UNUSED_ID,
 	MCTX_11_UNUSED_ID,
diff --git a/src/test/regress/expected/sysviews.out b/src/test/regress/expected/sysviews.out
index fad7fc3a7e..7ee794ea29 100644
--- a/src/test/regress/expected/sysviews.out
+++ b/src/test/regress/expected/sysviews.out
@@ -145,6 +145,7 @@ select count(*) = 0 as ok from pg_stat_wal_receiver;
 
 -- This is to record the prevailing planner enable_foo settings during
 -- a regression test run.
+set enable_malloc_context = off;
 select name, setting from pg_settings where name like 'enable%';
               name              | setting 
 --------------------------------+---------
@@ -157,6 +158,7 @@ select name, setting from pg_settings where name like 'enable%';
  enable_incremental_sort        | on
  enable_indexonlyscan           | on
  enable_indexscan               | on
+ enable_malloc_context          | off
  enable_material                | on
  enable_memoize                 | on
  enable_mergejoin               | on
@@ -170,8 +172,9 @@ select name, setting from pg_settings where name like 'enable%';
  enable_seqscan                 | on
  enable_sort                    | on
  enable_tidscan                 | on
-(22 rows)
+(23 rows)
 
+reset enable_malloc_context;
 -- There are always wait event descriptions for various types.  InjectionPoint
 -- may be present or absent, depending on history since last postmaster start.
 select type, count(*) > 0 as ok FROM pg_wait_events
diff --git a/src/test/regress/sql/sysviews.sql b/src/test/regress/sql/sysviews.sql
index b2a7923754..38b297b69d 100644
--- a/src/test/regress/sql/sysviews.sql
+++ b/src/test/regress/sql/sysviews.sql
@@ -78,7 +78,9 @@ select count(*) = 0 as ok from pg_stat_wal_receiver;
 
 -- This is to record the prevailing planner enable_foo settings during
 -- a regression test run.
+set enable_malloc_context = off;
 select name, setting from pg_settings where name like 'enable%';
+reset enable_malloc_context;
 
 -- There are always wait event descriptions for various types.  InjectionPoint
 -- may be present or absent, depending on history since last postmaster start.
-- 
2.43.1

