Make MemoryContextMemAllocated() more precise

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,

diff --git a/src/backend/executor/nodeAgg.c b/src/backend/executor/nodeAgg.c
index 4e100e5755f..fd205071ca1 100644
--- a/src/backend/executor/nodeAgg.c
+++ b/src/backend/executor/nodeAgg.c
@@ -414,6 +414,7 @@ static void agg_fill_hash_table(AggState *aggstate);
 static bool agg_refill_hash_table(AggState *aggstate);
 static TupleTableSlot *agg_retrieve_hash_table(AggState *aggstate);
 static TupleTableSlot *agg_retrieve_hash_table_in_memory(AggState *aggstate);
+static Size hash_agg_check_memory(MemoryContext context);
 static void hash_agg_check_limits(AggState *aggstate);
 static void hash_agg_enter_spill_mode(AggState *aggstate);
 static void hash_agg_update_metrics(AggState *aggstate, bool from_tape,
@@ -1792,6 +1793,25 @@ hash_agg_set_limits(double hashentrysize, uint64 input_groups, int used_bits,
 		*ngroups_limit = 1;
 }
 
+/*
+ * Get the memory consumed by the memory context. Consider all memory except
+ * "newspace", which is the untouched part of a new block allocation by the
+ * memory context itself. The algorithm has little control over how eagerly
+ * the memory context allocates new memory, so it doesn't make sense to count
+ * it against work_mem.
+ */
+static Size
+hash_agg_check_memory(MemoryContext context)
+{
+	MemoryContextCounters	counters;
+	uint32					flags = (MCXT_COUNT_TOTALSPACE |
+									 MCXT_COUNT_NEWSPACE);
+
+	MemoryContextGetCounters(context, flags, true, &counters);
+
+	return counters.totalspace - counters.newspace;
+}
+
 /*
  * hash_agg_check_limits
  *
@@ -1802,11 +1822,17 @@ 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);
+	uint64	ngroups = aggstate->hash_ngroups_current;
+	Size	meta_mem;
+	Size	hash_mem;
+
+	/* already in spill mode; nothing to do */
+	if (aggstate->hash_spill_mode)
+		return;
+
+	meta_mem = hash_agg_check_memory(aggstate->hash_metacxt);
+	hash_mem = hash_agg_check_memory(
+		aggstate->hashcontext->ecxt_per_tuple_memory);
 
 	/*
 	 * Don't spill unless there's at least one group in the hash table so we
@@ -1874,11 +1900,11 @@ hash_agg_update_metrics(AggState *aggstate, bool from_tape, int npartitions)
 		return;
 
 	/* memory for the hash table itself */
-	meta_mem = MemoryContextMemAllocated(aggstate->hash_metacxt, true);
+	meta_mem = hash_agg_check_memory(aggstate->hash_metacxt);
 
 	/* memory for the group keys and transition states */
-	hash_mem = MemoryContextMemAllocated(
-		aggstate->hashcontext->ecxt_per_tuple_memory, true);
+	hash_mem = hash_agg_check_memory(
+		aggstate->hashcontext->ecxt_per_tuple_memory);
 
 	/* 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..1dfc6126ad4 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 void AllocSetGetCounters(MemoryContext context, uint32 flags,
+								MemoryContextCounters *counters);
 
 #ifdef MEMORY_CONTEXT_CHECKING
 static void AllocSetCheck(MemoryContext context);
@@ -292,7 +293,7 @@ static const MemoryContextMethods AllocSetMethods = {
 	AllocSetDelete,
 	AllocSetGetChunkSpace,
 	AllocSetIsEmpty,
-	AllocSetStats
+	AllocSetGetCounters
 #ifdef MEMORY_CONTEXT_CHECKING
 	,AllocSetCheck
 #endif
@@ -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;
 
@@ -1356,65 +1366,53 @@ 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.
+ * AllocSetGetCounters
+ *		Compute and return memory consumption counters. Specify the needed
+ *		counters by setting the appropriate flags.
  */
 static void
-AllocSetStats(MemoryContext context,
-			  MemoryStatsPrintFunc printfunc, void *passthru,
-			  MemoryContextCounters *totals)
+AllocSetGetCounters(MemoryContext context, uint32 flags,
+					MemoryContextCounters *counters)
 {
-	AllocSet	set = (AllocSet) context;
-	Size		nblocks = 0;
+	AllocSet	set		   = (AllocSet) context;
+	Size		nblocks	   = 0;
 	Size		freechunks = 0;
-	Size		totalspace;
-	Size		freespace = 0;
-	AllocBlock	block;
-	int			fidx;
+	Size		freespace  = 0;
 
-	/* Include context header in totalspace */
-	totalspace = MAXALIGN(sizeof(AllocSetContext));
-
-	for (block = set->blocks; block != NULL; block = block->next)
-	{
-		nblocks++;
-		totalspace += block->endptr - ((char *) block);
-		freespace += block->endptr - block->freeptr;
-	}
-	for (fidx = 0; fidx < ALLOCSET_NUM_FREELISTS; fidx++)
+	if (flags & (MCXT_COUNT_NBLOCKS | MCXT_COUNT_FREESPACE))
 	{
-		AllocChunk	chunk;
+		AllocBlock	block;
 
-		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_COUNT_FREECHUNKS | MCXT_COUNT_FREESPACE))
 	{
-		char		stats_string[200];
+		int		fidx;
 
-		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 (fidx = 0; fidx < ALLOCSET_NUM_FREELISTS; fidx++)
+		{
+			AllocChunk	chunk;
 
-	if (totals)
-	{
-		totals->nblocks += nblocks;
-		totals->freechunks += freechunks;
-		totals->totalspace += totalspace;
-		totals->freespace += freespace;
+			for (chunk = set->freelist[fidx]; chunk != NULL;
+				 chunk = (AllocChunk) chunk->aset)
+			{
+				freechunks++;
+				freespace += chunk->size + ALLOC_CHUNKHDRSZ;
+			}
+		}
 	}
+
+	counters->nblocks = nblocks;
+	counters->nchunks = set->nChunks;
+	counters->freechunks = freechunks;
+	counters->totalspace = set->memAllocated;
+	counters->freespace = freespace;
+	counters->newspace = set->blocks->endptr - set->blocks->freeptr;
 }
 
 
@@ -1436,6 +1434,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 +1528,7 @@ AllocSetCheck(MemoryContext context)
 
 			blk_data += chsize;
 			nchunks++;
+			total_nchunks++;
 
 			bpoz += ALLOC_CHUNKHDRSZ + chsize;
 		}
