Move memory management away from writetup() and tuplesort_put*()

This commit puts some generic work away from sort-variant-specific function.
In particular, tuplesort_put*() now doesn't need to decrease available memory
and switch to sort context before calling puttuple_common().  writetup()
doesn't need to free SortTuple.tuple and increase available memory.

Discussion: https://postgr.es/m/CAPpHfdvjix0Ahx-H3Jp1M2R%2B_74P-zKnGGygx4OWr%3DbUQ8BNdw%40mail.gmail.com
Author: Alexander Korotkov
Reviewed-by: Pavel Borisov, Maxim Orlov, Matthias van de Meent
Reviewed-by: Andres Freund, John Naylor

src/backend/utils/sort/tuplesort.c

@@ -288,11 +288,7 @@ struct Tuplesortstate
 
 	/*
 	 * Function to write a stored tuple onto tape.  The representation of the
-	 * tuple on tape need not be the same as it is in memory; requirements on
+	 * tuple on tape need not be the same as it is in memory.
-	 * the tape representation are given below.  Unless the slab allocator is
-	 * used, after writing the tuple, pfree() the out-of-line data (not the
-	 * SortTuple struct!), and increase state->availMem by the amount of
-	 * memory space thereby released.
 	 */
 	void		(*writetup) (Tuplesortstate *state, LogicalTape *tape,
 							 SortTuple *stup);
@@ -549,7 +545,7 @@ struct Sharedsort
 
 #define REMOVEABBREV(state,stup,count)	((*(state)->removeabbrev) (state, stup, count))
 #define COMPARETUP(state,a,b)	((*(state)->comparetup) (a, b, state))
-#define WRITETUP(state,tape,stup)	((*(state)->writetup) (state, tape, stup))
+#define WRITETUP(state,tape,stup)	(writetuple(state, tape, stup))
 #define READTUP(state,stup,tape,len) ((*(state)->readtup) (state, stup, tape, len))
 #define LACKMEM(state)		((state)->availMem < 0 && !(state)->slabAllocatorUsed)
 #define USEMEM(state,amt)	((state)->availMem -= (amt))
@@ -618,6 +614,8 @@ static Tuplesortstate *tuplesort_begin_common(int workMem,
 static void tuplesort_begin_batch(Tuplesortstate *state);
 static void puttuple_common(Tuplesortstate *state, SortTuple *tuple,
 							bool useAbbrev);
+static void writetuple(Tuplesortstate *state, LogicalTape *tape,
+					   SortTuple *stup);
 static bool consider_abort_common(Tuplesortstate *state);
 static void inittapes(Tuplesortstate *state, bool mergeruns);
 static void inittapestate(Tuplesortstate *state, int maxTapes);
@@ -1848,7 +1846,6 @@ tuplesort_puttupleslot(Tuplesortstate *state, TupleTableSlot *slot)
 	/* copy the tuple into sort storage */
 	tuple = ExecCopySlotMinimalTuple(slot);
 	stup.tuple = (void *) tuple;
-	USEMEM(state, GetMemoryChunkSpace(tuple));
 	/* set up first-column key value */
 	htup.t_len = tuple->t_len + MINIMAL_TUPLE_OFFSET;
 	htup.t_data = (HeapTupleHeader) ((char *) tuple - MINIMAL_TUPLE_OFFSET);
@@ -1857,8 +1854,6 @@ tuplesort_puttupleslot(Tuplesortstate *state, TupleTableSlot *slot)
 							   state->tupDesc,
 							   &stup.isnull1);
 
-	MemoryContextSwitchTo(state->sortcontext);
-
 	puttuple_common(state, &stup,
 					state->sortKeys->abbrev_converter && !stup.isnull1);
 
@@ -1879,9 +1874,6 @@ tuplesort_putheaptuple(Tuplesortstate *state, HeapTuple tup)
 	/* copy the tuple into sort storage */
 	tup = heap_copytuple(tup);
 	stup.tuple = (void *) tup;
-	USEMEM(state, GetMemoryChunkSpace(tup));
-
-	MemoryContextSwitchTo(state->sortcontext);
 
 	/*
 	 * set up first-column key value, and potentially abbreviate, if it's a
@@ -1910,7 +1902,6 @@ tuplesort_putindextuplevalues(Tuplesortstate *state, Relation rel,
 							  ItemPointer self, Datum *values,
 							  bool *isnull)
 {
-	MemoryContext oldcontext;
 	SortTuple	stup;
 	IndexTuple	tuple;
 
@@ -1918,19 +1909,14 @@ tuplesort_putindextuplevalues(Tuplesortstate *state, Relation rel,
 										  isnull, state->tuplecontext);
 	tuple = ((IndexTuple) stup.tuple);
 	tuple->t_tid = *self;
-	USEMEM(state, GetMemoryChunkSpace(stup.tuple));
 	/* set up first-column key value */
 	stup.datum1 = index_getattr(tuple,
 								1,
 								RelationGetDescr(state->indexRel),
 								&stup.isnull1);
 
-	oldcontext = MemoryContextSwitchTo(state->sortcontext);
-
 	puttuple_common(state, &stup,
 					state->sortKeys && state->sortKeys->abbrev_converter && !stup.isnull1);
