On Wed, 2020-03-18 at 15:41 -0700, Jeff Davis wrote:
> In an off-list discussion, Andres suggested that MemoryContextStats
> could be refactored to achieve this purpose, perhaps with flags to
> avoid walking through the blocks and freelists when those are not
> needed.

Attached refactoring patch. There's enough in here that warrants
discussion that I don't think this makes sense for v13 and I'm adding
it to the July commitfest.

I still think we should do something for v13, such as the originally-
proposed patch[1]. It's not critical, but it simply reports a better
number for memory consumption. Currently, the memory usage appears to
jump, often right past work mem (by a reasonable but noticable amount),
which could be confusing.

Regarding the attached patch (target v14):

  * there's a new MemoryContextCount() that simply calculates the
    statistics without printing anything, and returns a struct
    - it supports flags to indicate which stats should be
      calculated, so that some callers can avoid walking through
      blocks/freelists
  * it adds a new statistic for "new space" (i.e. untouched)
  * it eliminates specialization of the memory context printing
    - the only specialization was for generation.c to output the
      number of chunks, which can be done easily enough for the
      other types, too

Regards,
        Jeff Davis


[1] 
https://postgr.es/m/ec63d70b668818255486a83ffadc3aec492c1f57.camel%40j-davis.com
diff --git a/src/backend/executor/nodeAgg.c b/src/backend/executor/nodeAgg.c
index 2a6f44a6274..9ae62dda9cd 100644
--- a/src/backend/executor/nodeAgg.c
+++ b/src/backend/executor/nodeAgg.c
@@ -1776,11 +1776,24 @@ hash_agg_set_limits(double hashentrysize, uint64 input_groups, int used_bits,
 static void
 hash_agg_check_limits(AggState *aggstate)
 {
-	uint64 ngroups = aggstate->hash_ngroups_current;
-	Size meta_mem = MemoryContextMemAllocated(
-		aggstate->hash_metacxt, true);
-	Size hash_mem = MemoryContextMemAllocated(
-		aggstate->hashcontext->ecxt_per_tuple_memory, true);
+	MemoryContextCounters	counters;
+	uint64					ngroups = aggstate->hash_ngroups_current;
+	uint32					flags	= (MCXT_STAT_TOTALSPACE |
+									   MCXT_STAT_NEWSPACE);
+	Size					meta_mem;
+	Size					hash_mem;
+
+	/*
+	 * Consider all memory except "newspace", which is part of a block
+	 * allocation by the memory context itself that aggregation has little
+	 * control over. It doesn't make sense to count that against work_mem.
+	 */
+	counters = MemoryContextCount(aggstate->hash_metacxt, flags, true);
+	meta_mem = counters.totalspace - counters.newspace;
+
+	counters = MemoryContextCount(
+		aggstate->hashcontext->ecxt_per_tuple_memory, flags, true);
+	hash_mem = counters.totalspace - counters.newspace;
 
 	/*
 	 * Don't spill unless there's at least one group in the hash table so we
@@ -1838,21 +1851,25 @@ hash_agg_enter_spill_mode(AggState *aggstate)
 static void
 hash_agg_update_metrics(AggState *aggstate, bool from_tape, int npartitions)
 {
-	Size	meta_mem;
-	Size	hash_mem;
-	Size	buffer_mem;
-	Size	total_mem;
+	MemoryContextCounters	counters;
+	uint32					flags = MCXT_STAT_TOTALSPACE | MCXT_STAT_NEWSPACE;
+	Size					meta_mem;
+	Size					hash_mem;
+	Size					buffer_mem;
+	Size					total_mem;
 
 	if (aggstate->aggstrategy != AGG_MIXED &&
 		aggstate->aggstrategy != AGG_HASHED)
 		return;
 
 	/* memory for the hash table itself */
-	meta_mem = MemoryContextMemAllocated(aggstate->hash_metacxt, true);
+	counters = MemoryContextCount(aggstate->hash_metacxt, flags, true);
+	meta_mem = counters.totalspace - counters.newspace;
 
 	/* memory for the group keys and transition states */
-	hash_mem = MemoryContextMemAllocated(
-		aggstate->hashcontext->ecxt_per_tuple_memory, true);
+	counters = MemoryContextCount(
+		aggstate->hashcontext->ecxt_per_tuple_memory, flags, true);
+	hash_mem = counters.totalspace - counters.newspace;
 
 	/* memory for read/write tape buffers, if spilled */
 	buffer_mem = npartitions * HASHAGG_WRITE_BUFFER_SIZE;
diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index c0623f106d2..b4f3cb33ff9 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -132,6 +132,8 @@ typedef struct AllocSetContext
 	Size		maxBlockSize;	/* maximum block size */
 	Size		nextBlockSize;	/* next block size to allocate */
 	Size		allocChunkLimit;	/* effective chunk size limit */
+	Size		memAllocated;	/* track memory allocated for this context */
+	uint64		nChunks;		/* total number of chunks */
 	AllocBlock	keeper;			/* keep this block over resets */
 	/* freelist this context could be put in, or -1 if not a candidate: */
 	int			freeListIndex;	/* index in context_freelists[], or -1 */
@@ -273,9 +275,8 @@ static void AllocSetReset(MemoryContext context);
 static void AllocSetDelete(MemoryContext context);
 static Size AllocSetGetChunkSpace(MemoryContext context, void *pointer);
 static bool AllocSetIsEmpty(MemoryContext context);
-static void AllocSetStats(MemoryContext context,
-						  MemoryStatsPrintFunc printfunc, void *passthru,
-						  MemoryContextCounters *totals);
+static MemoryContextCounters AllocSetStats(MemoryContext context,
+										   uint32 flags);
 
 #ifdef MEMORY_CONTEXT_CHECKING
 static void AllocSetCheck(MemoryContext context);
@@ -464,8 +465,8 @@ AllocSetContextCreateInternal(MemoryContext parent,
 								parent,
 								name);
 
-			((MemoryContext) set)->mem_allocated =
-				set->keeper->endptr - ((char *) set);
+			set->memAllocated = set->keeper->endptr - ((char *) set);
+			set->nChunks = 0;
 
 			return (MemoryContext) set;
 		}
@@ -555,7 +556,8 @@ AllocSetContextCreateInternal(MemoryContext parent,
 						parent,
 						name);
 
-	((MemoryContext) set)->mem_allocated = firstBlockSize;
+	set->memAllocated = firstBlockSize;
+	set->nChunks = 0;
 
 	return (MemoryContext) set;
 }
