From 186cc5f60d8ffb11ca5150a3ece53e19368448d1 Mon Sep 17 00:00:00 2001
From: Jehan-Guillaume de Rorthais <jgdr@dalibo.com>
Date: Fri, 7 Apr 2023 19:24:20 +0200
Subject: [PATCH 4/4] Limit BufFile memory explosion with bad HashJoin

When hash join rely on a largely underestimated
statistics, splitting batches could result of a
huge memory consumption because BufFile buffers
associated to batches were not accounted.

This patch tries to keep the memory consummed
by batches and actual data balanced by rising
the memory limit to make more room for the data
before deciding to split again the batches.
---
 src/backend/executor/nodeHash.c     | 101 +++++++++++++++++++++++-----
 src/backend/executor/nodeHashjoin.c |  26 +++++--
 src/include/executor/nodeHashjoin.h |   2 +-
 3 files changed, 107 insertions(+), 22 deletions(-)

diff --git a/src/backend/executor/nodeHash.c b/src/backend/executor/nodeHash.c
index ec6b80121b..3a6e0bcfb7 100644
--- a/src/backend/executor/nodeHash.c
+++ b/src/backend/executor/nodeHash.c
@@ -81,6 +81,8 @@ static bool ExecParallelHashTuplePrealloc(HashJoinTable hashtable,
 static void ExecParallelHashMergeCounters(HashJoinTable hashtable);
 static void ExecParallelHashCloseBatchAccessors(HashJoinTable hashtable);
 
+static void ExecHashUpdateSpacePeak(HashJoinTable hashtable);
+
 static void debugIncreaseBatches(HashJoinTable hashtable, char * label)
 {
 	int allocInnerBufFiles = 0;
@@ -224,10 +226,8 @@ MultiExecPrivateHash(HashState *node)
 	if (hashtable->nbuckets != hashtable->nbuckets_optimal)
 		ExecHashIncreaseNumBuckets(hashtable);
 
-	/* Account for the buckets in spaceUsed (reported in EXPLAIN ANALYZE) */
-	hashtable->spaceUsed += hashtable->nbuckets * sizeof(HashJoinTuple);
-	if (hashtable->spaceUsed > hashtable->spacePeak)
-		hashtable->spacePeak = hashtable->spaceUsed;
+	/* refresh info about peak used memory */
+	ExecHashUpdateSpacePeak(hashtable);
 
 	hashtable->partialTuples = hashtable->totalTuples;
 }
@@ -969,6 +969,56 @@ ExecHashIncreaseNumBatches(HashJoinTable hashtable)
 	nbatch = oldnbatch * 2;
 	Assert(nbatch > 1);
 
+	/*
+	 * Each batch requires a non-trivial amount of memory, because BufFile
+	 * includes a PGAlignedBlock (typically 8kB buffer). So when doubling
+	 * the number of batches, we need to be careful and only allow that if
+	 * it actually has a chance of reducing memory usage.
+	 *
+	 * When doubling the number of batches, we expect to save roughly 1/2
+	 * of memory currently used for data (rows) at the price of doubling
+	 * the memory used for BufFile.
+	 * This becomes pointless when memory used for BufFile is greater than
+	 * the memory we expect to save:
+	 *
+	 *		(2 * nbatches * sizeof(BufFile)) > spaceUsed/2
+	 *
+	 * As far as this is confined inside work_mem, that's fine. But if
+	 * we significantly underestimate the number of batches, we may end
+	 * up in a situation where BufFile alone exceed work_mem.
+	 *
+	 * In such situation, move the threshold a bit, until the next point
+	 * where it'll make sense to consider adding batches again.
+	 * 
+	 * We can't stop adding batches entirely, because that would just mean
+	 * the batches would need more and more memory. So we need to increase
+	 * the number of batches, even if we can't enforce work_mem properly.
+	 * 
+	 * Note: This applies mostly to cases of significant underestimates,
+	 * resulting in an explosion of the number of batches. The properly
+	 * estimated cases should generally end up using merge join based on
+	 * high cost of the batched hash join.
+	 */
+	/*
+	 * fileCxt size is good enough estimation of BufFiles consumption.
+	 * Keep in mind spaceUsed includes real BufFile consumption as well
+	 */
+	if (2 * hashtable->fileCxt->mem_allocated > hashtable->spaceUsed / 2)
+	{
+		Size new_limit = nbatch * sizeof(PGAlignedBlock) * 3;
+		if (hashtable->spaceAllowed < new_limit)
+		{
+			ereport(WARNING, (
+				errmsg("Hash Join node must grow outside of work_mem"),
+				errdetail("Rising memory limit from %ld to %ld",
+						  hashtable->spaceAllowed, new_limit),
+				errhint("You might need to ANALYZE your table or tune its statistics collection.")));
+			hashtable->spaceAllowed = new_limit;
+		}
+
+		return;
+	}
+
 #ifdef HJDEBUG
 	printf("Hashjoin %p: increasing nbatch to %d because space = %zu\n",
 		   hashtable, nbatch, hashtable->spaceUsed);
@@ -1067,7 +1117,7 @@ ExecHashIncreaseNumBatches(HashJoinTable hashtable)
 				ExecHashJoinSaveTuple(HJTUPLE_MINTUPLE(hashTuple),
 									  hashTuple->hashvalue,
 									  &hashtable->innerBatchFile[batchno],
-									  hashtable->fileCxt);
+									  hashtable);
 
 				hashtable->spaceUsed -= hashTupleSize;
 				nfreed++;
