Hi!

On 28.12.2024 04:48, Tomas Vondra wrote:
On 12/27/24 20:14, James Hunter wrote:
Reviving this thread, because I am thinking about something related --
please ignore the "On Fri, Dec 27, 2024" date, this seems to be an
artifact of me re-sending the message, from the list archive. The
original message was from January 28, 2024.

On Fri, Dec 27, 2024 at 11:02 AM Tomas Vondra
<tomas.von...@enterprisedb.com> wrote:

Firstly, I agree with the goal of having a way to account for memory
used by the backends, and also ability to enforce some sort of limit.
It's difficult to track the memory at the OS level (interpreting RSS
values is not trivial), and work_mem is not sufficient to enforce a
backend-level limit, not even talking about a global limit.

But as I said earlier, it seems quite strange to start by introducing
some sort of global limit, combining memory for all backends. I do
understand that the intent is to have such global limit in order to
prevent issues with the OOM killer and/or not to interfere with other
stuff running on the same machine. And while I'm not saying we should
not have such limit, every time I wished to have a memory limit it was a
backend-level one. Ideally a workmem-like limit that would "adjust" the
work_mem values used by the optimizer (but that's not what this patch
aims to do), or at least a backstop in case something goes wrong (say, a
memory leak, OLTP application issuing complex queries, etc.).

I think what Tomas suggests is the right strategy.

I'm also interested in this topic. And agreed that it's best to move from the 
limit
for a separate backend to the global one. In more details let me suggest
the following steps or parts:
1) realize memory limitation for a separate backend independent from the 
work_mem GUC;
2) add workmem-like limit that would "adjust" the work_mem values used by
the optimize as Thomas suggested;
3) add global limit for all backends.

As for p.1 there is a patch that was originally suggested by my colleague
Maxim Orlov <orlo...@gmail.com> and which i modified for the current master.
This patch introduces the only max_backend_memory GUC that specifies
the maximum amount of memory that can be allocated to a backend.
Zero value means no limit.
If the allocated memory size is exceeded, a standard "out of memory" error will 
be issued.
Also the patch introduces the pg_get_backend_memory_contexts_total_bytes() 
function,
which allows to know how many bytes have already been allocated
to the process in contexts. And in the build with asserts it adds
the pg_get_backend_memory_allocation_stats() function that allows
to get additional information about memory allocations for debug purposes.

This strategy solves the ongoing problem of how to set work_mem, if
some queries have lots of operators and others don't -- now we just
set backend_work_mem, as a limit on the entire query's total work_mem.
And a bit of integration with the optimizer will allow us to
distribute the total backend_work_mem to individual execution nodes,
with the goal of minimizing spilling, without exceeding the
backend_work_mem limit.

As for p.2 maybe one can set a maximum number of parallel sort or
hash table operations before writing to disk instead of absolute
value in the work_mem GUC? E.g. introduce а max_query_operations GUC
or a variable in such a way that old work_mem will be equal
to max_backend_memory divided by max_query_operations.

What do you think about such approach?


With the best wishes,

--
Anton A. Melnikov
Postgres Professional: http://www.postgrespro.com
The Russian Postgres Company
From 8ec3548604bf6a1f5f4c290495843e7080d09660 Mon Sep 17 00:00:00 2001
From: "Anton A. Melnikov" <a.melni...@postgrespro.ru>
Date: Sat, 28 Dec 2024 12:47:24 +0300
Subject: [PATCH]     Limit backend heap memory allocation

    This per-backend GUC can limit the amount of heap allocated
    memory.  Upon reaching this limit "out of memory" error will
    be triggered.

    GUC is diabled by default.

    Authors: Maxim Orlov <orlo...@gmail.com>, Anton A. Melnikov <a.melni...@postgrespro.ru>
---
 doc/src/sgml/config.sgml                      |  17 +
 src/backend/utils/activity/backend_status.c   | 335 ++++++++++++++++++
 src/backend/utils/misc/guc_tables.c           |  10 +
 src/backend/utils/misc/postgresql.conf.sample |   1 +
 src/backend/utils/mmgr/aset.c                 |  36 +-
 src/backend/utils/mmgr/generation.c           |  18 +-
 src/backend/utils/mmgr/slab.c                 |  17 +-
 src/include/catalog/pg_proc.dat               |  14 +
 src/include/utils/backend_status.h            |  10 +
 src/test/modules/Makefile                     |   4 +
 .../modules/test_backend_memory/.gitignore    |   4 +
 src/test/modules/test_backend_memory/Makefile |  14 +
 src/test/modules/test_backend_memory/README   |   1 +
 .../t/001_max_backend_memory.pl               |  85 +++++
 14 files changed, 540 insertions(+), 26 deletions(-)
 create mode 100644 src/test/modules/test_backend_memory/.gitignore
 create mode 100644 src/test/modules/test_backend_memory/Makefile
 create mode 100644 src/test/modules/test_backend_memory/README
 create mode 100644 src/test/modules/test_backend_memory/t/001_max_backend_memory.pl