@@ -1538,7 +1538,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..f1a989dd1e6 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 void GenerationGetCounters(MemoryContext context, uint32 flags,
+								  MemoryContextCounters *counters);
 
 #ifdef MEMORY_CONTEXT_CHECKING
 static void GenerationCheck(MemoryContext context);
@@ -172,7 +172,7 @@ static const MemoryContextMethods GenerationMethods = {
 	GenerationDelete,
 	GenerationGetChunkSpace,
 	GenerationIsEmpty,
-	GenerationStats
+	GenerationGetCounters
 #ifdef MEMORY_CONTEXT_CHECKING
 	,GenerationCheck
 #endif
@@ -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);
 }
 
@@ -679,61 +680,43 @@ 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).
+ * GenerationGetCounters
+ *		Compute and return memory consumption counters. Specify the needed
+ *		counters by setting the appropriate flags.
  */
 static void
-GenerationStats(MemoryContext context,
-				MemoryStatsPrintFunc printfunc, void *passthru,
-				MemoryContextCounters *totals)
+GenerationGetCounters(MemoryContext context, uint32 flags,
+					  MemoryContextCounters *counters)
 {
-	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)
+	GenerationContext		*set	  = (GenerationContext *) context;
+	uint64					 nblocks = 0;
+	uint64					 nchunks = 0;
+	uint64					 freechunks = 0;
+	Size					 freespace = 0;
+
+	if (flags & (MCXT_COUNT_NBLOCKS | MCXT_COUNT_NCHUNKS |
+				 MCXT_COUNT_FREECHUNKS | MCXT_COUNT_FREESPACE))
 	{
-		GenerationBlock *block = dlist_container(GenerationBlock, node, iter.cur);
+		dlist_iter	iter;
 
-		nblocks++;
-		nchunks += block->nchunks;
-		nfreechunks += block->nfree;
-		totalspace += block->blksize;
-		freespace += (block->endptr - block->freeptr);
-	}
-
-	if (printfunc)
-	{
-		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;
-	}
+	counters->nblocks = nblocks;
+	counters->nchunks = nchunks;
+	counters->freechunks = freechunks;
+	counters->totalspace = set->memAllocated;
+	counters->freespace = freespace;
+	counters->newspace = set->block->endptr - set->block->freeptr;
 }
 
 
