diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index 743455e..d556f0b 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -242,6 +242,8 @@ typedef struct AllocChunkData
 #define AllocChunkGetPointer(chk)	\
 					((AllocPointer)(((char *)(chk)) + ALLOC_CHUNKHDRSZ))
 
+static void update_allocation(MemoryContext context, Size size);
+
 /*
  * These functions implement the MemoryContext API for AllocSet contexts.
  */
@@ -430,6 +432,9 @@ randomize_mem(char *ptr, size_t size)
  * minContextSize: minimum context size
  * initBlockSize: initial allocation block size
  * maxBlockSize: maximum allocation block size
+ *
+ * The flag determining whether this context tracks memory usage is inherited
+ * from the parent context.
  */
 MemoryContext
 AllocSetContextCreate(MemoryContext parent,
@@ -438,6 +443,26 @@ AllocSetContextCreate(MemoryContext parent,
 					  Size initBlockSize,
 					  Size maxBlockSize)
 {
+	return AllocSetContextCreateTracked(
+		parent, name, minContextSize, initBlockSize, maxBlockSize,
+		false);
+}
+
+/*
+ * AllocSetContextCreateTracked
+ *		Create a new AllocSet context.
+ *
+ * Implementation for AllocSetContextCreate, but also allows the caller to
+ * specify whether memory usage should be tracked or not.
+ */
+MemoryContext
+AllocSetContextCreateTracked(MemoryContext parent,
+							 const char *name,
+							 Size minContextSize,
+							 Size initBlockSize,
+							 Size maxBlockSize,
+							 bool track_mem)
+{
 	AllocSet	context;
 
 	/* Do the type-independent part of context creation */
@@ -445,7 +470,8 @@ AllocSetContextCreate(MemoryContext parent,
 											 sizeof(AllocSetContext),
 											 &AllocSetMethods,
 											 parent,
-											 name);
+											 name,
+											 track_mem);
 
 	/*
 	 * Make sure alloc parameters are reasonable, and save them.
@@ -500,6 +526,9 @@ AllocSetContextCreate(MemoryContext parent,
 					 errdetail("Failed while creating memory context \"%s\".",
 							   name)));
 		}
+
+		update_allocation((MemoryContext) context, blksize);
+
 		block->aset = context;
 		block->freeptr = ((char *) block) + ALLOC_BLOCKHDRSZ;
 		block->endptr = ((char *) block) + blksize;
@@ -590,6 +619,7 @@ AllocSetReset(MemoryContext context)
 		else
 		{
 			/* Normal case, release the block */
+			update_allocation(context, -(block->endptr - ((char*) block)));
 #ifdef CLOBBER_FREED_MEMORY
 			wipe_mem(block, block->freeptr - ((char *) block));
 #endif
@@ -616,6 +646,8 @@ AllocSetDelete(MemoryContext context)
 	AllocSet	set = (AllocSet) context;
 	AllocBlock	block = set->blocks;
 
+	MemoryAccounting accounting;
+
 	AssertArg(AllocSetIsValid(set));
 
 #ifdef MEMORY_CONTEXT_CHECKING
@@ -623,6 +655,17 @@ AllocSetDelete(MemoryContext context)
 	AllocSetCheck(context);
 #endif
 
+	if (context->accounting != NULL) {
+
+		accounting = context->accounting->parent;
+
+		while (accounting != NULL)
+		{
+			accounting->total_allocated -= context->accounting->total_allocated;
+			accounting = accounting->parent;
+		}
+	}
+
 	/* Make it look empty, just in case... */
 	MemSetAligned(set->freelist, 0, sizeof(set->freelist));
 	set->blocks = NULL;
@@ -678,6 +721,9 @@ AllocSetAlloc(MemoryContext context, Size size)
 					 errmsg("out of memory"),
 					 errdetail("Failed on request of size %zu.", size)));
 		}
+
+		update_allocation(context, blksize);
+
 		block->aset = set;
 		block->freeptr = block->endptr = ((char *) block) + blksize;
 
