Adjust tuplestore stats API

1eff8279d added an API to tuplestore.c to allow callers to obtain
storage telemetry data.  That API wasn't quite good enough for callers
that perform tuplestore_clear() as the telemetry functions only
accounted for the current state of the tuplestore, not the maximums
before tuplestore_clear() was called.

There's a pending patch that would like to add tuplestore telemetry
output to EXPLAIN ANALYZE for WindowAgg.  That node type uses
tuplestore_clear() before moving to the next window partition and we
want to show the maximum space used, not the space used for the final
partition.

Reviewed-by: Tatsuo Ishii, Ashutosh Bapat
Discussion: https://postgres/m/CAApHDvoY8cibGcicLV0fNh=9JVx9PANcWvhkdjBnDCc9Quqytg@mail.gmail.com

3 files changed, 48 insertions, 38 deletions

diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index 14cd36c87c5..2819e479f82 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -3350,8 +3350,9 @@ static void
 show_material_info(MaterialState *mstate, ExplainState *es)
 {
 	Tuplestorestate *tupstore = mstate->tuplestorestate;
-	const char *storageType;
-	int64		spaceUsedKB;
+	char	   *maxStorageType;
+	int64		maxSpaceUsed,
+				maxSpaceUsedKB;
 
 	/*
 	 * Nothing to show if ANALYZE option wasn't used or if execution didn't
@@ -3360,21 +3361,21 @@ show_material_info(MaterialState *mstate, ExplainState *es)
 	if (!es->analyze || tupstore == NULL)
 		return;
 
-	storageType = tuplestore_storage_type_name(tupstore);
-	spaceUsedKB = BYTES_TO_KILOBYTES(tuplestore_space_used(tupstore));
+	tuplestore_get_stats(tupstore, &maxStorageType, &maxSpaceUsed);
+	maxSpaceUsedKB = BYTES_TO_KILOBYTES(maxSpaceUsed);
 
 	if (es->format != EXPLAIN_FORMAT_TEXT)
 	{
-		ExplainPropertyText("Storage", storageType, es);
-		ExplainPropertyInteger("Maximum Storage", "kB", spaceUsedKB, es);
+		ExplainPropertyText("Storage", maxStorageType, es);
+		ExplainPropertyInteger("Maximum Storage", "kB", maxSpaceUsedKB, es);
 	}
 	else
 	{
 		ExplainIndentText(es);
 		appendStringInfo(es->str,
 						 "Storage: %s  Maximum Storage: " INT64_FORMAT "kB\n",
-						 storageType,
-						 spaceUsedKB);
+						 maxStorageType,
+						 maxSpaceUsedKB);
 	}
 }
 
diff --git a/src/backend/utils/sort/tuplestore.c b/src/backend/utils/sort/tuplestore.c
index 444c8e25c21..a720d702001 100644
--- a/src/backend/utils/sort/tuplestore.c
+++ b/src/backend/utils/sort/tuplestore.c
@@ -107,9 +107,10 @@ struct Tuplestorestate
 	bool		backward;		/* store extra length words in file? */
 	bool		interXact;		/* keep open through transactions? */
 	bool		truncated;		/* tuplestore_trim has removed tuples? */
+	bool		usedDisk;		/* used by tuplestore_get_stats() */
+	int64		maxSpace;		/* used by tuplestore_get_stats() */
 	int64		availMem;		/* remaining memory available, in bytes */
 	int64		allowedMem;		/* total memory allowed, in bytes */
-	int64		maxSpace;		/* maximum space used in memory */
 	int64		tuples;			/* number of tuples added */
 	BufFile    *myfile;			/* underlying file, or NULL if none */
 	MemoryContext context;		/* memory context for holding tuples */
@@ -262,9 +263,10 @@ tuplestore_begin_common(int eflags, bool interXact, int maxKBytes)
 	state->eflags = eflags;
 	state->interXact = interXact;
 	state->truncated = false;
+	state->usedDisk = false;
+	state->maxSpace = 0;
 	state->allowedMem = maxKBytes * 1024L;
 	state->availMem = state->allowedMem;
-	state->maxSpace = 0;
 	state->myfile = NULL;
 
 	/*
@@ -870,6 +872,14 @@ tuplestore_puttuple_common(Tuplestorestate *state, void *tuple)
 			 * though callers might drop the requirement.
 			 */
 			state->backward = (state->eflags & EXEC_FLAG_BACKWARD) != 0;
+
+			/*
+			 * Update the maximum space used before dumping the tuples.  It's
+			 * possible that more space will be used by the tuples in memory
+			 * than the space that will be used on disk.
+			 */
+			tuplestore_updatemax(state);
+
 			state->status = TSS_WRITEFILE;
 			dumptuples(state);
 			break;
@@ -1444,7 +1454,7 @@ tuplestore_trim(Tuplestorestate *state)
 	Assert(nremove >= state->memtupdeleted);
 	Assert(nremove <= state->memtupcount);
 
-	/* before freeing any memory, update maxSpace */
+	/* before freeing any memory, update the statistics */
 	tuplestore_updatemax(state);
 
 	/* Release no-longer-needed tuples */
@@ -1491,7 +1501,8 @@ tuplestore_trim(Tuplestorestate *state)
 
 /*
  * tuplestore_updatemax
- *		Update maxSpace field
+ *		Update the maximum space used by this tuplestore and the method used
+ *		for storage.
  */
 static void
 tuplestore_updatemax(Tuplestorestate *state)
@@ -1499,37 +1510,37 @@ tuplestore_updatemax(Tuplestorestate *state)
 	if (state->status == TSS_INMEM)
 		state->maxSpace = Max(state->maxSpace,
 							  state->allowedMem - state->availMem);
-}
-
-/*
- * tuplestore_storage_type_name
- *		Return a string description of the storage method used to store the
- *		tuples.
- */
-const char *
-tuplestore_storage_type_name(Tuplestorestate *state)
-{
-	if (state->status == TSS_INMEM)
-		return "Memory";
 	else
-		return "Disk";
+	{
+		state->maxSpace = Max(state->maxSpace,
+							  BufFileSize(state->myfile));
+
+		/*
+		 * usedDisk never gets set to false again after spilling to disk, even
+		 * if tuplestore_clear() is called and new tuples go to memory again.
+		 */
+		state->usedDisk = true;
+	}
 }
 
 /*
- * tuplestore_space_used
- *		Return the maximum space used in memory unless the tuplestore has spilled
- *		to disk, in which case, return the disk space used.
+ * tuplestore_get_stats
+ *		Obtain statistics about the maximum space used by the tuplestore.
+ *		These statistics are the maximums and are not reset by calls to
+ *		tuplestore_trim() or tuplestore_clear().
  */