diff --git a/doc/src/sgml/config.sgml b/doc/src/sgml/config.sgml
index fbdd6ce574..4099d6e1a8 100644
--- a/doc/src/sgml/config.sgml
+++ b/doc/src/sgml/config.sgml
@@ -2337,6 +2337,23 @@ include_dir 'conf.d'
       </listitem>
      </varlistentry>
 
+    <varlistentry id="guc-max-backend-memory" xreflabel="max_backend_memory">
+      <term><varname>max_backend_memory</varname> (<type>integer</type>)
+      <indexterm>
+       <primary><varname>max_backend_memory</varname> configuration parameter</primary>
+      </indexterm>
+      </term>
+      <listitem>
+       <para>
+        Specifies the maximum amount of memory that can be allocated to a backend.
+        Zero means no limit.
+       </para>
+       <para>
+        Default: 0
+       </para>
+      </listitem>
+     </varlistentry>
+
      </variablelist>
      </sect2>
 
diff --git a/src/backend/utils/activity/backend_status.c b/src/backend/utils/activity/backend_status.c
index bf33e33a4e..a41f004cf7 100644
--- a/src/backend/utils/activity/backend_status.c
+++ b/src/backend/utils/activity/backend_status.c
@@ -12,6 +12,7 @@
 #include "postgres.h"
 
 #include "access/xact.h"
+#include "funcapi.h"
 #include "libpq/libpq-be.h"
 #include "miscadmin.h"
 #include "pg_trace.h"
@@ -20,6 +21,7 @@
 #include "storage/proc.h"		/* for MyProc */
 #include "storage/procarray.h"
 #include "utils/ascii.h"
+#include "utils/builtins.h"
 #include "utils/guc.h"			/* for application_name */
 #include "utils/memutils.h"
 
@@ -73,6 +75,17 @@ static void pgstat_beshutdown_hook(int code, Datum arg);
 static void pgstat_read_current_status(void);
 static void pgstat_setup_backend_status_context(void);
 