@@ -617,7 +619,7 @@ AllocSetReset(MemoryContext context)
 		else
 		{
 			/* Normal case, release the block */
-			context->mem_allocated -= block->endptr - ((char*) block);
+			set->memAllocated -= block->endptr - ((char*) block);
 
 #ifdef CLOBBER_FREED_MEMORY
 			wipe_mem(block, block->freeptr - ((char *) block));
@@ -627,7 +629,9 @@ AllocSetReset(MemoryContext context)
 		block = next;
 	}
 
-	Assert(context->mem_allocated == keepersize);
+	Assert(set->memAllocated == keepersize);
+
+	set->nChunks = 0;
 
 	/* Reset block size allocation sequence, too */
 	set->nextBlockSize = set->initBlockSize;
@@ -703,7 +707,7 @@ AllocSetDelete(MemoryContext context)
 		AllocBlock	next = block->next;
 
 		if (block != set->keeper)
-			context->mem_allocated -= block->endptr - ((char *) block);
+			set->memAllocated -= block->endptr - ((char *) block);
 
 #ifdef CLOBBER_FREED_MEMORY
 		wipe_mem(block, block->freeptr - ((char *) block));
@@ -715,7 +719,7 @@ AllocSetDelete(MemoryContext context)
 		block = next;
 	}
 
-	Assert(context->mem_allocated == keepersize);
+	Assert(set->memAllocated == keepersize);
 
 	/* Finally, free the context header, including the keeper block */
 	free(set);
@@ -758,7 +762,7 @@ AllocSetAlloc(MemoryContext context, Size size)
 		if (block == NULL)
 			return NULL;
 
-		context->mem_allocated += blksize;
+		set->memAllocated += blksize;
 
 		block->aset = set;
 		block->freeptr = block->endptr = ((char *) block) + blksize;
@@ -805,6 +809,7 @@ AllocSetAlloc(MemoryContext context, Size size)
 		/* Disallow external access to private part of chunk header. */
 		VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
 
+		set->nChunks++;
 		return AllocChunkGetPointer(chunk);
 	}
 
@@ -844,6 +849,8 @@ AllocSetAlloc(MemoryContext context, Size size)
 		/* Disallow external access to private part of chunk header. */
 		VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
 
+		/* chunk already existed; don't increment nChunks */
+
 		return AllocChunkGetPointer(chunk);
 	}
 
@@ -906,6 +913,7 @@ AllocSetAlloc(MemoryContext context, Size size)
 #endif
 				chunk->aset = (void *) set->freelist[a_fidx];
 				set->freelist[a_fidx] = chunk;
+				set->nChunks++;
 			}
 
 			/* Mark that we need to create a new block */
@@ -955,7 +963,7 @@ AllocSetAlloc(MemoryContext context, Size size)
 		if (block == NULL)
 			return NULL;
 
-		context->mem_allocated += blksize;
+		set->memAllocated += blksize;
 
 		block->aset = set;
 		block->freeptr = ((char *) block) + ALLOC_BLOCKHDRSZ;
@@ -1005,6 +1013,7 @@ AllocSetAlloc(MemoryContext context, Size size)
 	/* Disallow external access to private part of chunk header. */
 	VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
 
+	set->nChunks++;
 	return AllocChunkGetPointer(chunk);
 }
 
@@ -1058,7 +1067,8 @@ AllocSetFree(MemoryContext context, void *pointer)
 		if (block->next)
 			block->next->prev = block->prev;
 