@@ -1711,14 +1761,19 @@ ExecHashTableInsert(HashJoinTable hashtable,
 
 		/* Account for space used, and back off if we've used too much */
 		hashtable->spaceUsed += hashTupleSize;
-		if (hashtable->spaceUsed > hashtable->spacePeak)
-			hashtable->spacePeak = hashtable->spaceUsed;
+
+		/* refresh info about peak used memory */
+		ExecHashUpdateSpacePeak(hashtable);
+
+		/* Consider increasing number of batches. */
 		if (hashtable->spaceUsed +
 			hashtable->nbuckets_optimal * sizeof(HashJoinTuple)
 			> hashtable->spaceAllowed)
 		{
 			debugIncreaseBatches(hashtable, "trying to save memory");
+
 			ExecHashIncreaseNumBatches(hashtable);
+
 			debugIncreaseBatches(hashtable, "memory rescue done");
 		}
 	}
@@ -1732,7 +1787,7 @@ ExecHashTableInsert(HashJoinTable hashtable,
 		ExecHashJoinSaveTuple(tuple,
 							  hashvalue,
 							  &hashtable->innerBatchFile[batchno],
-							  hashtable->fileCxt);
+							  hashtable);
 	}
 
 	if (shouldFree)
@@ -1982,6 +2037,18 @@ ExecHashGetBucketAndBatch(HashJoinTable hashtable,
 	}
 }
 
+static void
+ExecHashUpdateSpacePeak(HashJoinTable hashtable)
+{
+	Size	spaceUsed = hashtable->spaceUsed;
+
+	/* Account for the buckets in spaceUsed (reported in EXPLAIN ANALYZE) */
+	spaceUsed += hashtable->nbuckets * sizeof(HashJoinTuple);
+
+	if (spaceUsed > hashtable->spacePeak)
+		hashtable->spacePeak = spaceUsed;
+}
+
 /*
  * ExecScanHashBucket
  *		scan a hash bucket for matches to the current outer tuple
@@ -2489,8 +2556,9 @@ ExecHashBuildSkewHash(HashJoinTable hashtable, Hash *node, int mcvsToUse)
 			+ mcvsToUse * sizeof(int);
 		hashtable->spaceUsedSkew += nbuckets * sizeof(HashSkewBucket *)
 			+ mcvsToUse * sizeof(int);
-		if (hashtable->spaceUsed > hashtable->spacePeak)
-			hashtable->spacePeak = hashtable->spaceUsed;
+
+		/* refresh info about peak used memory */
+		ExecHashUpdateSpacePeak(hashtable);
 
 		/*
 		 * Create a skew bucket for each MCV hash value.
@@ -2540,8 +2608,9 @@ ExecHashBuildSkewHash(HashJoinTable hashtable, Hash *node, int mcvsToUse)
 			hashtable->nSkewBuckets++;
 			hashtable->spaceUsed += SKEW_BUCKET_OVERHEAD;
 			hashtable->spaceUsedSkew += SKEW_BUCKET_OVERHEAD;
-			if (hashtable->spaceUsed > hashtable->spacePeak)
-				hashtable->spacePeak = hashtable->spaceUsed;
+
+			/* refresh info about peak used memory */
+			ExecHashUpdateSpacePeak(hashtable);
 		}
 
 		free_attstatsslot(&sslot);
@@ -2630,8 +2699,10 @@ ExecHashSkewTableInsert(HashJoinTable hashtable,
 	/* Account for space used, and back off if we've used too much */
 	hashtable->spaceUsed += hashTupleSize;
 	hashtable->spaceUsedSkew += hashTupleSize;
-	if (hashtable->spaceUsed > hashtable->spacePeak)
-		hashtable->spacePeak = hashtable->spaceUsed;
+
+	/* refresh info about peak used memory */
+	ExecHashUpdateSpacePeak(hashtable);
+
 	while (hashtable->spaceUsedSkew > hashtable->spaceAllowedSkew)
 		ExecHashRemoveNextSkewBucket(hashtable);
 
@@ -2715,7 +2786,7 @@ ExecHashRemoveNextSkewBucket(HashJoinTable hashtable)
 
 			ExecHashJoinSaveTuple(tuple, hashvalue,
 								  &hashtable->innerBatchFile[batchno],
-								  hashtable->fileCxt);
+								  hashtable);
 
 			pfree(hashTuple);
 			hashtable->spaceUsed -= tupleSize;