+/*
+ * Total memory size allocated by backend.
+ */
+static Size my_allocated_bytes = 0;
+
+/*
+ * GUC variable
+ *
+ * Max backend memory allocation allowed (MB). 0 = disabled
+ */
+int max_backend_memory_size_mb = 0;
 
 /*
  * Report shared-memory space needed by BackendStatusShmemInit.
@@ -1222,3 +1235,325 @@ pgstat_clip_activity(const char *raw_activity)
 
 	return activity;
 }
+
+#ifdef USE_ASSERT_CHECKING
+/*
+ * Single memory allocation.
+ */
+typedef struct
+{
+	void   *addr;
+	Size	size;
+	bool	deleted;
+} HeapBlock;
+
+/*
+ * Dynamic array of all allocated memory.
+ */
+typedef struct
+{
+	HeapBlock	   *blocks;
+	Size			size;
+	Size			capacity;
+	Size			ndeleted;
+} HeapBlockStat;
+
+static HeapBlockStat heapstat =
+{
+	.blocks = NULL,
+	.size = 0,
+	.capacity = 0,
+	.ndeleted = 0
+};
+
+/*
+ * Extend heapstat allocation stat array.
+ */
+static inline void
+HeapBlockStatExtend(void)
+{
+	if (heapstat.capacity == 0)
+	{
+#define INIT_MEMORY_CAPACITY 1024
+		heapstat.capacity = INIT_MEMORY_CAPACITY;
+		heapstat.blocks = (HeapBlock *) malloc(heapstat.capacity *
+											sizeof(heapstat.blocks[0]));
+	}
+	else if (heapstat.capacity == heapstat.size)
+	{
+		heapstat.capacity *= 2;
+		heapstat.blocks = (HeapBlock *) realloc(heapstat.blocks,
+											 heapstat.capacity *
+											 sizeof(heapstat.blocks[0]));
+	}
+}
+
+/*
+ * Compare HeapBlocks. Put not deleted in the beginning of an array.
+ */
+static int
+HeapBlockCmp(const void *p, const void *q)
+{
+	const HeapBlock    *a = (const HeapBlock *) p,
+					   *b = (const HeapBlock *) q;
+	int					arg1 = a->deleted,
+						arg2 = b->deleted;
+
+	if (arg1 < arg2)
+		return -1;
+	if (arg1 > arg2)
+		return 1;
+	return 0;
+}
+
+/*
+ * Put all not deleted items in the beginning.
+ */
+static void
+HeapBlockStatCompactify(void)
+{
+	if (heapstat.ndeleted < heapstat.size / 2)
+		return;
+
+	qsort(heapstat.blocks, heapstat.size, sizeof(heapstat.blocks[0]), HeapBlockCmp);
+	heapstat.size -= heapstat.ndeleted;
+	heapstat.ndeleted = 0;
+}
+
+/*
+ * Push new memory allocation.
+ */
+static void
+HeapBlockStatPush(void *ptr, Size size)
+{
+	Size	i;
+
+	HeapBlockStatExtend();
+	HeapBlockStatCompactify();
+
+	for (i = 0; i < heapstat.size; ++i)
+		Assert(heapstat.blocks[i].deleted || heapstat.blocks[i].addr != ptr);
+
+	Assert(heapstat.ndeleted <= heapstat.size);
+
+	/* Try to find deleted item to reuse it... */
+	for (i = 0; i < heapstat.size; ++i)
+	{
+		if (heapstat.blocks[i].deleted == true)
+		{
+			heapstat.blocks[i].addr = ptr;
+			heapstat.blocks[i].size = size;
+			heapstat.blocks[i].deleted = false;
+			--heapstat.ndeleted;
+			Assert(heapstat.ndeleted <= heapstat.size);
+			return;
+		}
+	}
+
+	/* ... no empty places, append! */
+	heapstat.blocks[i].addr = ptr;
+	heapstat.blocks[i].size = size;
+	heapstat.blocks[i].deleted = false;
+
+	++heapstat.size;
+}
+
+/*
+ * Pop memory allocation.
+ */
+static Size
+HeapBlockStatPop(void *ptr)
+{
+	Size	i;
+
+	for (i = 0; i < heapstat.size; ++i)
+	{
+		if (heapstat.blocks[i].addr == ptr)
+		{
+			Assert(heapstat.blocks[i].deleted == false);
+			heapstat.blocks[i].deleted = true;
+			++heapstat.ndeleted;
+			Assert(heapstat.ndeleted <= heapstat.size);
+			return heapstat.blocks[i].size;
+		}
+	}
+
+	/* should not get here... */
+	Assert(i != heapstat.size);
+	return 0;
+}
+#endif /* USE_ASSERT_CHECKING */
+
+/*
+ * Count backend allocated memory.
+ * If limit is set (max_backend_memory_size_mb) and check=true, then need to
+ * check my_allocated_bytes for going beyond the limit.
+ */
+static bool
+pgstat_alloc(Size size, bool check)
+{
+	/* Exclude auxiliary processes from the check */
+	switch (MyBackendType)
+	{
+		case B_STARTUP:
+		case B_ARCHIVER:
+		case B_BG_WRITER:
+		case B_CHECKPOINTER:
+		case B_WAL_WRITER:
+		case B_WAL_RECEIVER:
+		case B_WAL_SUMMARIZER:
+			return true;
+		default:
+			break;
+	}
+
+	my_allocated_bytes += size;
+
+	if (max_backend_memory_size_mb == 0 || !check)
+		return true;
+
+	/* Check for going beyond the limit */
+#define TO_BYTES(mb)	((Size)(mb) * 1024 * 1024)
+	if (my_allocated_bytes > TO_BYTES(max_backend_memory_size_mb))
+	{
+		my_allocated_bytes -= size;
+		return false;
+	}
+
+	return true;
+}
+
+/*
+ * Count deallocated backend memory.
+ */
+static void
+pgstat_free(Size size)
+{
+	/* Exclude auxiliary processes from the check */
+	switch (MyBackendType)
+	{
+		case B_STARTUP:
+		case B_ARCHIVER:
+		case B_BG_WRITER:
+		case B_CHECKPOINTER:
+		case B_WAL_WRITER:
+		case B_WAL_RECEIVER:
+		case B_WAL_SUMMARIZER:
+			return;
+		default:
+			break;
+	}
+
+	Assert(my_allocated_bytes >= size);
+	my_allocated_bytes -= size;
+}
+
+/*
+ * malloc() with bytes counter
+ */
+void *
+malloc_and_count(Size size)
+{
+	void   *ptr;
+
+	if (!pgstat_alloc(size, true))
+		return NULL;
+
+	ptr = malloc(size);
+	if (ptr == NULL)
+		pgstat_free(size);
+#ifdef USE_ASSERT_CHECKING
+	else
+		HeapBlockStatPush(ptr, size);
+#endif
+
+	return ptr;
+}
+
+/*
+ * realloc() with bytes counter
+ */
+void *
+realloc_and_count(void *ptr, Size new_size, Size old_size)
+{
+	Assert(old_size == HeapBlockStatPop(ptr));
+
+	pgstat_free(old_size);
+
+	if (!pgstat_alloc(new_size, true))
+	{
+		/* New buffer can not be alloced, so need to restore old one */
+		pgstat_alloc(old_size, false);
+#ifdef USE_ASSERT_CHECKING
+		HeapBlockStatPush(ptr, old_size);
+#endif
+		return NULL;
+	}
+
+	ptr = realloc(ptr, new_size);
+	if (ptr == NULL)
+		pgstat_free(new_size);
+#ifdef USE_ASSERT_CHECKING
+	else
+		HeapBlockStatPush(ptr, new_size);
+#endif
+
+	return ptr;
+}
+
+/*
+ * free() with bytes counter
+ */
+void
+free_and_count(void *ptr, Size size)
+{
+#ifdef USE_ASSERT_CHECKING
+	Assert(size == HeapBlockStatPop(ptr));
+#endif
+
+	pgstat_free(size);
+	free(ptr);
+}
+
+/*
+ * pg_get_backend_memory_contexts_total_bytes
+ *		Total amount of bytes in all allocated contexts.
+ */
+Datum
+pg_get_backend_memory_contexts_total_bytes(PG_FUNCTION_ARGS)
+{
+	PG_RETURN_INT64(my_allocated_bytes);
+}
+
+/*
+ * pg_get_backend_memory_allocation_stats
+ *		SQL SRF showing backend allocated memory.
+ */
+Datum
+pg_get_backend_memory_allocation_stats(PG_FUNCTION_ARGS)
+{
+#ifdef USE_ASSERT_CHECKING
+#define MEMORY_ALLOCATION_COLS 3
+	ReturnSetInfo  *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	Datum			values[MEMORY_ALLOCATION_COLS];
+	bool			nulls[MEMORY_ALLOCATION_COLS];
+	Size			i;
+	char			text[16];
+
+	InitMaterializedSRF(fcinfo, 0);
+
+	memset(nulls, 0, sizeof(nulls));
+	for (i = 0; i < heapstat.size; ++i)
+	{
+		pg_snprintf(text, sizeof(text), "%p", heapstat.blocks[i].addr);
+		values[0] = CStringGetTextDatum(text);
+		values[1] = Int64GetDatum(heapstat.blocks[i].size);
+		values[2] = BoolGetDatum(heapstat.blocks[i].deleted);
+		tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc, values, nulls);
+	}
+#else
+	elog(ERROR, "this only works for builds with asserts enabled");
+#endif
+
+	return (Datum) 0;
+}
diff --git a/src/backend/utils/misc/guc_tables.c b/src/backend/utils/misc/guc_tables.c
index 8cf1afbad2..f888d417de 100644
--- a/src/backend/utils/misc/guc_tables.c
+++ b/src/backend/utils/misc/guc_tables.c
@@ -3732,6 +3732,16 @@ struct config_int ConfigureNamesInt[] =
 		SCRAM_SHA_256_DEFAULT_ITERATIONS, 1, INT_MAX,
 		NULL, NULL, NULL
 	},