-		context->mem_allocated -= block->endptr - ((char*) block);
+		set->memAllocated -= block->endptr - ((char*) block);
+		set->nChunks--;
 
 #ifdef CLOBBER_FREED_MEMORY
 		wipe_mem(block, block->freeptr - ((char *) block));
@@ -1161,8 +1171,8 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 		}
 
 		/* updated separately, not to underflow when (oldblksize > blksize) */
-		context->mem_allocated -= oldblksize;
-		context->mem_allocated += blksize;
+		set->memAllocated -= oldblksize;
+		set->memAllocated += blksize;
 
 		block->freeptr = block->endptr = ((char *) block) + blksize;
 
@@ -1358,63 +1368,55 @@ AllocSetIsEmpty(MemoryContext context)
 /*
  * AllocSetStats
  *		Compute stats about memory consumption of an allocset.
- *
- * printfunc: if not NULL, pass a human-readable stats string to this.
- * passthru: pass this pointer through to printfunc.
- * totals: if not NULL, add stats about this context into *totals.
  */
-static void
-AllocSetStats(MemoryContext context,
-			  MemoryStatsPrintFunc printfunc, void *passthru,
-			  MemoryContextCounters *totals)
+static MemoryContextCounters
+AllocSetStats(MemoryContext context, uint32 flags)
 {
-	AllocSet	set = (AllocSet) context;
-	Size		nblocks = 0;
-	Size		freechunks = 0;
-	Size		totalspace;
-	Size		freespace = 0;
-	AllocBlock	block;
-	int			fidx;
-
-	/* Include context header in totalspace */
-	totalspace = MAXALIGN(sizeof(AllocSetContext));
-
-	for (block = set->blocks; block != NULL; block = block->next)
+	AllocSet				set		 = (AllocSet) context;
+	MemoryContextCounters	counters = {0};
+	uint64					nblocks = 0;
+	uint64					freechunks = 0;
+	Size					freespace = 0;
+	AllocBlock				block;
+	int						fidx;
+
+	if (flags & (MCXT_STAT_NBLOCKS | MCXT_STAT_FREESPACE))
 	{
-		nblocks++;
-		totalspace += block->endptr - ((char *) block);
-		freespace += block->endptr - block->freeptr;
-	}
-	for (fidx = 0; fidx < ALLOCSET_NUM_FREELISTS; fidx++)
-	{
-		AllocChunk	chunk;
-
-		for (chunk = set->freelist[fidx]; chunk != NULL;
-			 chunk = (AllocChunk) chunk->aset)
+		for (block = set->blocks; block != NULL; block = block->next)
 		{
-			freechunks++;
-			freespace += chunk->size + ALLOC_CHUNKHDRSZ;
+			nblocks++;
+			freespace += block->endptr - block->freeptr;
 		}
 	}
-
-	if (printfunc)
+	if (flags & (MCXT_STAT_FREECHUNKS | MCXT_STAT_FREESPACE))
 	{
-		char		stats_string[200];
+		for (fidx = 0; fidx < ALLOCSET_NUM_FREELISTS; fidx++)
+		{
+			AllocChunk	chunk;
 
-		snprintf(stats_string, sizeof(stats_string),
-				 "%zu total in %zd blocks; %zu free (%zd chunks); %zu used",
-				 totalspace, nblocks, freespace, freechunks,
-				 totalspace - freespace);
-		printfunc(context, passthru, stats_string);
+			for (chunk = set->freelist[fidx]; chunk != NULL;
+				 chunk = (AllocChunk) chunk->aset)
+			{
+				freechunks++;
+				freespace += chunk->size + ALLOC_CHUNKHDRSZ;
+			}
+		}
 	}
 
-	if (totals)
-	{
-		totals->nblocks += nblocks;
-		totals->freechunks += freechunks;
-		totals->totalspace += totalspace;
-		totals->freespace += freespace;
-	}
+	if (flags & MCXT_STAT_NBLOCKS)
+		counters.nblocks = nblocks;
+	if (flags & MCXT_STAT_NCHUNKS)
+		counters.nchunks = set->nChunks;
+	if (flags & MCXT_STAT_FREECHUNKS)
+		counters.freechunks = freechunks;
+	if (flags & MCXT_STAT_TOTALSPACE)
+		counters.totalspace = set->memAllocated;
+	if (flags & MCXT_STAT_FREESPACE)
+		counters.freespace = freespace;
+	if (flags & MCXT_STAT_NEWSPACE)
+		counters.newspace = set->blocks->endptr - set->blocks->freeptr;
+
+	return counters;
 }
 
 
@@ -1436,6 +1438,7 @@ AllocSetCheck(MemoryContext context)
 	AllocBlock	prevblock;
 	AllocBlock	block;
 	Size		total_allocated = 0;