-
-	MemoryContextSwitchTo(oldcontext);
 }
 
 /*
@@ -1965,15 +1951,12 @@ tuplesort_putdatum(Tuplesortstate *state, Datum val, bool isNull)
 		stup.datum1 = !isNull ? val : (Datum) 0;
 		stup.isnull1 = isNull;
 		stup.tuple = NULL;		/* no separate storage */
-		MemoryContextSwitchTo(state->sortcontext);
 	}
 	else
 	{
 		stup.isnull1 = false;
 		stup.datum1 = datumCopy(val, false, state->datumTypeLen);
 		stup.tuple = DatumGetPointer(stup.datum1);
-		USEMEM(state, GetMemoryChunkSpace(stup.tuple));
-		MemoryContextSwitchTo(state->sortcontext);
 	}
 
 	puttuple_common(state, &stup,
@@ -1988,8 +1971,14 @@ tuplesort_putdatum(Tuplesortstate *state, Datum val, bool isNull)
 static void
 puttuple_common(Tuplesortstate *state, SortTuple *tuple, bool useAbbrev)
 {
+	MemoryContext oldcontext = MemoryContextSwitchTo(state->sortcontext);
+
 	Assert(!LEADER(state));
 
+	/* Count the size of the out-of-line data */
+	if (tuple->tuple != NULL)
+		USEMEM(state, GetMemoryChunkSpace(tuple->tuple));
+
 	if (!useAbbrev)
 	{
 		/*
@@ -2062,6 +2051,7 @@ puttuple_common(Tuplesortstate *state, SortTuple *tuple, bool useAbbrev)
 						 pg_rusage_show(&state->ru_start));
 #endif
 				make_bounded_heap(state);
+				MemoryContextSwitchTo(oldcontext);
 				return;
 			}
 
@@ -2069,7 +2059,10 @@ puttuple_common(Tuplesortstate *state, SortTuple *tuple, bool useAbbrev)
 			 * Done if we still fit in available memory and have array slots.
 			 */
 			if (state->memtupcount < state->memtupsize && !LACKMEM(state))
+			{
+				MemoryContextSwitchTo(oldcontext);
 				return;
+			}
 
 			/*
 			 * Nope; time to switch to tape-based operation.
@@ -2123,6 +2116,25 @@ puttuple_common(Tuplesortstate *state, SortTuple *tuple, bool useAbbrev)
 			elog(ERROR, "invalid tuplesort state");
 			break;
 	}
+	MemoryContextSwitchTo(oldcontext);
+}
+
+/*
+ * Write a stored tuple onto tape.tuple.  Unless the slab allocator is
+ * used, after writing the tuple, pfree() the out-of-line data (not the
+ * SortTuple struct!), and increase state->availMem by the amount of
+ * memory space thereby released.
+ */
+static void
+writetuple(Tuplesortstate *state, LogicalTape *tape, SortTuple *stup)
+{
+	state->writetup(state, tape, stup);
+
+	if (!state->slabAllocatorUsed && stup->tuple)
+	{
+		FREEMEM(state, GetMemoryChunkSpace(stup->tuple));
+		pfree(stup->tuple);
+	}
 }
 
 static bool
@@ -3960,12 +3972,6 @@ writetup_heap(Tuplesortstate *state, LogicalTape *tape, SortTuple *stup)
 	if (state->sortopt & TUPLESORT_RANDOMACCESS)	/* need trailing length
 													 * word? */
 		LogicalTapeWrite(tape, (void *) &tuplen, sizeof(tuplen));
-
-	if (!state->slabAllocatorUsed)
-	{
-		FREEMEM(state, GetMemoryChunkSpace(tuple));
-		heap_free_minimal_tuple(tuple);
-	}
 }
 
 static void
@@ -4141,12 +4147,6 @@ writetup_cluster(Tuplesortstate *state, LogicalTape *tape, SortTuple *stup)
 	if (state->sortopt & TUPLESORT_RANDOMACCESS)	/* need trailing length
 													 * word? */
 		LogicalTapeWrite(tape, &tuplen, sizeof(tuplen));
-
-	if (!state->slabAllocatorUsed)
-	{
-		FREEMEM(state, GetMemoryChunkSpace(tuple));
-		heap_freetuple(tuple);
-	}
 }
 
 static void
@@ -4403,12 +4403,6 @@ writetup_index(Tuplesortstate *state, LogicalTape *tape, SortTuple *stup)
 	if (state->sortopt & TUPLESORT_RANDOMACCESS)	/* need trailing length
 													 * word? */
 		LogicalTapeWrite(tape, (void *) &tuplen, sizeof(tuplen));
-
-	if (!state->slabAllocatorUsed)
-	{
-		FREEMEM(state, GetMemoryChunkSpace(tuple));
-		pfree(tuple);
-	}
 }
 
 static void
@@ -4495,12 +4489,6 @@ writetup_datum(Tuplesortstate *state, LogicalTape *tape, SortTuple *stup)
 	if (state->sortopt & TUPLESORT_RANDOMACCESS)	/* need trailing length
 													 * word? */
 		LogicalTapeWrite(tape, (void *) &writtenlen, sizeof(writtenlen));
-
-	if (!state->slabAllocatorUsed && stup->tuple)
-	{
-		FREEMEM(state, GetMemoryChunkSpace(stup->tuple));
-		pfree(stup->tuple);
-	}
 }
 
 static void