-int64
-tuplestore_space_used(Tuplestorestate *state)
+void
+tuplestore_get_stats(Tuplestorestate *state, char **max_storage_type,
+					 int64 *max_space)
 {
-	/* First, update the maxSpace field */
 	tuplestore_updatemax(state);
 
-	if (state->status == TSS_INMEM)
-		return state->maxSpace;
+	if (state->usedDisk)
+		*max_storage_type = "Disk";
 	else
-		return BufFileSize(state->myfile);
+		*max_storage_type = "Memory";
+
+	*max_space = state->maxSpace;
 }
 
 /*
@@ -1601,7 +1612,6 @@ writetup_heap(Tuplestorestate *state, void *tup)
 	if (state->backward)		/* need trailing length word? */
 		BufFileWrite(state->myfile, &tuplen, sizeof(tuplen));
 
-	/* no need to call tuplestore_updatemax() when not in TSS_INMEM */
 	FREEMEM(state, GetMemoryChunkSpace(tuple));
 	heap_free_minimal_tuple(tuple);
 }
diff --git a/src/include/utils/tuplestore.h b/src/include/utils/tuplestore.h
index 3d8a90caaf9..e4782ac5a7f 100644
--- a/src/include/utils/tuplestore.h
+++ b/src/include/utils/tuplestore.h
@@ -65,9 +65,8 @@ extern void tuplestore_copy_read_pointer(Tuplestorestate *state,
 
 extern void tuplestore_trim(Tuplestorestate *state);
 
-extern const char *tuplestore_storage_type_name(Tuplestorestate *state);
-
-extern int64 tuplestore_space_used(Tuplestorestate *state);
+extern void tuplestore_get_stats(Tuplestorestate *state, char **storage_type,
+								 int64 *max_space);
 
 extern bool tuplestore_in_memory(Tuplestorestate *state);