+	uint64		total_nchunks = 0;
 
 	for (prevblock = NULL, block = set->blocks;
 		 block != NULL;
@@ -1529,6 +1532,7 @@ AllocSetCheck(MemoryContext context)
 
 			blk_data += chsize;
 			nchunks++;
+			total_nchunks++;
 
 			bpoz += ALLOC_CHUNKHDRSZ + chsize;
 		}
@@ -1538,7 +1542,8 @@ AllocSetCheck(MemoryContext context)
 				 name, block);
 	}
 
-	Assert(total_allocated == context->mem_allocated);
+	Assert(total_allocated == set->memAllocated);
+	Assert(total_nchunks == set->nChunks);
 }
 
 #endif							/* MEMORY_CONTEXT_CHECKING */
diff --git a/src/backend/utils/mmgr/generation.c b/src/backend/utils/mmgr/generation.c
index 56651d06931..1f3713cb27e 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -61,6 +61,7 @@ typedef struct GenerationContext
 
 	/* Generational context parameters */
 	Size		blockSize;		/* standard block size */
+	Size		memAllocated;	/* track memory allocated for this context */
 
 	GenerationBlock *block;		/* current (most recently allocated) block */
 	dlist_head	blocks;			/* list of blocks */
@@ -153,9 +154,8 @@ static void GenerationReset(MemoryContext context);
 static void GenerationDelete(MemoryContext context);
 static Size GenerationGetChunkSpace(MemoryContext context, void *pointer);
 static bool GenerationIsEmpty(MemoryContext context);
-static void GenerationStats(MemoryContext context,
-							MemoryStatsPrintFunc printfunc, void *passthru,
-							MemoryContextCounters *totals);
+static MemoryContextCounters GenerationStats(MemoryContext context,
+											 uint32 flags);
 
 #ifdef MEMORY_CONTEXT_CHECKING
 static void GenerationCheck(MemoryContext context);
@@ -258,6 +258,7 @@ GenerationContextCreate(MemoryContext parent,
 
 	/* Fill in GenerationContext-specific header fields */
 	set->blockSize = blockSize;
+	set->memAllocated = 0;
 	set->block = NULL;
 	dlist_init(&set->blocks);
 
@@ -297,7 +298,7 @@ GenerationReset(MemoryContext context)
 
 		dlist_delete(miter.cur);
 
-		context->mem_allocated -= block->blksize;
+		set->memAllocated -= block->blksize;
 
 #ifdef CLOBBER_FREED_MEMORY
 		wipe_mem(block, block->blksize);
@@ -354,7 +355,7 @@ GenerationAlloc(MemoryContext context, Size size)
 		if (block == NULL)
 			return NULL;
 
-		context->mem_allocated += blksize;
+		set->memAllocated += blksize;
 
 		/* block with a single (used) chunk */
 		block->blksize = blksize;
@@ -411,7 +412,7 @@ GenerationAlloc(MemoryContext context, Size size)
 		if (block == NULL)
 			return NULL;
 
-		context->mem_allocated += blksize;
+		set->memAllocated += blksize;
 
 		block->blksize = blksize;
 		block->nchunks = 0;
@@ -528,7 +529,7 @@ GenerationFree(MemoryContext context, void *pointer)
 	if (set->block == block)
 		set->block = NULL;
 
-	context->mem_allocated -= block->blksize;
+	set->memAllocated -= block->blksize;
 	free(block);
 }
 
@@ -681,59 +682,47 @@ GenerationIsEmpty(MemoryContext context)
 /*
  * GenerationStats
  *		Compute stats about memory consumption of a Generation context.
- *
- * printfunc: if not NULL, pass a human-readable stats string to this.
- * passthru: pass this pointer through to printfunc.
- * totals: if not NULL, add stats about this context into *totals.
- *
- * XXX freespace only accounts for empty space at the end of the block, not
- * space of freed chunks (which is unknown).
  */