+	{
+		{"max_backend_memory", PGC_SU_BACKEND, RESOURCES_MEM,
+			gettext_noop("Restrict total backend memory allocations to this max (rounded up to the nearest MB)."),
+			NULL,
+			GUC_UNIT_MB
+		},
+		&max_backend_memory_size_mb,
+		0, 0, INT_MAX,
+		NULL, NULL, NULL
+	},
 
 	/* End-of-list marker */
 	{
diff --git a/src/backend/utils/misc/postgresql.conf.sample b/src/backend/utils/misc/postgresql.conf.sample
index a2ac7575ca..1854601ca9 100644
--- a/src/backend/utils/misc/postgresql.conf.sample
+++ b/src/backend/utils/misc/postgresql.conf.sample
@@ -171,6 +171,7 @@
 #serializable_buffers = 32		# memory for pg_serial
 #subtransaction_buffers = 0		# memory for pg_subtrans (0 = auto)
 #transaction_buffers = 0		# memory for pg_xact (0 = auto)
+#max_backend_memory = 0			# limit amount of per-backend allocated memory in megabytes
 
 # - Disk -
 
diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index dede30dd86..e304da7c3c 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -47,6 +47,7 @@
 #include "postgres.h"
 
 #include "port/pg_bitutils.h"
+#include "utils/backend_status.h"
 #include "utils/memdebug.h"
 #include "utils/memutils.h"
 #include "utils/memutils_internal.h"
@@ -441,7 +442,7 @@ AllocSetContextCreateInternal(MemoryContext parent,
 	 * Allocate the initial block.  Unlike other aset.c blocks, it starts with
 	 * the context header and its block header follows that.
 	 */
-	set = (AllocSet) malloc(firstBlockSize);
+	set = (AllocSet) malloc_and_count(firstBlockSize);
 	if (set == NULL)
 	{
 		if (TopMemoryContext)
@@ -579,13 +580,16 @@ AllocSetReset(MemoryContext context)
 		}
 		else
 		{
+			Size	deallocation_size;
+
 			/* Normal case, release the block */
-			context->mem_allocated -= block->endptr - ((char *) block);
+			deallocation_size = block->endptr - ((char *) block);
+			context->mem_allocated -= deallocation_size;
 
 #ifdef CLOBBER_FREED_MEMORY
 			wipe_mem(block, block->freeptr - ((char *) block));
 #endif
-			free(block);
+			free_and_count(block, deallocation_size);
 		}
 		block = next;
 	}
@@ -649,7 +653,7 @@ AllocSetDelete(MemoryContext context)
 				freelist->num_free--;
 
 				/* All that remains is to free the header/initial block */
-				free(oldset);
+				free_and_count(oldset, oldset->header.mem_allocated);
 			}
 			Assert(freelist->num_free == 0);
 		}
