1 files changed, 99 insertions, 19 deletions

diff --git a/src/backend/utils/sort/tuplesort.c b/src/backend/utils/sort/tuplesort.c
index d63c24d416c..38077ec9c4e 100644
--- a/src/backend/utils/sort/tuplesort.c
+++ b/src/backend/utils/sort/tuplesort.c
@@ -276,6 +276,7 @@ struct Tuplesortstate
 	SortTuple  *memtuples;		/* array of SortTuple structs */
 	int			memtupcount;	/* number of tuples currently present */
 	int			memtupsize;		/* allocated length of memtuples array */
+	bool		growmemtuples;	/* memtuples' growth still underway? */
 
 	/*
 	 * While building initial runs, this is the current output run number
@@ -570,6 +571,7 @@ tuplesort_begin_common(int workMem, bool randomAccess)
 
 	state->memtupcount = 0;
 	state->memtupsize = 1024;	/* initial guess */
+	state->growmemtuples = true;
 	state->memtuples = (SortTuple *) palloc(state->memtupsize * sizeof(SortTuple));
 
 	USEMEM(state, GetMemoryChunkSpace(state->memtuples));
@@ -955,37 +957,110 @@ tuplesort_end(Tuplesortstate *state)
 
 /*
  * Grow the memtuples[] array, if possible within our memory constraint.
- * Return TRUE if able to enlarge the array, FALSE if not.
+ * Return TRUE if we were able to enlarge the array, FALSE if not.
  *
- * At each increment we double the size of the array.  When we are short
- * on memory we could consider smaller increases, but because availMem
- * moves around with tuple addition/removal, this might result in thrashing.
- * Small increases in the array size are likely to be pretty inefficient.
+ * Normally, at each increment we double the size of the array.  When we no
+ * longer have enough memory to do that, we attempt one last, smaller increase
+ * (and then clear the growmemtuples flag so we don't try any more).  That
+ * allows us to use allowedMem as fully as possible; sticking to the pure
+ * doubling rule could result in almost half of allowedMem going unused.
+ * Because availMem moves around with tuple addition/removal, we need some
+ * rule to prevent making repeated small increases in memtupsize, which would
+ * just be useless thrashing.  The growmemtuples flag accomplishes that and
+ * also prevents useless recalculations in this function.
  */
 static bool
 grow_memtuples(Tuplesortstate *state)
 {
+	int			newmemtupsize;
+	int			memtupsize = state->memtupsize;
+	long		memNowUsed = state->allowedMem - state->availMem;
+
+	/* Forget it if we've already maxed out memtuples, per comment above */
+	if (!state->growmemtuples)
+		return false;
+
+	/* Select new value of memtupsize */
+	if (memNowUsed <= state->availMem)
+	{
+		/*
+		 * It is surely safe to double memtupsize if we've used no more than
+		 * half of allowedMem.
+		 *
+		 * Note: it might seem that we need to worry about memtupsize * 2
+		 * overflowing an int, but the MaxAllocSize clamp applied below
+		 * ensures the existing memtupsize can't be large enough for that.
+		 */
+		newmemtupsize = memtupsize * 2;
+	}
+	else
+	{
+		/*
+		 * This will be the last increment of memtupsize.  Abandon doubling
+		 * strategy and instead increase as much as we safely can.
+		 *
+		 * To stay within allowedMem, we can't increase memtupsize by more
+		 * than availMem / sizeof(SortTuple) elements.	In practice, we want
+		 * to increase it by considerably less, because we need to leave some
+		 * space for the tuples to which the new array slots will refer.  We
+		 * assume the new tuples will be about the same size as the tuples
+		 * we've already seen, and thus we can extrapolate from the space
+		 * consumption so far to estimate an appropriate new size for the
+		 * memtuples array.  The optimal value might be higher or lower than
+		 * this estimate, but it's hard to know that in advance.
+		 *
+		 * This calculation is safe against enlarging the array so much that
+		 * LACKMEM becomes true, because the memory currently used includes
+		 * the present array; thus, there would be enough allowedMem for the
+		 * new array elements even if no other memory were currently used.
+		 *
+		 * We do the arithmetic in float8, because otherwise the product of
+		 * memtupsize and allowedMem could overflow.  (A little algebra shows
+		 * that grow_ratio must be less than 2 here, so we are not risking
+		 * integer overflow this way.)	Any inaccuracy in the result should be
+		 * insignificant; but even if we computed a completely insane result,
+		 * the checks below will prevent anything really bad from happening.
+		 */
+		double		grow_ratio;
+
+		grow_ratio = (double) state->allowedMem / (double) memNowUsed;
+		newmemtupsize = (int) (memtupsize * grow_ratio);
+
+		/* We won't make any further enlargement attempts */
+		state->growmemtuples = false;
+	}
+
+	/* Must enlarge array by at least one element, else report failure */
+	if (newmemtupsize <= memtupsize)
+		goto noalloc;
+
 	/*
-	 * We need to be sure that we do not cause LACKMEM to become true, else
-	 * the space management algorithm will go nuts.  We assume here that the
-	 * memory chunk overhead associated with the memtuples array is constant
-	 * and so there will be no unexpected addition to what we ask for.	(The
-	 * minimum array size established in tuplesort_begin_common is large
-	 * enough to force palloc to treat it as a separate chunk, so this
-	 * assumption should be good.  But let's check it.)
+	 * On a 64-bit machine, allowedMem could be more than MaxAllocSize.  Clamp
+	 * to ensure our request won't be rejected by palloc.
 	 */
-	if (state->availMem <= (long) (state->memtupsize * sizeof(SortTuple)))
-		return false;
+	if ((Size) newmemtupsize >= MaxAllocSize / sizeof(SortTuple))
+	{
+		newmemtupsize = (int) (MaxAllocSize / sizeof(SortTuple));
+		state->growmemtuples = false;	/* can't grow any more */
+	}
 
 	/*
-	 * On a 64-bit machine, allowedMem could be high enough to get us into
-	 * trouble with MaxAllocSize, too.
+	 * We need to be sure that we do not cause LACKMEM to become true, else
+	 * the space management algorithm will go nuts.  The code above should
+	 * never generate a dangerous request, but to be safe, check explicitly
+	 * that the array growth fits within availMem.	(We could still cause
+	 * LACKMEM if the memory chunk overhead associated with the memtuples
+	 * array were to increase.	That shouldn't happen with any sane value of
+	 * allowedMem, because at any array size large enough to risk LACKMEM,
+	 * palloc would be treating both old and new arrays as separate chunks.
+	 * But we'll check LACKMEM explicitly below just in case.)
 	 */
-	if ((Size) (state->memtupsize * 2) >= MaxAllocSize / sizeof(SortTuple))
-		return false;
+	if (state->availMem < (long) ((newmemtupsize - memtupsize) * sizeof(SortTuple)))
+		goto noalloc;
 
+	/* OK, do it */
 	FREEMEM(state, GetMemoryChunkSpace(state->memtuples));
-	state->memtupsize *= 2;
+	state->memtupsize = newmemtupsize;
 	state->memtuples = (SortTuple *)
 		repalloc(state->memtuples,
 				 state->memtupsize * sizeof(SortTuple));
@@ -993,6 +1068,11 @@ grow_memtuples(Tuplesortstate *state)
 	if (LACKMEM(state))
 		elog(ERROR, "unexpected out-of-memory situation during sort");
 	return true;
+
+noalloc:
+	/* If for any reason we didn't realloc, shut off future attempts */
+	state->growmemtuples = false;
+	return false;
 }
 
 /*