-static void
-GenerationStats(MemoryContext context,
-				MemoryStatsPrintFunc printfunc, void *passthru,
-				MemoryContextCounters *totals)
+static MemoryContextCounters
+GenerationStats(MemoryContext context, uint32 flags)
 {
-	GenerationContext *set = (GenerationContext *) context;
-	Size		nblocks = 0;
-	Size		nchunks = 0;
-	Size		nfreechunks = 0;
-	Size		totalspace;
-	Size		freespace = 0;
-	dlist_iter	iter;
-
-	/* Include context header in totalspace */
-	totalspace = MAXALIGN(sizeof(GenerationContext));
-
-	dlist_foreach(iter, &set->blocks)
-	{
-		GenerationBlock *block = dlist_container(GenerationBlock, node, iter.cur);
-
-		nblocks++;
-		nchunks += block->nchunks;
-		nfreechunks += block->nfree;
-		totalspace += block->blksize;
-		freespace += (block->endptr - block->freeptr);
-	}
-
-	if (printfunc)
+	GenerationContext		*set	  = (GenerationContext *) context;
+	MemoryContextCounters	 counters = {0};
+	uint64					 nblocks = 0;
+	uint64					 nchunks = 0;
+	uint64					 freechunks = 0;
+	Size					 freespace = 0;
+	dlist_iter				 iter;
+
+	if (flags & (MCXT_STAT_NBLOCKS | MCXT_STAT_NCHUNKS |
+				 MCXT_STAT_FREECHUNKS | MCXT_STAT_FREESPACE))
 	{
-		char		stats_string[200];
+		dlist_foreach(iter, &set->blocks)
+		{
+			GenerationBlock *block = dlist_container(
+				GenerationBlock, node, iter.cur);
 
-		snprintf(stats_string, sizeof(stats_string),
-				 "%zu total in %zd blocks (%zd chunks); %zu free (%zd chunks); %zu used",
-				 totalspace, nblocks, nchunks, freespace,
-				 nfreechunks, totalspace - freespace);
-		printfunc(context, passthru, stats_string);
+			nblocks++;
+			nchunks += block->nchunks;
+			freechunks += block->nfree;
+			freespace += (block->endptr - block->freeptr);
+		}
 	}
 
-	if (totals)
-	{
-		totals->nblocks += nblocks;
-		totals->freechunks += nfreechunks;
-		totals->totalspace += totalspace;
-		totals->freespace += freespace;
-	}
+	if (flags & MCXT_STAT_NBLOCKS)
+		counters.nblocks = nblocks;
+	if (flags & MCXT_STAT_NCHUNKS)
+		counters.nchunks = nchunks;
+	if (flags & MCXT_STAT_FREECHUNKS)
+		counters.freechunks = freechunks;
+	if (flags & MCXT_STAT_TOTALSPACE)
+		counters.totalspace = set->memAllocated;
+	if (flags & MCXT_STAT_FREESPACE)
+		counters.freespace = freespace;
+	if (flags & MCXT_STAT_NEWSPACE)
+		counters.newspace = set->block->endptr - set->block->freeptr;
+
+	return counters;
 }
 
 
@@ -844,7 +833,7 @@ GenerationCheck(MemoryContext context)
 				 name, nfree, block, block->nfree);
 	}
 
-	Assert(total_allocated == context->mem_allocated);
+	Assert(total_allocated == gen->memAllocated);
 }
 
 #endif							/* MEMORY_CONTEXT_CHECKING */
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 9e24fec72d6..bd8ee42405f 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -56,7 +56,7 @@ static void MemoryContextCallResetCallbacks(MemoryContext context);
 static void MemoryContextStatsInternal(MemoryContext context, int level,
 									   bool print, int max_children,
 									   MemoryContextCounters *totals);