@@ -666,16 +670,20 @@ AllocSetDelete(MemoryContext context)
 	while (block != NULL)
 	{
 		AllocBlock	next = block->next;
+		Size		deallocation_size = 0;
 
 		if (!IsKeeperBlock(set, block))
-			context->mem_allocated -= block->endptr - ((char *) block);
+		{
+			deallocation_size = block->endptr - ((char *) block);
+			context->mem_allocated -= deallocation_size;
+		}
 
 #ifdef CLOBBER_FREED_MEMORY
 		wipe_mem(block, block->freeptr - ((char *) block));
 #endif
 
 		if (!IsKeeperBlock(set, block))
-			free(block);
+			free_and_count(block, deallocation_size);
 
 		block = next;
 	}
@@ -683,7 +691,7 @@ AllocSetDelete(MemoryContext context)
 	Assert(context->mem_allocated == keepersize);
 
 	/* Finally, free the context header, including the keeper block */
-	free(set);
+	free_and_count(set, keepersize);
 }
 
 /*
@@ -712,7 +720,7 @@ AllocSetAllocLarge(MemoryContext context, Size size, int flags)
 #endif
 
 	blksize = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
-	block = (AllocBlock) malloc(blksize);
+	block = (AllocBlock) malloc_and_count(blksize);
 	if (block == NULL)
 		return MemoryContextAllocationFailure(context, size, flags);
 
@@ -905,7 +913,7 @@ AllocSetAllocFromNewBlock(MemoryContext context, Size size, int flags,
 		blksize <<= 1;
 
 	/* Try to allocate it */
-	block = (AllocBlock) malloc(blksize);
+	block = (AllocBlock) malloc_and_count(blksize);
 
 	/*
 	 * We could be asking for pretty big blocks here, so cope if malloc fails.
@@ -916,7 +924,7 @@ AllocSetAllocFromNewBlock(MemoryContext context, Size size, int flags,
 		blksize >>= 1;
 		if (blksize < required_size)
 			break;
-		block = (AllocBlock) malloc(blksize);
+		block = (AllocBlock) malloc_and_count(blksize);
 	}
 
 	if (block == NULL)
@@ -1071,6 +1079,7 @@ AllocSetFree(void *pointer)
 	{
 		/* Release single-chunk block. */
 		AllocBlock	block = ExternalChunkGetBlock(chunk);
+		Size		deallocation_size;
 
 		/*
 		 * Try to verify that we have a sane block pointer: the block header
@@ -1099,12 +1108,13 @@ AllocSetFree(void *pointer)
 		if (block->next)
 			block->next->prev = block->prev;
 
-		set->header.mem_allocated -= block->endptr - ((char *) block);
+		deallocation_size = block->endptr - ((char *) block);
+		set->header.mem_allocated -= deallocation_size;
 
 #ifdef CLOBBER_FREED_MEMORY
 		wipe_mem(block, block->freeptr - ((char *) block));
 #endif
-		free(block);
+		free_and_count(block, deallocation_size);
 	}
 	else
 	{
@@ -1223,7 +1233,7 @@ AllocSetRealloc(void *pointer, Size size, int flags)
 		blksize = chksize + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
 		oldblksize = block->endptr - ((char *) block);
 
-		block = (AllocBlock) realloc(block, blksize);
+		block = (AllocBlock) realloc_and_count(block, blksize, oldblksize);
 		if (block == NULL)
 		{
 			/* Disallow access to the chunk header. */