@@ -844,7 +827,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..1708413c46c 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,27 +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)
+void
+MemoryContextGetCounters(MemoryContext context, uint32 flags, bool recurse,
+						 MemoryContextCounters *counters)
 {
-	Size	total = context->mem_allocated;
-
 	AssertArg(MemoryContextIsValid(context));
 
+	context->methods->get_counters(context, flags, counters);
+
 	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->methods->get_counters(child, flags, &child_counters);
 
-	return total;
+			counters->nblocks += child_counters.nblocks;
+			counters->nchunks += child_counters.nchunks;
+			counters->freechunks += child_counters.freechunks;
+			counters->totalspace += child_counters.totalspace;
+			counters->freespace += child_counters.freespace;
+			counters->newspace += child_counters.newspace;
+		}
+	}
 }
 
 /*
@@ -516,9 +528,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 +547,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);
+	context->methods->get_counters(context, MCXT_COUNT_ALL, &current);
+
+	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 +594,7 @@ MemoryContextStatsInternal(MemoryContext context, int level,
 		else
 			MemoryContextStatsInternal(child, level + 1,
 									   false, max_children,
-									   &local_totals);
+									   &excess_summary);
 	}
 
 	/* Deal with excess children */
@@ -578,35 +609,33 @@ 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;
 		}
 	}
 }
 
 /*
  * MemoryContextStatsPrint
- *		Print callback used by MemoryContextStatsInternal
- *
- * For now, the passthru pointer just points to "int level"; later we might
- * make that more complicated.
+ *		Print function used by MemoryContextStatsInternal
  */
 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 +789,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..acab22ffa49 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,8 @@ 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 void SlabGetCounters(MemoryContext context, uint32 flags,
+							MemoryContextCounters *counters);
 #ifdef MEMORY_CONTEXT_CHECKING
 static void SlabCheck(MemoryContext context);
 #endif
@@ -151,7 +151,7 @@ static const MemoryContextMethods SlabMethods = {
 	SlabDelete,
 	SlabGetChunkSpace,
 	SlabIsEmpty,
-	SlabStats
+	SlabGetCounters
 #ifdef MEMORY_CONTEXT_CHECKING
 	,SlabCheck
 #endif
@@ -262,6 +262,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 +323,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 +408,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 +502,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 +579,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);
 }
 
 /*
@@ -645,61 +646,48 @@ 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.
+ * SlabGetCounters
+ *		Compute and return memory consumption counters. Specify the needed
+ *		counters by setting the appropriate flags.
  */
 static void