-static void MemoryContextStatsPrint(MemoryContext context, void *passthru,
+static void MemoryContextStatsPrint(MemoryContext context, int level,
 									const char *stats_string);
 
 /*
@@ -463,24 +463,39 @@ MemoryContextIsEmpty(MemoryContext context)
 }
 
 /*
- * Find the memory allocated to blocks for this memory context. If recurse is
- * true, also include children.
+ * MemoryContextCount
+ *		Return statistics about this memory context, optionally recursing to
+ *		children. Flags are defined in memnodes.h and specify which statistics
+ *		are required.
  */
-Size
-MemoryContextMemAllocated(MemoryContext context, bool recurse)
+MemoryContextCounters
+MemoryContextCount(MemoryContext context, uint32 flags, bool recurse)
 {
-	Size	total = context->mem_allocated;
+	MemoryContextCounters	total;
 
 	AssertArg(MemoryContextIsValid(context));
 
+	total = context->methods->count(context, flags);
+
 	if (recurse)
 	{
-		MemoryContext child = context->firstchild;
+		MemoryContext			child;
 
 		for (child = context->firstchild;
 			 child != NULL;
 			 child = child->nextchild)
-			total += MemoryContextMemAllocated(child, true);
+		{
+			MemoryContextCounters	child_counters;
+			
+			child_counters = MemoryContextCount(child, flags, true);
+
+			total.nblocks += child_counters.nblocks;
+			total.nchunks += child_counters.nchunks;
+			total.freechunks += child_counters.freechunks;
+			total.totalspace += child_counters.totalspace;
+			total.freespace += child_counters.freespace;
+			total.newspace += child_counters.newspace;
+		}
 	}
 
 	return total;
@@ -516,9 +531,10 @@ MemoryContextStatsDetail(MemoryContext context, int max_children)
 	MemoryContextStatsInternal(context, 0, true, max_children, &grand_totals);
 
 	fprintf(stderr,
-			"Grand total: %zu bytes in %zd blocks; %zu free (%zd chunks); %zu used\n",
+			"Grand total: %zu bytes in %zd blocks (%zd chunks); %zu free (%zd chunks); %zu used\n",
 			grand_totals.totalspace, grand_totals.nblocks,
-			grand_totals.freespace, grand_totals.freechunks,
+			grand_totals.nchunks, grand_totals.freespace,
+			grand_totals.freechunks,
 			grand_totals.totalspace - grand_totals.freespace);
 }
 
@@ -534,24 +550,42 @@ MemoryContextStatsInternal(MemoryContext context, int level,
 						   bool print, int max_children,
 						   MemoryContextCounters *totals)
 {
-	MemoryContextCounters local_totals;
+	MemoryContextCounters excess_summary = {0};
+	MemoryContextCounters current;
 	MemoryContext child;
 	int			ichild;
 
+
 	AssertArg(MemoryContextIsValid(context));
 
 	/* Examine the context itself */
-	context->methods->stats(context,
-							print ? MemoryContextStatsPrint : NULL,
-							(void *) &level,
-							totals);
+	current = context->methods->count(context, MCXT_STAT_ALL);
+
+	if (print)
+	{
+		char		stats_string[200];
+		snprintf(stats_string, sizeof(stats_string),
+				 "%zu total in %zd blocks (%zd chunks); %zu free (%zd chunks); %zu used",
+				 current.totalspace, current.nblocks, current.nchunks,
+				 current.freespace, current.freechunks,
+				 current.totalspace - current.freespace);
+		MemoryContextStatsPrint(context, level, stats_string);
+	}
+
+	if (totals)
+	{
+		totals->nblocks += current.nblocks;
+		totals->nchunks += current.nchunks;
+		totals->freechunks += current.freechunks;
+		totals->totalspace += current.totalspace;
+		totals->freespace += current.freespace;
+		totals->newspace += current.newspace;
+	}
 
 	/*
 	 * Examine children.  If there are more than max_children of them, we do
 	 * not print the rest explicitly, but just summarize them.
 	 */
-	memset(&local_totals, 0, sizeof(local_totals));
-
 	for (child = context->firstchild, ichild = 0;
 		 child != NULL;
 		 child = child->nextchild, ichild++)
@@ -563,7 +597,7 @@ MemoryContextStatsInternal(MemoryContext context, int level,
 		else
 			MemoryContextStatsInternal(child, level + 1,
 									   false, max_children,
-									   &local_totals);
+									   &excess_summary);
 	}
 
 	/* Deal with excess children */
@@ -578,19 +612,21 @@ MemoryContextStatsInternal(MemoryContext context, int level,
 			fprintf(stderr,
 					"%d more child contexts containing %zu total in %zd blocks; %zu free (%zd chunks); %zu used\n",
 					ichild - max_children,
-					local_totals.totalspace,
-					local_totals.nblocks,
-					local_totals.freespace,
-					local_totals.freechunks,
-					local_totals.totalspace - local_totals.freespace);
+					excess_summary.totalspace,
+					excess_summary.nblocks,
+					excess_summary.freespace,
+					excess_summary.freechunks,
+					excess_summary.totalspace - excess_summary.freespace);
 		}
 
 		if (totals)
 		{
-			totals->nblocks += local_totals.nblocks;
-			totals->freechunks += local_totals.freechunks;
-			totals->totalspace += local_totals.totalspace;
-			totals->freespace += local_totals.freespace;
+			totals->nblocks += excess_summary.nblocks;
+			totals->nchunks += excess_summary.nchunks;
+			totals->freechunks += excess_summary.freechunks;
+			totals->totalspace += excess_summary.totalspace;
+			totals->freespace += excess_summary.freespace;
+			totals->newspace += excess_summary.newspace;
 		}
 	}
 }