diff --git a/src/backend/utils/mmgr/generation.c b/src/backend/utils/mmgr/generation.c
index 0238c111d2..c7ec138110 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -37,6 +37,7 @@
 
 #include "lib/ilist.h"
 #include "port/pg_bitutils.h"
+#include "utils/backend_status.h"
 #include "utils/memdebug.h"
 #include "utils/memutils.h"
 #include "utils/memutils_internal.h"
@@ -206,7 +207,7 @@ GenerationContextCreate(MemoryContext parent,
 	 * Allocate the initial block.  Unlike other generation.c blocks, it
 	 * starts with the context header and its block header follows that.
 	 */
-	set = (GenerationContext *) malloc(allocSize);
+	set = (GenerationContext *) malloc_and_count(allocSize);
 	if (set == NULL)
 	{
 		MemoryContextStats(TopMemoryContext);
@@ -330,7 +331,7 @@ GenerationDelete(MemoryContext context)
 	/* Reset to release all releasable GenerationBlocks */
 	GenerationReset(context);
 	/* And free the context header and keeper block */
-	free(context);
+	free_and_count(context, context->mem_allocated + MAXALIGN(sizeof(GenerationContext)));
 }
 
 /*
@@ -361,7 +362,7 @@ GenerationAllocLarge(MemoryContext context, Size size, int flags)
 	required_size = chunk_size + Generation_CHUNKHDRSZ;
 	blksize = required_size + Generation_BLOCKHDRSZ;
 
-	block = (GenerationBlock *) malloc(blksize);
+	block = (GenerationBlock *) malloc_and_count(blksize);
 	if (block == NULL)
 		return MemoryContextAllocationFailure(context, size, flags);
 
@@ -482,7 +483,7 @@ GenerationAllocFromNewBlock(MemoryContext context, Size size, int flags,
 	if (blksize < required_size)
 		blksize = pg_nextpower2_size_t(required_size);
 
-	block = (GenerationBlock *) malloc(blksize);
+	block = (GenerationBlock *) malloc_and_count(blksize);
 
 	if (block == NULL)
 		return MemoryContextAllocationFailure(context, size, flags);
@@ -663,6 +664,9 @@ GenerationBlockFreeBytes(GenerationBlock *block)
 static inline void
 GenerationBlockFree(GenerationContext *set, GenerationBlock *block)
 {
+	/* Have to store the value locally, block will be wiped. */
+	Size	blksize = block->blksize;
+
 	/* Make sure nobody tries to free the keeper block */
 	Assert(!IsKeeperBlock(set, block));
 	/* We shouldn't be freeing the freeblock either */
@@ -671,13 +675,13 @@ GenerationBlockFree(GenerationContext *set, GenerationBlock *block)
 	/* release the block from the list of blocks */
 	dlist_delete(&block->node);
 
-	((MemoryContext) set)->mem_allocated -= block->blksize;
+	((MemoryContext) set)->mem_allocated -= blksize;
 
 #ifdef CLOBBER_FREED_MEMORY
-	wipe_mem(block, block->blksize);
+	wipe_mem(block, blksize);
 #endif
 
-	free(block);
+	free_and_count(block, blksize);
 }
 
 /*
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index 3e15d59683..b8baea6c98 100644
--- a/src/backend/utils/mmgr/slab.c
+++ b/src/backend/utils/mmgr/slab.c
@@ -69,6 +69,7 @@
 #include "postgres.h"
 
 #include "lib/ilist.h"
+#include "utils/backend_status.h"
 #include "utils/memdebug.h"
 #include "utils/memutils.h"
 #include "utils/memutils_internal.h"
@@ -361,7 +362,7 @@ SlabContextCreate(MemoryContext parent,
 
 
 
-	slab = (SlabContext *) malloc(Slab_CONTEXT_HDRSZ(chunksPerBlock));
+	slab = (SlabContext *) malloc_and_count(Slab_CONTEXT_HDRSZ(chunksPerBlock));
 	if (slab == NULL)
 	{
 		MemoryContextStats(TopMemoryContext);
@@ -451,7 +452,7 @@ SlabReset(MemoryContext context)
 #ifdef CLOBBER_FREED_MEMORY
 		wipe_mem(block, slab->blockSize);
 #endif
-		free(block);
+		free_and_count(block, slab->blockSize);
 		context->mem_allocated -= slab->blockSize;
 	}
 
@@ -467,7 +468,7 @@ SlabReset(MemoryContext context)
 #ifdef CLOBBER_FREED_MEMORY
 			wipe_mem(block, slab->blockSize);
 #endif
-			free(block);
+			free_and_count(block, slab->blockSize);
 			context->mem_allocated -= slab->blockSize;
 		}
 	}
@@ -484,10 +485,14 @@ SlabReset(MemoryContext context)
 void
 SlabDelete(MemoryContext context)
 {
+#ifdef MEMORY_CONTEXT_CHECKING
+	SlabContext *slab = (SlabContext *) context;
+#endif
+
 	/* Reset to release all the SlabBlocks */
 	SlabReset(context);
 	/* And free the context header */