diff --git a/src/backend/executor/nodeHashjoin.c b/src/backend/executor/nodeHashjoin.c
index 5bc7f814c6..36fb024fab 100644
--- a/src/backend/executor/nodeHashjoin.c
+++ b/src/backend/executor/nodeHashjoin.c
@@ -488,7 +488,7 @@ ExecHashJoinImpl(PlanState *pstate, bool parallel)
 
 					ExecHashJoinSaveTuple(mintuple, hashvalue,
 										  &hashtable->outerBatchFile[batchno],
-										  hashtable->fileCxt);
+										  hashtable);
 
 					if (shouldFree)
 						heap_free_minimal_tuple(mintuple);
@@ -1041,8 +1041,11 @@ ExecHashJoinNewBatch(HashJoinState *hjstate)
 		 * away to free disk space.
 		 */
 		if (hashtable->outerBatchFile[curbatch])
+		{
 			BufFileClose(hashtable->outerBatchFile[curbatch]);
-		hashtable->outerBatchFile[curbatch] = NULL;
+			hashtable->outerBatchFile[curbatch] = NULL;
+			hashtable->spaceUsed -= sizeof(PGAlignedBlock);
+		}
 	}
 	else						/* we just finished the first batch */
 	{
@@ -1096,11 +1099,19 @@ ExecHashJoinNewBatch(HashJoinState *hjstate)
 		/* We can ignore this batch. */
 		/* Release associated temp files right away. */
 		if (hashtable->innerBatchFile[curbatch])
+		{
 			BufFileClose(hashtable->innerBatchFile[curbatch]);
-		hashtable->innerBatchFile[curbatch] = NULL;
+			hashtable->innerBatchFile[curbatch] = NULL;
+			hashtable->spaceUsed -= sizeof(PGAlignedBlock);
+		}
+
 		if (hashtable->outerBatchFile[curbatch])
+		{
 			BufFileClose(hashtable->outerBatchFile[curbatch]);
-		hashtable->outerBatchFile[curbatch] = NULL;
+			hashtable->outerBatchFile[curbatch] = NULL;
+			hashtable->spaceUsed -= sizeof(PGAlignedBlock);
+		}
+
 		curbatch++;
 	}
 
@@ -1141,6 +1152,7 @@ ExecHashJoinNewBatch(HashJoinState *hjstate)
 		 */
 		BufFileClose(innerFile);
 		hashtable->innerBatchFile[curbatch] = NULL;
+		hashtable->spaceUsed -= sizeof(PGAlignedBlock);
 	}
 
 	/*
@@ -1305,7 +1317,7 @@ ExecParallelHashJoinNewBatch(HashJoinState *hjstate)
  */
 void
 ExecHashJoinSaveTuple(MinimalTuple tuple, uint32 hashvalue,
-					  BufFile **fileptr, MemoryContext filecxt)
+					  BufFile **fileptr, HashJoinTable hashtable)
 {
 	BufFile    *file = *fileptr;
 
@@ -1313,13 +1325,15 @@ ExecHashJoinSaveTuple(MinimalTuple tuple, uint32 hashvalue,
 	{
 		MemoryContext oldctx;
 
-		oldctx = MemoryContextSwitchTo(filecxt);
+		oldctx = MemoryContextSwitchTo(hashtable->fileCxt);
 
 		/* First write to this batch file, so open it. */
 		file = BufFileCreateTemp(false);
 		*fileptr = file;
 
 		MemoryContextSwitchTo(oldctx);
+
+		hashtable->spaceUsed += sizeof(PGAlignedBlock);
 	}
 
 	BufFileWrite(file, &hashvalue, sizeof(uint32));
diff --git a/src/include/executor/nodeHashjoin.h b/src/include/executor/nodeHashjoin.h
index a8f9ae1989..ccb704ede1 100644
--- a/src/include/executor/nodeHashjoin.h
+++ b/src/include/executor/nodeHashjoin.h
@@ -29,6 +29,6 @@ extern void ExecHashJoinInitializeWorker(HashJoinState *state,
 										 ParallelWorkerContext *pwcxt);
 
 extern void ExecHashJoinSaveTuple(MinimalTuple tuple, uint32 hashvalue,
-								  BufFile **fileptr, MemoryContext filecxt);
+								  BufFile **fileptr, HashJoinTable hashtable);
 
 #endif							/* NODEHASHJOIN_H */
-- 
2.39.2