-SlabStats(MemoryContext context,
-		  MemoryStatsPrintFunc printfunc, void *passthru,
-		  MemoryContextCounters *totals)
+SlabGetCounters(MemoryContext context, uint32 flags,
+				MemoryContextCounters *counters)
 {
-	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);
+	uint64					 nblocks = 0;
+	uint64					 nchunks = 0;
+	uint64					 freechunks = 0;
+	Size					 freespace = 0;
+
+	if (flags & (MCXT_COUNT_NBLOCKS | MCXT_COUNT_NCHUNKS |
+				 MCXT_COUNT_FREECHUNKS | MCXT_COUNT_FREESPACE))
 	{
-		dlist_iter	iter;
+		int	i;
 
-		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;
-	}
+	counters->nblocks = nblocks;
+	counters->nchunks = nchunks;
+	counters->freechunks = freechunks;
+	counters->totalspace = slab->memAllocated;
+	counters->freespace = freespace;
+	/* new memory is already sliced into chunks, so newspace is always 0 */
+	counters->newspace = 0;
 }
 
 
@@ -804,7 +792,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..798854da99d 100644
--- a/src/include/nodes/memnodes.h
+++ b/src/include/nodes/memnodes.h
@@ -25,15 +25,30 @@
  * we might need more or different counters here.  A possible API spec then
  * would be to print only nonzero counters, but for now we just summarize in
  * the format historically used by AllocSet.
+ *
+ * The 'newspace' counter represents memory that is allocated by malloc, but
+ * not touched yet. This can be significant for context types that malloc
+ * large blocks of memory in anticipation of chunks that have not yet been
+ * requested.
  */
 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_COUNT_NBLOCKS		(1 << 0)
+#define MCXT_COUNT_NCHUNKS		(1 << 1)
+#define MCXT_COUNT_FREECHUNKS	(1 << 2)
+#define MCXT_COUNT_TOTALSPACE	(1 << 3)
+#define MCXT_COUNT_FREESPACE	(1 << 4)
+#define MCXT_COUNT_NEWSPACE		(1 << 5)
+#define MCXT_COUNT_ALL			(~0)
+
 /*
  * MemoryContext
  *		A logical context in which memory allocations occur.
@@ -51,9 +66,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 +76,8 @@ 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);
+	void		(*get_counters) (MemoryContext context, uint32 flags,
+								 MemoryContextCounters *counters);
 #ifdef MEMORY_CONTEXT_CHECKING
 	void		(*check) (MemoryContext context);
 #endif
@@ -79,7 +90,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..e787e43b3fb 100644
--- a/src/include/utils/memutils.h
+++ b/src/include/utils/memutils.h
@@ -82,7 +82,9 @@ 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 void MemoryContextGetCounters(MemoryContext context, uint32 flags,
+									 bool recurse,
+									 MemoryContextCounters *counters);
 extern void MemoryContextStats(MemoryContext context);
 extern void MemoryContextStatsDetail(MemoryContext context, int max_children);
 extern void MemoryContextAllowInCriticalSection(MemoryContext context,

diff --git a/src/backend/executor/execUtils.c b/src/backend/executor/execUtils.c
index cc5177cc2b9..ca973882d01 100644
--- a/src/backend/executor/execUtils.c
+++ b/src/backend/executor/execUtils.c
@@ -53,6 +53,7 @@
 #include "executor/executor.h"
 #include "jit/jit.h"
 #include "mb/pg_wchar.h"
+#include "miscadmin.h"
 #include "nodes/nodeFuncs.h"
 #include "parser/parsetree.h"
 #include "partitioning/partdesc.h"
@@ -227,21 +228,13 @@ FreeExecutorState(EState *estate)
 	MemoryContextDelete(estate->es_query_cxt);
 }
 
-/* ----------------
- *		CreateExprContext
- *
- *		Create a context for expression evaluation within an EState.
- *
- * An executor run may require multiple ExprContexts (we usually make one
- * for each Plan node, and a separate one for per-output-tuple processing
- * such as constraint checking).  Each ExprContext has its own "per-tuple"
- * memory context.
- *
- * Note we make no assumption about the caller's memory context.
- * ----------------
+/*
+ * Internal implementation for CreateExprContext() and CreateWorkExprContext()
+ * that allows control over the AllocSet parameters.
  */