-	free(context);
+	free_and_count(context, Slab_CONTEXT_HDRSZ(slab->chunksPerBlock));
 }
 
 /*
@@ -562,7 +567,7 @@ SlabAllocFromNewBlock(MemoryContext context, Size size, int flags)
 	}
 	else
 	{
-		block = (SlabBlock *) malloc(slab->blockSize);
+		block = (SlabBlock *) malloc_and_count(slab->blockSize);
 
 		if (unlikely(block == NULL))
 			return MemoryContextAllocationFailure(context, size, flags);
@@ -795,7 +800,7 @@ SlabFree(void *pointer)
 #ifdef CLOBBER_FREED_MEMORY
 			wipe_mem(block, slab->blockSize);
 #endif
-			free(block);
+			free_and_count(block, slab->blockSize);
 			slab->header.mem_allocated -= slab->blockSize;
 		}
 
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 2dcc2d42da..16dd4a0d97 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -8442,6 +8442,20 @@
   proargmodes => '{o,o,o,o,o,o,o,o,o,o}',
   proargnames => '{name, ident, type, level, path, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes}',
   prosrc => 'pg_get_backend_memory_contexts' },
+{ oid => '8085',
+  descr => 'get the allocated memory in all memory contexts of local backend',
+  proname => 'pg_get_backend_memory_contexts_total_bytes', provolatile => 'v',
+  proparallel => 'r', prorettype => 'int8', proargtypes => '',
+  prosrc => 'pg_get_backend_memory_contexts_total_bytes' },
+{ oid => '8005',
+  descr => 'information about memory allocation of local backend',
+  proname => 'pg_get_backend_memory_allocation_stats', prorows => '1024',
+  proretset => 't', provolatile => 'v', proparallel => 'r',
+  prorettype => 'record', proargtypes => '',
+  proallargtypes => '{text,int8,bool}',
+  proargmodes => '{o,o,o}',
+  proargnames => '{pointer, size, deleted}',
+  prosrc => 'pg_get_backend_memory_allocation_stats' },
 
 # logging memory contexts of the specified backend
 { oid => '4543', descr => 'log memory contexts of the specified backend',
diff --git a/src/include/utils/backend_status.h b/src/include/utils/backend_status.h
index 4e8b39a66d..26b9338b2f 100644
--- a/src/include/utils/backend_status.h
+++ b/src/include/utils/backend_status.h
@@ -11,6 +11,7 @@
 #define BACKEND_STATUS_H
 
 #include "datatype/timestamp.h"
+#include "fmgr.h"
 #include "libpq/pqcomm.h"
 #include "miscadmin.h"			/* for BackendType */
 #include "storage/procnumber.h"
@@ -285,6 +286,7 @@ typedef struct LocalPgBackendStatus
  */
 extern PGDLLIMPORT bool pgstat_track_activities;
 extern PGDLLIMPORT int pgstat_track_activity_query_size;
+extern PGDLLIMPORT int max_backend_memory_size_mb;
 
 
 /* ----------
@@ -337,5 +339,13 @@ extern LocalPgBackendStatus *pgstat_get_local_beentry_by_proc_number(ProcNumber
 extern LocalPgBackendStatus *pgstat_get_local_beentry_by_index(int idx);
 extern char *pgstat_clip_activity(const char *raw_activity);
 
+/*
+ * Fuctions to count memory
+ */
+extern void *malloc_and_count(Size size);
+extern void *realloc_and_count(void *ptr, Size new_size, Size old_size);
+extern void free_and_count(void *ptr, Size size);
+extern Datum pg_get_backend_memory_contexts_total_bytes(PG_FUNCTION_ARGS);
+extern Datum pg_get_backend_memory_allocation_stats(PG_FUNCTION_ARGS);
 
 #endif							/* BACKEND_STATUS_H */
diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index c0d3cf0e14..f011f29b27 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -54,6 +54,10 @@ else
 ALWAYS_SUBDIRS += ssl_passphrase_callback
 endif
 
