From de570047993bd5fd65ad2bdf6b0acf5b8939bcb3 Mon Sep 17 00:00:00 2001
From: Tom Lane <tgl@sss.pgh.pa.us>
Date: Thu, 14 Mar 2019 18:36:26 -0400
Subject: [PATCH] Fix some oversights in commit 2455ab488.

The idea was to generate all the junk in a destroyable subcontext rather
than leaking it in the caller's context, but partition_bounds_create was
still being called in the caller's context, allowing plenty of scope for
leakage.  Also, get_rel_relkind() was still being called in the rel's
rd_pdcxt, creating a risk of session-lifespan memory wastage.

Simplify the logic a bit while at it.  Also, reduce rd_pdcxt to
ALLOCSET_SMALL_SIZES, since it seems likely to not usually be big.

Probably something like this needs to be back-patched into v11,
but for now let's get some buildfarm testing on this.

Discussion: https://postgr.es/m/15943.1552601288@sss.pgh.pa.us
---
 src/backend/partitioning/partdesc.c | 72 +++++++++++++----------------
 1 file changed, 33 insertions(+), 39 deletions(-)

diff --git a/src/backend/partitioning/partdesc.c b/src/backend/partitioning/partdesc.c
index 0593f044c2..5f349ad46a 100644
--- a/src/backend/partitioning/partdesc.c
+++ b/src/backend/partitioning/partdesc.c
@@ -67,8 +67,8 @@ RelationBuildPartitionDesc(Relation rel)
 				nparts;
 	PartitionKey key = RelationGetPartitionKey(rel);
 	MemoryContext oldcxt;
+	MemoryContext rbcontext;
 	int		   *mapping;
-	MemoryContext rbcontext = NULL;
 
 	/*
 	 * While building the partition descriptor, we create various temporary
@@ -187,56 +187,50 @@ RelationBuildPartitionDesc(Relation rel)
 	/* Now build the actual relcache partition descriptor */
 	rel->rd_pdcxt = AllocSetContextCreate(CacheMemoryContext,
 										  "partition descriptor",
-										  ALLOCSET_DEFAULT_SIZES);
+										  ALLOCSET_SMALL_SIZES);
 	MemoryContextCopyAndSetIdentifier(rel->rd_pdcxt,
 									  RelationGetRelationName(rel));
 
-	MemoryContextSwitchTo(rel->rd_pdcxt);
-	partdesc = (PartitionDescData *) palloc0(sizeof(PartitionDescData));
+	partdesc = (PartitionDescData *)
+		MemoryContextAllocZero(rel->rd_pdcxt, sizeof(PartitionDescData));
 	partdesc->nparts = nparts;
-	/* oids and boundinfo are allocated below. */
-
-	MemoryContextSwitchTo(oldcxt);
-
-	if (nparts == 0)
+	/* If there are no partitions, the rest of the partdesc can stay zero */
+	if (nparts > 0)
 	{
-		/* We can exit early in this case. */
-		rel->rd_partdesc = partdesc;
+		/* Create PartitionBoundInfo, using the temporary context. */
+		boundinfo = partition_bounds_create(boundspecs, nparts, key, &mapping);
 
-		/* Blow away the temporary context. */
-		MemoryContextDelete(rbcontext);
-		return;
-	}
+		/* Now copy all info into relcache's partdesc. */
+		MemoryContextSwitchTo(rel->rd_pdcxt);
+		partdesc->boundinfo = partition_bounds_copy(boundinfo, key);
+		partdesc->oids = (Oid *) palloc(nparts * sizeof(Oid));
+		partdesc->is_leaf = (bool *) palloc(nparts * sizeof(bool));
 
-	/* First create PartitionBoundInfo */
-	boundinfo = partition_bounds_create(boundspecs, nparts, key, &mapping);
-
-	/* Now copy boundinfo and oids into partdesc. */
-	oldcxt = MemoryContextSwitchTo(rel->rd_pdcxt);
-	partdesc->boundinfo = partition_bounds_copy(boundinfo, key);
-	partdesc->oids = (Oid *) palloc(partdesc->nparts * sizeof(Oid));
-	partdesc->is_leaf = (bool *) palloc(partdesc->nparts * sizeof(bool));
-
-	/*
-	 * Now assign OIDs from the original array into mapped indexes of the
-	 * result array.  The order of OIDs in the former is defined by the
-	 * catalog scan that retrieved them, whereas that in the latter is defined
-	 * by canonicalized representation of the partition bounds.
-	 */
-	for (i = 0; i < partdesc->nparts; i++)
-	{
-		int			index = mapping[i];
+		/*
+		 * Assign OIDs from the original array into mapped indexes of the
+		 * result array.  The order of OIDs in the former is defined by the
+		 * catalog scan that retrieved them, whereas that in the latter is
+		 * defined by canonicalized representation of the partition bounds.
+		 *
+		 * Also record leaf-ness of each partition.  For this we use
+		 * get_rel_relkind() which may leak memory, so be sure to run it in
+		 * the temporary context.
+		 */
+		MemoryContextSwitchTo(rbcontext);
+		for (i = 0; i < nparts; i++)
+		{
+			int			index = mapping[i];
 
-		partdesc->oids[index] = oids[i];
-		/* Record if the partition is a leaf partition */
-		partdesc->is_leaf[index] =
-			(get_rel_relkind(oids[i]) != RELKIND_PARTITIONED_TABLE);
+			partdesc->oids[index] = oids[i];
+			partdesc->is_leaf[index] =
+				(get_rel_relkind(oids[i]) != RELKIND_PARTITIONED_TABLE);
+		}
 	}
-	MemoryContextSwitchTo(oldcxt);
 
 	rel->rd_partdesc = partdesc;
 
-	/* Blow away the temporary context. */
+	/* Return to caller's context, and blow away the temporary context. */
+	MemoryContextSwitchTo(oldcxt);
 	MemoryContextDelete(rbcontext);
 }
 
-- 
2.40.0