@@ -603,10 +639,9 @@ MemoryContextStatsInternal(MemoryContext context, int level,
  * make that more complicated.
  */
 static void
-MemoryContextStatsPrint(MemoryContext context, void *passthru,
+MemoryContextStatsPrint(MemoryContext context, int level,
 						const char *stats_string)
 {
-	int			level = *(int *) passthru;
 	const char *name = context->name;
 	const char *ident = context->ident;
 	int			i;
@@ -760,7 +795,6 @@ MemoryContextCreate(MemoryContext node,
 	node->methods = methods;
 	node->parent = parent;
 	node->firstchild = NULL;
-	node->mem_allocated = 0;
 	node->prevchild = NULL;
 	node->name = name;
 	node->ident = NULL;
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index c928476c479..e3921991248 100644
--- a/src/backend/utils/mmgr/slab.c
+++ b/src/backend/utils/mmgr/slab.c
@@ -67,6 +67,7 @@ typedef struct SlabContext
 	Size		fullChunkSize;	/* chunk size including header and alignment */
 	Size		blockSize;		/* block size */
 	Size		headerSize;		/* allocated size of context header */
+	Size		memAllocated;	/* track memory allocated for this context */
 	int			chunksPerBlock; /* number of chunks per block */
 	int			minFreeChunks;	/* min number of free chunks in any block */
 	int			nblocks;		/* number of blocks allocated */
@@ -133,9 +134,7 @@ static void SlabReset(MemoryContext context);
 static void SlabDelete(MemoryContext context);
 static Size SlabGetChunkSpace(MemoryContext context, void *pointer);
 static bool SlabIsEmpty(MemoryContext context);
-static void SlabStats(MemoryContext context,
-					  MemoryStatsPrintFunc printfunc, void *passthru,
-					  MemoryContextCounters *totals);
+static MemoryContextCounters SlabStats(MemoryContext context, uint32 flags);
 #ifdef MEMORY_CONTEXT_CHECKING
 static void SlabCheck(MemoryContext context);
 #endif
@@ -262,6 +261,7 @@ SlabContextCreate(MemoryContext parent,
 	slab->fullChunkSize = fullChunkSize;
 	slab->blockSize = blockSize;
 	slab->headerSize = headerSize;
+	slab->memAllocated = 0;
 	slab->chunksPerBlock = chunksPerBlock;
 	slab->minFreeChunks = 0;
 	slab->nblocks = 0;
@@ -322,14 +322,14 @@ SlabReset(MemoryContext context)
 #endif
 			free(block);
 			slab->nblocks--;
-			context->mem_allocated -= slab->blockSize;
+			slab->memAllocated -= slab->blockSize;
 		}
 	}
 
 	slab->minFreeChunks = 0;
 
 	Assert(slab->nblocks == 0);
-	Assert(context->mem_allocated == 0);
+	Assert(slab->memAllocated == 0);
 }
 
 /*
@@ -407,7 +407,7 @@ SlabAlloc(MemoryContext context, Size size)
 
 		slab->minFreeChunks = slab->chunksPerBlock;
 		slab->nblocks += 1;
-		context->mem_allocated += slab->blockSize;
+		slab->memAllocated += slab->blockSize;
 	}
 
 	/* grab the block from the freelist (even the new block is there) */
@@ -501,7 +501,7 @@ SlabAlloc(MemoryContext context, Size size)
 
 	SlabAllocInfo(slab, chunk);
 
-	Assert(slab->nblocks * slab->blockSize == context->mem_allocated);
+	Assert(slab->nblocks * slab->blockSize == slab->memAllocated);
 
 	return SlabChunkGetPointer(chunk);
 }
@@ -578,13 +578,13 @@ SlabFree(MemoryContext context, void *pointer)
 	{
 		free(block);
 		slab->nblocks--;
-		context->mem_allocated -= slab->blockSize;
+		slab->memAllocated -= slab->blockSize;
 	}
 	else
 		dlist_push_head(&slab->freelist[block->nfree], &block->node);
 
 	Assert(slab->nblocks >= 0);
-	Assert(slab->nblocks * slab->blockSize == context->mem_allocated);
+	Assert(slab->nblocks * slab->blockSize == slab->memAllocated);
 }
 
 /*
@@ -647,59 +647,50 @@ SlabIsEmpty(MemoryContext context)
 /*
  * SlabStats
  *		Compute stats about memory consumption of a Slab context.
- *
- * printfunc: if not NULL, pass a human-readable stats string to this.
- * passthru: pass this pointer through to printfunc.
- * totals: if not NULL, add stats about this context into *totals.
  */
-static void
-SlabStats(MemoryContext context,
-		  MemoryStatsPrintFunc printfunc, void *passthru,
-		  MemoryContextCounters *totals)
+static MemoryContextCounters
+SlabStats(MemoryContext context, uint32 flags)
 {
-	SlabContext *slab = castNode(SlabContext, context);
-	Size		nblocks = 0;
-	Size		freechunks = 0;
-	Size		totalspace;
-	Size		freespace = 0;
-	int			i;
-
-	/* Include context header in totalspace */
-	totalspace = slab->headerSize;
-
-	for (i = 0; i <= slab->chunksPerBlock; i++)
+	SlabContext				*slab	  = castNode(SlabContext, context);
+	MemoryContextCounters	 counters = {0};
+	uint64					 nblocks = 0;
+	uint64					 nchunks = 0;
+	uint64					 freechunks = 0;
+	Size					 freespace = 0;
+	int						 i;
+
+	if (flags & (MCXT_STAT_NBLOCKS | MCXT_STAT_NCHUNKS |
+				 MCXT_STAT_FREECHUNKS | MCXT_STAT_FREESPACE))
 	{
-		dlist_iter	iter;
-
-		dlist_foreach(iter, &slab->freelist[i])
+		for (i = 0; i <= slab->chunksPerBlock; i++)
 		{
-			SlabBlock  *block = dlist_container(SlabBlock, node, iter.cur);
+			dlist_iter	iter;
 
-			nblocks++;
-			totalspace += slab->blockSize;
-			freespace += slab->fullChunkSize * block->nfree;
-			freechunks += block->nfree;
-		}
-	}
-
-	if (printfunc)
-	{
-		char		stats_string[200];
+			dlist_foreach(iter, &slab->freelist[i])
+			{
+				SlabBlock  *block = dlist_container(SlabBlock, node, iter.cur);
 
-		snprintf(stats_string, sizeof(stats_string),
-				 "%zu total in %zd blocks; %zu free (%zd chunks); %zu used",
-				 totalspace, nblocks, freespace, freechunks,
-				 totalspace - freespace);
-		printfunc(context, passthru, stats_string);
+				nblocks++;
+				nchunks += slab->chunksPerBlock;
+				freespace += slab->fullChunkSize * block->nfree;
+				freechunks += block->nfree;
+			}
+		}
 	}
 