-ExprContext *
-CreateExprContext(EState *estate)
+static ExprContext *
+CreateExprContextInternal(EState *estate, Size minContextSize,
+						  Size initBlockSize, Size maxBlockSize)
 {
 	ExprContext *econtext;
 	MemoryContext oldcontext;
@@ -264,7 +257,9 @@ CreateExprContext(EState *estate)
 	econtext->ecxt_per_tuple_memory =
 		AllocSetContextCreate(estate->es_query_cxt,
 							  "ExprContext",
-							  ALLOCSET_DEFAULT_SIZES);
+							  minContextSize,
+							  initBlockSize,
+							  maxBlockSize);
 
 	econtext->ecxt_param_exec_vals = estate->es_param_exec_vals;
 	econtext->ecxt_param_list_info = estate->es_param_list_info;
@@ -294,6 +289,52 @@ CreateExprContext(EState *estate)
 	return econtext;
 }
 
+/* ----------------
+ *		CreateExprContext
+ *
+ *		Create a context for expression evaluation within an EState.
+ *
+ * An executor run may require multiple ExprContexts (we usually make one
+ * for each Plan node, and a separate one for per-output-tuple processing
+ * such as constraint checking).  Each ExprContext has its own "per-tuple"
+ * memory context.
+ *
+ * Note we make no assumption about the caller's memory context.
+ * ----------------
+ */
+ExprContext *
+CreateExprContext(EState *estate)
+{
+	return CreateExprContextInternal(estate, ALLOCSET_DEFAULT_SIZES);
+}
+
+
+/* ----------------
+ *		CreateWorkExprContext
+ *
+ * Like CreateExprContext, but specifies the AllocSet sizes to be reasonable
+ * in proportion to work_mem. If the maximum block allocation size is too
+ * large, it's easy to skip right past work_mem with a single allocation.
+ * ----------------
+ */
+ExprContext *
+CreateWorkExprContext(EState *estate)
+{
+	Size minContextSize = ALLOCSET_DEFAULT_MINSIZE;
+	Size initBlockSize = ALLOCSET_DEFAULT_INITSIZE;
+	Size maxBlockSize = ALLOCSET_DEFAULT_MAXSIZE;
+
+	/* choose the maxBlockSize to be no larger than 1/16 of work_mem */
+	while (16 * maxBlockSize > work_mem * 1024L)
+		maxBlockSize >>= 1;
+
+	if (maxBlockSize < ALLOCSET_DEFAULT_INITSIZE)
+		maxBlockSize = ALLOCSET_DEFAULT_INITSIZE;
+
+	return CreateExprContextInternal(estate, minContextSize,
+									 initBlockSize, maxBlockSize);
+}
+
 /* ----------------
  *		CreateStandaloneExprContext
  *
diff --git a/src/backend/executor/nodeAgg.c b/src/backend/executor/nodeAgg.c
index 4e100e5755f..44587a84bae 100644
--- a/src/backend/executor/nodeAgg.c
+++ b/src/backend/executor/nodeAgg.c
@@ -3277,10 +3277,7 @@ ExecInitAgg(Agg *node, EState *estate, int eflags)
 	}
 
 	if (use_hashing)
-	{
-		ExecAssignExprContext(estate, &aggstate->ss.ps);
-		aggstate->hashcontext = aggstate->ss.ps.ps_ExprContext;
-	}
+		aggstate->hashcontext = CreateWorkExprContext(estate);
 
 	ExecAssignExprContext(estate, &aggstate->ss.ps);
 
diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h
index 94890512dc8..c7deeac662f 100644
--- a/src/include/executor/executor.h
+++ b/src/include/executor/executor.h
@@ -493,6 +493,7 @@ extern void end_tup_output(TupOutputState *tstate);
 extern EState *CreateExecutorState(void);
 extern void FreeExecutorState(EState *estate);
 extern ExprContext *CreateExprContext(EState *estate);
+extern ExprContext *CreateWorkExprContext(EState *estate);
 extern ExprContext *CreateStandaloneExprContext(void);
 extern void FreeExprContext(ExprContext *econtext, bool isCommit);
 extern void ReScanExprContext(ExprContext *econtext);