+ifneq (test_backend_memory,$(filter test_backend_memory,$(PG_TEST_SKIP)))
+SUBDIRS += test_backend_memory
+endif
+
 # Test runs an LDAP server, so only run if ldap is in PG_TEST_EXTRA
 ifeq ($(with_ldap),yes)
 ifneq (,$(filter ldap,$(PG_TEST_EXTRA)))
diff --git a/src/test/modules/test_backend_memory/.gitignore b/src/test/modules/test_backend_memory/.gitignore
new file mode 100644
index 0000000000..5dcb3ff972
--- /dev/null
+++ b/src/test/modules/test_backend_memory/.gitignore
@@ -0,0 +1,4 @@
+# Generated subdirectories
+/log/
+/results/
+/tmp_check/
diff --git a/src/test/modules/test_backend_memory/Makefile b/src/test/modules/test_backend_memory/Makefile
new file mode 100644
index 0000000000..c6afaedd0a
--- /dev/null
+++ b/src/test/modules/test_backend_memory/Makefile
@@ -0,0 +1,14 @@
+# src/test/modules/test_backend_memory/Makefile
+
+TAP_TESTS = 1
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_backend_memory
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_backend_memory/README b/src/test/modules/test_backend_memory/README
new file mode 100644
index 0000000000..95ffe2c438
--- /dev/null
+++ b/src/test/modules/test_backend_memory/README
@@ -0,0 +1 @@
+This directory contain tests for max_backend_memory GUC
diff --git a/src/test/modules/test_backend_memory/t/001_max_backend_memory.pl b/src/test/modules/test_backend_memory/t/001_max_backend_memory.pl
new file mode 100644
index 0000000000..c479f6ec4f
--- /dev/null
+++ b/src/test/modules/test_backend_memory/t/001_max_backend_memory.pl
@@ -0,0 +1,85 @@
+# Tests to check backend memory limit (max_backend_memory GUC)
+
+use strict;
+use warnings;
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+my $node = PostgreSQL::Test::Cluster->new('main');
+$node->init;
+$node->start;
+
+# Initialize backend memory limit
+$node->safe_psql('postgres', qq{
+	ALTER SYSTEM SET max_backend_memory = 15;
+	SELECT pg_reload_conf();
+});
+
+# Check backend memory limit after SET.
+my $psql_stdout;
+$psql_stdout = $node->safe_psql('postgres',
+								"SELECT current_setting('max_backend_memory');");
+is($psql_stdout, '15MB', "max_backend_memory is SET correctly");
+
+# Create test table and function.
+$node->safe_psql('postgres', q{
+	CREATE TABLE test(t text);
+	INSERT INTO test VALUES (repeat('1234567890', 400000));
+
+	-- Recursive function that should cause overflow
+	CREATE FUNCTION test_func() RETURNS void LANGUAGE plpgsql AS $$
+	DECLARE
+		bt text;
+	BEGIN
+		SELECT t || 'x' FROM test INTO bt;
+		PERFORM test_func();
+	END;
+$$;
+});
+
+# Call function "test_func" several times for memory leak check.
+my $psql_stdout_first;
+for (my $i = 0; $i < 4; $i++)
+{
+	# Call function and check that it finishes with 'out of memory' error.
+	my ($ret, $stdout, $stderr) = $node->psql('postgres',
+											  "SELECT test_func();");
+	is($ret, 3, 'recursive function call causes overflow');
+	like($stderr, qr/out of memory/, 'expected out of memory error');
+
+	$psql_stdout = $node->safe_psql('postgres',
+									"SELECT pg_get_backend_memory_contexts_total_bytes();");
+	if ($i eq 0)
+	{
+		# Store first value of backend_memory_contexts_total_bytes.
+		$psql_stdout_first = $psql_stdout;
+		is($psql_stdout_first > 0, 1,
+		   "backend has allocated $psql_stdout_first (greater than 0) bytes");
+		next;
+	}
+	# Check other values of backend_memory_contexts_total_bytes.
+	# They should be the same as first value.
+	is($psql_stdout, $psql_stdout_first, "memory does not leak");
+}
+
+# Drop test table and function.
+$node->safe_psql('postgres', q{
+	DROP FUNCTION test_func();
+	DROP TABLE test;
+});
+
+# Deinitialize backend memory limit.
+$node->safe_psql('postgres', q{
+	ALTER SYSTEM RESET max_backend_memory;
+	SELECT pg_reload_conf();
+});
+
+# Check backend memory limit after RESET.
+$psql_stdout = $node->safe_psql('postgres',
+								"SELECT current_setting('max_backend_memory');");
+is($psql_stdout, '0', "max_backend_memory is RESET correctly");
+
+$node->stop;
+
+done_testing();
-- 
2.47.1

Reply via email to