-	if (totals)
-	{
-		totals->nblocks += nblocks;
-		totals->freechunks += freechunks;
-		totals->totalspace += totalspace;
-		totals->freespace += freespace;
-	}
+	if (flags & MCXT_STAT_NBLOCKS)
+		counters.nblocks = nblocks;
+	if (flags & MCXT_STAT_NCHUNKS)
+		counters.nchunks = nchunks;
+	if (flags & MCXT_STAT_FREECHUNKS)
+		counters.freechunks = freechunks;
+	if (flags & MCXT_STAT_TOTALSPACE)
+		counters.totalspace = slab->memAllocated;
+	if (flags & MCXT_STAT_FREESPACE)
+		counters.freespace = freespace;
+	/* new memory is already sliced into chunks, so newspace is always 0 */
+
+	return counters;
 }
 
 
@@ -804,7 +795,7 @@ SlabCheck(MemoryContext context)
 		}
 	}
 
-	Assert(slab->nblocks * slab->blockSize == context->mem_allocated);
+	Assert(slab->nblocks * slab->blockSize == slab->memAllocated);
 }
 
 #endif							/* MEMORY_CONTEXT_CHECKING */
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index c9f2bbcb367..cc545852968 100644
--- a/src/include/nodes/memnodes.h
+++ b/src/include/nodes/memnodes.h
@@ -29,11 +29,21 @@
 typedef struct MemoryContextCounters
 {
 	Size		nblocks;		/* Total number of malloc blocks */
+	Size		nchunks;		/* Total number of chunks (used+free) */
 	Size		freechunks;		/* Total number of free chunks */
 	Size		totalspace;		/* Total bytes requested from malloc */
 	Size		freespace;		/* The unused portion of totalspace */
+	Size		newspace;		/* Allocated but never held any chunks */
 } MemoryContextCounters;
 
+#define MCXT_STAT_NBLOCKS		(1 << 0)
+#define MCXT_STAT_NCHUNKS		(1 << 1)
+#define MCXT_STAT_FREECHUNKS	(1 << 2)
+#define MCXT_STAT_TOTALSPACE	(1 << 3)
+#define MCXT_STAT_FREESPACE		(1 << 4)
+#define MCXT_STAT_NEWSPACE		(1 << 5)
+#define MCXT_STAT_ALL			((1 << 6) - 1)
+
 /*
  * MemoryContext
  *		A logical context in which memory allocations occur.
@@ -51,9 +61,6 @@ typedef struct MemoryContextCounters
  * to the context struct rather than the struct type itself.
  */
 
-typedef void (*MemoryStatsPrintFunc) (MemoryContext context, void *passthru,
-									  const char *stats_string);
-
 typedef struct MemoryContextMethods
 {
 	void	   *(*alloc) (MemoryContext context, Size size);
@@ -64,9 +71,7 @@ typedef struct MemoryContextMethods
 	void		(*delete_context) (MemoryContext context);
 	Size		(*get_chunk_space) (MemoryContext context, void *pointer);
 	bool		(*is_empty) (MemoryContext context);
-	void		(*stats) (MemoryContext context,
-						  MemoryStatsPrintFunc printfunc, void *passthru,
-						  MemoryContextCounters *totals);
+	MemoryContextCounters (*count) (MemoryContext context, uint32 flags);
 #ifdef MEMORY_CONTEXT_CHECKING
 	void		(*check) (MemoryContext context);
 #endif
@@ -79,7 +84,6 @@ typedef struct MemoryContextData
 	/* these two fields are placed here to minimize alignment wastage: */
 	bool		isReset;		/* T = no space alloced since last reset */
 	bool		allowInCritSection; /* allow palloc in critical section */
-	Size		mem_allocated;	/* track memory allocated for this context */
 	const MemoryContextMethods *methods;	/* virtual function table */
 	MemoryContext parent;		/* NULL if no parent (toplevel context) */
 	MemoryContext firstchild;	/* head of linked list of children */
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index 909bc2e9888..fb2d960e333 100644
--- a/src/include/utils/memutils.h
+++ b/src/include/utils/memutils.h
@@ -82,7 +82,8 @@ extern void MemoryContextSetParent(MemoryContext context,
 extern Size GetMemoryChunkSpace(void *pointer);
 extern MemoryContext MemoryContextGetParent(MemoryContext context);
 extern bool MemoryContextIsEmpty(MemoryContext context);
-extern Size MemoryContextMemAllocated(MemoryContext context, bool recurse);
+extern MemoryContextCounters MemoryContextCount(MemoryContext context,
+												uint32 flags, bool recurse);
 extern void MemoryContextStats(MemoryContext context);
 extern void MemoryContextStatsDetail(MemoryContext context, int max_children);
 extern void MemoryContextAllowInCriticalSection(MemoryContext context,

Reply via email to