Fix memory accounting (#9739)

1. `localAllocated` is not set to 0 when resetting the memory context, which
makes the statistics data incorrect. 
2. Size is an unsigned integer, and `MEMORY_ACCOUNT_UPDATE_ALLOCATED`
uses the opposite value for subtraction. Since they have the same bit width
and the same bit results. We'd avoid using it in this way.
3. Fix memory accounting when setting a new parent. The accounting
parent must be one of the ancestors of the child. Re-setting the parent
might break the rule, see below:

If we set the new parent to the sibling or the ancestor of the old parent,
all children's accounting parent may have to change, because
the old accountingParent may not be the ancestor of the children's AllocSet.
We must have a loop to find it. For example:

root <- A <- B <- H <- C <- D
               \
                 <-E
root <- A <- B <- H
               \
                 <-E <- C <- D

We want to change the parent of C to E, and both C and D have the B
to be the accountingParent. After we set the new parent of C to E,
the accountingParent has to update. But if the new parent of C is B,
it still satisfies the above rule. So, we don't have to update the accountingParent.

The issue 1 is fixed in PR(https://github.com/greenplum-db/gpdb/pull/10106).
Reviewed-by: NNing Yu <nyu@pivotal.io>

src/backend/utils/mmgr/aset.c

@@ -239,7 +239,7 @@ IS_MEMORY_ACCOUNT(AllocSet set)
 }
 
 static inline void
-MEMORY_ACCOUNT_UPDATE_ALLOCATED(AllocSet set, Size newbytes)
+MEMORY_ACCOUNT_INC_ALLOCATED(AllocSet set, Size newbytes)
 {
 	AllocSet	parent = set->accountingParent;
 
@@ -248,6 +248,23 @@ MEMORY_ACCOUNT_UPDATE_ALLOCATED(AllocSet set, Size newbytes)
 	parent->currentAllocated += newbytes;
 	parent->peakAllocated = Max(parent->peakAllocated,
 							  parent->currentAllocated);
+
+	/* Make sure these values are not overflow */
+	Assert(set->localAllocated >= newbytes);
+	Assert(parent->currentAllocated >= set->localAllocated);
+}
+
+static inline void
+MEMORY_ACCOUNT_DEC_ALLOCATED(AllocSet set, Size newbytes)
+{
+	AllocSet	parent = set->accountingParent;
+
+	Assert(set->localAllocated >= newbytes);
+	Assert(parent->currentAllocated >= set->localAllocated);
+
+	set->localAllocated -= newbytes;
+	parent->currentAllocated -= newbytes;
+
 }
 
 /*
@@ -662,8 +679,12 @@ AllocSetReset(MemoryContext context)
 	AllocSetCheck(context);
 #endif
 
-	set->accountingParent->currentAllocated -= set->localAllocated;
-	set->localAllocated = 0;
+	/*
+	 * Make sure all children have been deleted,
+	 * or the accounting data is incorrect.
+	 */
+	Assert(context->firstchild == NULL);
+	MEMORY_ACCOUNT_DEC_ALLOCATED(set, set->localAllocated);
 
 	/* Clear chunk freelists */
 	MemSetAligned(set->freelist, 0, sizeof(set->freelist));
@@ -728,6 +749,9 @@ AllocSetDelete(MemoryContext context, MemoryContext parent)
 	AllocSetCheck(context);
 #endif
 
+	/* Make sure all children have been deleted */
+	Assert(context->firstchild == NULL);
+	MEMORY_ACCOUNT_DEC_ALLOCATED(set, set->localAllocated);
 	if (IS_MEMORY_ACCOUNT(set))
 	{
 		/* Roll up our peak value to the parent, before this context goes away. */
@@ -737,9 +761,6 @@ AllocSetDelete(MemoryContext context, MemoryContext parent)
 			Max(set->peakAllocated,
 				parentset->accountingParent->peakAllocated);
 	}
-	else
-		set->accountingParent->currentAllocated -= set->localAllocated;
-	set->localAllocated = 0;
 
 	/* Make it look empty, just in case... */
 	MemSetAligned(set->freelist, 0, sizeof(set->freelist));
@@ -848,7 +869,7 @@ AllocSetAlloc(MemoryContext context, Size size)
 		VALGRIND_MAKE_MEM_NOACCESS((char *) chunk + ALLOC_CHUNK_PUBLIC,
 						 chunk_size + ALLOC_CHUNKHDRSZ - ALLOC_CHUNK_PUBLIC);
 
-		MEMORY_ACCOUNT_UPDATE_ALLOCATED(set, chunk->size);
+		MEMORY_ACCOUNT_INC_ALLOCATED(set, chunk->size);
 
 		return AllocChunkGetPointer(chunk);
 	}