@@ -873,6 +919,8 @@ AllocSetAlloc(MemoryContext context, Size size)
 					 errdetail("Failed on request of size %zu.", size)));
 		}
 
+		update_allocation(context, blksize);
+
 		block->aset = set;
 		block->freeptr = ((char *) block) + ALLOC_BLOCKHDRSZ;
 		block->endptr = ((char *) block) + blksize;
@@ -976,6 +1024,7 @@ AllocSetFree(MemoryContext context, void *pointer)
 			set->blocks = block->next;
 		else
 			prevblock->next = block->next;
+		update_allocation(context, -(block->endptr - ((char*) block)));
 #ifdef CLOBBER_FREED_MEMORY
 		wipe_mem(block, block->freeptr - ((char *) block));
 #endif
@@ -1088,6 +1137,7 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 		AllocBlock	prevblock = NULL;
 		Size		chksize;
 		Size		blksize;
+		Size		oldblksize;
 
 		while (block != NULL)
 		{
@@ -1105,6 +1155,8 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 		/* Do the realloc */
 		chksize = MAXALIGN(size);
 		blksize = chksize + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
+		oldblksize = block->endptr - ((char *)block);
+
 		block = (AllocBlock) realloc(block, blksize);
 		if (block == NULL)
 		{
@@ -1114,6 +1166,7 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 					 errmsg("out of memory"),
 					 errdetail("Failed on request of size %zu.", size)));
 		}
+		update_allocation(context, blksize - oldblksize);
 		block->freeptr = block->endptr = ((char *) block) + blksize;
 
 		/* Update pointers since block has likely been moved */
@@ -1277,6 +1330,45 @@ AllocSetStats(MemoryContext context, int level)
 }
 
 
+/*
+ * update_allocation
+ *
+ * Track newly-allocated or newly-freed memory (freed memory should be
+ * negative).
+ */
+static void
+update_allocation(MemoryContext context, Size size)
+{
+
+	MemoryAccounting accounting = context->accounting;
+
+	if (accounting == NULL)
+		return;
+
+	/*
+	 * Update self_allocated only if this accounting info is specific
+	 * for this context (i.e. if tracking was requested for the context).
+	 */
+	if (context->track_mem)
+		accounting->self_allocated += size;
+
+	/*
+	 * Update total_allocated for all contexts up the accounting tree
+	 * (including this one).
+	 */
+	while (accounting != NULL) {
+
+		accounting->total_allocated += size;
+
+		Assert(accounting->self_allocated >= 0);
+		Assert(accounting->total_allocated >= accounting->self_allocated);
+
+		accounting = accounting->parent;
+
+	}
+
+}
+
 #ifdef MEMORY_CONTEXT_CHECKING
 
 /*
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 4185a03..a70b296 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -202,7 +202,12 @@ MemoryContextDelete(MemoryContext context)
 	 */
 	MemoryContextSetParent(context, NULL);
 
+	/* pass the parent in case it's needed, however */
 	(*context->methods->delete_context) (context);
+
+	if (context->track_mem)
+		pfree(context->accounting);
+
 	VALGRIND_DESTROY_MEMPOOL(context);
 	pfree(context);
 }
