From 850af5f9f7b1b85f98a54f5b4a757bbd00df77ec Mon Sep 17 00:00:00 2001
From: Peter Geoghegan <pg@bowt.ie>
Date: Sun, 4 Sep 2016 16:05:49 -0700
Subject: [PATCH] Defend against faulty batch memtuples calculation

Add a new test to batchmemtuples() that must be passed in order to
proceed with actually growing memtuples.  This avoids problems with
spurious target allocation sizes in corner cases with low memory.

Backpatch to 9.6, where the tuplesort batch memory optimization was
added.
---
 src/backend/utils/sort/tuplesort.c | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/backend/utils/sort/tuplesort.c b/src/backend/utils/sort/tuplesort.c
index c8fbcf8..3478277 100644
--- a/src/backend/utils/sort/tuplesort.c
+++ b/src/backend/utils/sort/tuplesort.c
@@ -2868,6 +2868,7 @@ batchmemtuples(Tuplesortstate *state)
 
 	/* For simplicity, assume no memtuples are actually currently counted */
 	Assert(state->memtupcount == 0);
+	Assert(state->activeTapes > 0);
 
 	/*
 	 * Refund STANDARDCHUNKHEADERSIZE per tuple.
@@ -2880,6 +2881,20 @@ batchmemtuples(Tuplesortstate *state)
 	availMemLessRefund = state->availMem - refund;
 
 	/*
+	 * We need to be sure that we do not cause LACKMEM to become true, else a
+	 * negative number of bytes could be calculated as batch allocation size,
+	 * causing havoc (a negative availMemLessRefund cannot be accepted).  Note
+	 * that we are never reliant on a tape's batch allocation being large
+	 * enough to fit even a single tuple; the use of the constant
+	 * ALLOCSET_DEFAULT_INITSIZE here is not critical (a threshold like this
+	 * avoids a useless repalloc that only grows memtuples by a tiny amount,
+	 * which seems worth avoiding here in passing).
+	 */
+	if (availMemLessRefund <=
+		(int64) state->activeTapes * ALLOCSET_DEFAULT_INITSIZE)
+	   return;
+
+	/*
 	 * To establish balanced memory use after refunding palloc overhead,
 	 * temporarily have our accounting indicate that we've allocated all
 	 * memory we're allowed to less that refund, and call grow_memtuples() to
@@ -2889,8 +2904,11 @@ batchmemtuples(Tuplesortstate *state)
 	USEMEM(state, availMemLessRefund);
 	(void) grow_memtuples(state);
 	/* Should not matter, but be tidy */
+	state->growmemtuples = false;
+	/* availMem must stay accurate for spacePerTape calculation */
 	FREEMEM(state, availMemLessRefund);
-	state->growmemtuples = false;
+	if (LACKMEM(state))
+		elog(ERROR, "unexpected out-of-memory situation in tuplesort");
 
 #ifdef TRACE_SORT
 	if (trace_sort)
-- 
2.7.4