@@ -885,7 +906,7 @@ AllocSetAlloc(MemoryContext context, Size size)
 
 		AllocAllocInfo(set, chunk);
 
-		MEMORY_ACCOUNT_UPDATE_ALLOCATED(set, chunk->size);
+		MEMORY_ACCOUNT_INC_ALLOCATED(set, chunk->size);
 
 		return AllocChunkGetPointer(chunk);
 	}
@@ -1052,7 +1073,7 @@ AllocSetAlloc(MemoryContext context, Size size)
 #endif
 
 	AllocAllocInfo(set, chunk);
-	MEMORY_ACCOUNT_UPDATE_ALLOCATED(set, chunk->size);
+	MEMORY_ACCOUNT_INC_ALLOCATED(set, chunk->size);
 	return AllocChunkGetPointer(chunk);
 }
 
@@ -1068,7 +1089,7 @@ AllocSetFree(MemoryContext context, void *pointer)
 
 	AllocFreeInfo(set, chunk);
 
-	MEMORY_ACCOUNT_UPDATE_ALLOCATED(set, -chunk->size);
+	MEMORY_ACCOUNT_DEC_ALLOCATED(set, chunk->size);
 
 #ifdef USE_ASSERT_CHECKING
 	/*
@@ -1276,7 +1297,10 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 		/* Make any trailing alignment padding NOACCESS. */
 		VALGRIND_MAKE_MEM_NOACCESS((char *) pointer + size, chksize - size);
 
-		MEMORY_ACCOUNT_UPDATE_ALLOCATED(set, chksize - oldsize);
+		if (chksize > oldsize)
+			MEMORY_ACCOUNT_INC_ALLOCATED(set, chksize - oldsize);
+		else
+			MEMORY_ACCOUNT_DEC_ALLOCATED(set, oldsize - chksize);
 
 		return pointer;
 	}
@@ -1542,6 +1566,37 @@ AllocSetSetPeakUsage(MemoryContext context, Size nbytes)
 	return oldpeak;
 }
 
+void
+AllocSetTransferAccounting(MemoryContext context, MemoryContext new_parent)
+{
+	AllocSet set = (AllocSet)context;
+	AllocSet np = (AllocSet)new_parent;
+
+	if (set->accountingParent == set || set->accountingParent == np ||
+		(np && set->accountingParent == np->accountingParent))
+		return;
+
+	while (np && np != set->accountingParent)
+		np = (AllocSet)np->header.parent;
+
+	if (np == set->accountingParent)
+	{
+		/*
+		 * if set->accountingParent is the ancestor of the new parent,
+		 * the accoutingParent doesn't need to change.
+		 */
+	}
+	else
+	{
+		/* new_parent is NULL or new_parent is not the ancestor of context */
+		set->accountingParent->currentAllocated -= set->localAllocated;
+		set->accountingParent = set;
+		set->currentAllocated = set->localAllocated;
+		set->peakAllocated = set->localAllocated;
+	}
+
+}
+
 #ifdef MEMORY_CONTEXT_CHECKING
 
 /*

src/backend/utils/mmgr/mcxt.c

@@ -356,6 +356,9 @@ MemoryContextSetParent(MemoryContext context, MemoryContext new_parent)
 	if (new_parent == context->parent)
 		return;
 
+	/* update memory accounting */
+	AllocSetTransferAccounting(context, new_parent);
+
 	/* Delink from existing parent, if any */
 	if (context->parent)
 	{

src/include/utils/memutils.h

@@ -162,6 +162,10 @@ extern MemoryContext AllocSetContextCreate(MemoryContext parent,
 					  Size initBlockSize,
 					  Size maxBlockSize);
 
+/* this function should be only called by MemoryContextSetParent() */
+extern void AllocSetTransferAccounting(MemoryContext context,
+									   MemoryContext new_parent);
+
 /* mpool.c */
 typedef struct MPool MPool;
 extern MPool *mpool_create(MemoryContext parent,