@@ -324,6 +329,26 @@ MemoryContextAllowInCriticalSection(MemoryContext context, bool allow)
 }
 
 /*
+ * MemoryContextGetAllocated
+ *
+ * Return memory allocated by the system to this context. If total is true,
+ * include child contexts. Context must have track_mem set.
+ */
+Size
+MemoryContextGetAllocated(MemoryContext context, bool total)
+{
+	Assert(context->track_mem);
+
+	if (! context->track_mem)
+		return 0;
+
+	if (total)
+		return context->accounting->total_allocated;
+	else
+		return context->accounting->self_allocated;
+}
+
+/*
  * GetMemoryChunkSpace
  *		Given a currently-allocated chunk, determine the total space
  *		it occupies (including all memory-allocation overhead).
@@ -546,7 +571,8 @@ MemoryContext
 MemoryContextCreate(NodeTag tag, Size size,
 					MemoryContextMethods *methods,
 					MemoryContext parent,
-					const char *name)
+					const char *name,
+					bool track_mem)
 {
 	MemoryContext node;
 	Size		needed = size + strlen(name) + 1;
@@ -576,6 +602,8 @@ MemoryContextCreate(NodeTag tag, Size size,
 	node->firstchild = NULL;
 	node->nextchild = NULL;
 	node->isReset = true;
+	node->track_mem = track_mem;
+	node->accounting = NULL;
 	node->name = ((char *) node) + size;
 	strcpy(node->name, name);
 
@@ -596,6 +624,24 @@ MemoryContextCreate(NodeTag tag, Size size,
 #endif
 	}
 
+	/*
+	 * If accounting was requested for this context, create the struct
+	 * and link it to the accounting from parent context. Otherwise just
+	 * copy the accounting reference from parent.
+	 *
+	 * If both cases, accounting in parent may be NULL, which means
+	 * we don't need to update accounting in the upper memory contexts.
+	 */
+	if (track_mem)
+	{
+		node->accounting = (MemoryAccounting)MemoryContextAlloc(TopMemoryContext,
+												sizeof(MemoryAccountingData));
+		if (parent)
+			node->accounting->parent = parent->accounting;
+	} else if (parent) {
+		node->accounting = parent->accounting;
+	}
+
 	VALGRIND_CREATE_MEMPOOL(node, 0, false);
 
 	/* Return to type-specific creation routine to finish up */
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index ad77509..a6d9f8c 100644
--- a/src/include/nodes/memnodes.h
+++ b/src/include/nodes/memnodes.h
@@ -50,6 +50,18 @@ typedef struct MemoryContextMethods
 #endif
 } MemoryContextMethods;
 
+typedef struct MemoryAccountingData {
+
+	Size	total_allocated; /* including child contexts */
+	Size	self_allocated;  /* not including child contexts */
+
+	/* parent accounting (not parent context) */
+	struct MemoryAccountingData * parent;
+
+} MemoryAccountingData;
+
+typedef MemoryAccountingData * MemoryAccounting;
+
 
 typedef struct MemoryContextData
 {
@@ -60,6 +72,8 @@ typedef struct MemoryContextData
 	MemoryContext nextchild;	/* next child of same parent */
 	char	   *name;			/* context name (just for debugging) */
 	bool		isReset;		/* T = no space alloced since last reset */
+	bool		track_mem;		/* whether to track memory usage */
+	MemoryAccounting	accounting;
 #ifdef USE_ASSERT_CHECKING
 	bool		allowInCritSection;	/* allow palloc in critical section */
 #endif
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index 2fede86..b82f923 100644
--- a/src/include/utils/memutils.h
+++ b/src/include/utils/memutils.h
@@ -96,6 +96,7 @@ extern void MemoryContextDeleteChildren(MemoryContext context);
 extern void MemoryContextResetAndDeleteChildren(MemoryContext context);
 extern void MemoryContextSetParent(MemoryContext context,
 					   MemoryContext new_parent);
+extern Size MemoryContextGetAllocated(MemoryContext context, bool total);
 extern Size GetMemoryChunkSpace(void *pointer);
 extern MemoryContext GetMemoryChunkContext(void *pointer);
 extern MemoryContext MemoryContextGetParent(MemoryContext context);
@@ -117,7 +118,8 @@ extern bool MemoryContextContains(MemoryContext context, void *pointer);
 extern MemoryContext MemoryContextCreate(NodeTag tag, Size size,
 					MemoryContextMethods *methods,
 					MemoryContext parent,
-					const char *name);
+					const char *name,
+					bool track_mem);
 
 
 /*
@@ -130,6 +132,12 @@ extern MemoryContext AllocSetContextCreate(MemoryContext parent,
 					  Size minContextSize,
 					  Size initBlockSize,
 					  Size maxBlockSize);
+extern MemoryContext AllocSetContextCreateTracked(MemoryContext parent,
+					  const char *name,
+					  Size minContextSize,
+					  Size initBlockSize,
+					  Size maxBlockSize,
+					  bool track_mem);
 
 /*
  * Recommended default alloc parameters, suitable for "ordinary" contexts