From 3b1657a0ba51299fa0384a0032759f3661a2baca Mon Sep 17 00:00:00 2001
From: Melanie Plageman <melanieplageman@gmail.com>
Date: Mon, 31 Mar 2025 22:02:25 -0400
Subject: [PATCH v11 2/3] Refactor autoprewarm_database_main() in preparation
 for read stream

Autoprewarm prewarms blocks from a dump file representing the contents
of shared buffers last time it was dumped. It uses a sorted array of
BlockInfoRecords, each representing a block from one of the cluster's
databases and tables.

autoprewarm_database_main() prewarms all the blocks from a single
database. It is optimized to ensure we don't try to open the same
relation or fork over and over again if it has been dropped or is
invalid. The main loop handled this by carefully setting various local
variables to sentinel values when a run of blocks should be skipped.

This method won't work with the read stream API. A read stream can only
be created for a single relation and fork combination. The callback has
to be able to advance the position in the array to allow for reading
ahead additional blocks, however the callback cannot try to open another
relation or close the current relation. So, the main loop in
autoprewarm_database_main() must also advance the position in the array
of BlockInfoRecords.

To make it compatible with the read stream API, change
autoprewarm_database_main() to explicitly fast-forward in the array past
the blocks belonging to an invalid relation or fork.

This commit only implements the new control flow -- it does not use the
read stream API.

Co-authored-by: Nazir Bilal Yavuz <byavuz81@gmail.com>
Co-authored-by: Melanie Plageman <melanieplageman@gmail.com>
Discussion: https://postgr.es/m/flat/CAN55FZ3n8Gd%2BhajbL%3D5UkGzu_aHGRqnn%2BxktXq2fuds%3D1AOR6Q%40mail.gmail.com
---
 contrib/pg_prewarm/autoprewarm.c | 174 ++++++++++++++++++-------------
 1 file changed, 101 insertions(+), 73 deletions(-)

diff --git a/contrib/pg_prewarm/autoprewarm.c b/contrib/pg_prewarm/autoprewarm.c
index 172654fee25..b02859d8ec3 100644
--- a/contrib/pg_prewarm/autoprewarm.c
+++ b/contrib/pg_prewarm/autoprewarm.c
@@ -426,12 +426,11 @@ apw_load_buffers(void)
 void
 autoprewarm_database_main(Datum main_arg)
 {
-	int			pos;
 	BlockInfoRecord *block_info;
-	Relation	rel = NULL;
-	BlockNumber nblocks = 0;
-	BlockInfoRecord *old_blk = NULL;
+	int			i;
+	BlockInfoRecord blk;
 	dsm_segment *seg;
+	Oid			database;
 
 	/* Establish signal handlers; once that's done, unblock signals. */
 	pqsignal(SIGTERM, die);
@@ -446,104 +445,133 @@ autoprewarm_database_main(Datum main_arg)
 				 errmsg("could not map dynamic shared memory segment")));
 	BackgroundWorkerInitializeConnectionByOid(apw_state->database, InvalidOid, 0);
 	block_info = (BlockInfoRecord *) dsm_segment_address(seg);
-	pos = apw_state->prewarm_start_idx;
+
+	i = apw_state->prewarm_start_idx;
+	blk = block_info[i];
+
+	/*
+	 * apw_state->database may differ from blk.database if we are prewarming
+	 * blocks from global objects.
+	 */
+	database = blk.database;
 
 	/*
 	 * Loop until we run out of blocks to prewarm or until we run out of free
-	 * buffers.
+	 * buffers. We'll quit if we've reached records for another database,
 	 */
-	while (pos < apw_state->prewarm_stop_idx && have_free_buffer())
+	while (i < apw_state->prewarm_stop_idx &&
+		   blk.database == database &&
+		   have_free_buffer())
 	{
-		BlockInfoRecord *blk = &block_info[pos++];
-		Buffer		buf;
+		RelFileNumber filenumber = blk.filenumber;
+		Oid			reloid;
+		Relation	rel;
 
 		CHECK_FOR_INTERRUPTS();
 
-		/* Quit if we've reached records for another database. */
-		if (old_blk != NULL && old_blk->database != blk->database)
-			break;
+		StartTransactionCommand();
 
-		/*
-		 * As soon as we encounter a block of a new relation, close the old
-		 * relation. Note that rel will be NULL if try_relation_open failed
-		 * previously; in that case, there is nothing to close.
-		 */
-		if (old_blk != NULL && old_blk->filenumber != blk->filenumber &&
-			rel != NULL)
+		reloid = RelidByRelfilenumber(blk.tablespace, blk.filenumber);
+		if (!OidIsValid(reloid) ||
+			(rel = try_relation_open(reloid, AccessShareLock)) == NULL)
 		{
-			relation_close(rel, AccessShareLock);
-			rel = NULL;
+			/* We failed to open the relation, so there is nothing to close. */
 			CommitTransactionCommand();
-		}
 
-		/*
-		 * Try to open each new relation, but only once, when we first
-		 * encounter it. If it's been dropped, skip the associated blocks.
-		 */
-		if (old_blk == NULL || old_blk->filenumber != blk->filenumber)
-		{
-			Oid			reloid;
-
-			Assert(rel == NULL);
-			StartTransactionCommand();
-			reloid = RelidByRelfilenumber(blk->tablespace, blk->filenumber);
-			if (OidIsValid(reloid))
-				rel = try_relation_open(reloid, AccessShareLock);
+			/*
+			 * Fast-forward to the next relation. We want to skip all of the
+			 * other records referencing this relation since we know we can't
+			 * open it. That way, we avoid repeatedly trying and failing to
+			 * open the same relation.
+			 */
+			for (; i < apw_state->prewarm_stop_idx; i++)
+			{
+				blk = block_info[i];
+				if (blk.database != database ||
+					blk.filenumber != filenumber)
+					break;
+			}
 
-			if (!rel)
-				CommitTransactionCommand();
-		}
-		if (!rel)
-		{
-			old_blk = blk;
+			/* Time to try and open our new found relation */
 			continue;
 		}
 
-		/* Once per fork, check for fork existence and size. */
-		if (old_blk == NULL ||
-			old_blk->filenumber != blk->filenumber ||
-			old_blk->forknum != blk->forknum)
+		/*
+		 * We have a relation; now let's loop until we find a valid fork of
+		 * the relation or we run out of free buffers. Once we've read from
+		 * all valid forks or run out of options, we'll close the relation and
+		 * move on.
+		 */
+		while (i < apw_state->prewarm_stop_idx &&
+			   blk.database == database &&
+			   blk.filenumber == filenumber &&
+			   have_free_buffer())
 		{
+			ForkNumber	forknum = blk.forknum;
+			BlockNumber nblocks;
+			Buffer		buf;
+
 			/*
 			 * smgrexists is not safe for illegal forknum, hence check whether
 			 * the passed forknum is valid before using it in smgrexists.
 			 */
-			if (blk->forknum > InvalidForkNumber &&
-				blk->forknum <= MAX_FORKNUM &&
-				smgrexists(RelationGetSmgr(rel), blk->forknum))
-				nblocks = RelationGetNumberOfBlocksInFork(rel, blk->forknum);
-			else
-				nblocks = 0;
-		}
+			if (blk.forknum <= InvalidForkNumber ||
+				blk.forknum > MAX_FORKNUM ||
+				!smgrexists(RelationGetSmgr(rel), blk.forknum))
+			{
+				/*
+				 * Fast-forward to the next fork. We want to skip all of the
+				 * other records referencing this fork since we already know
+				 * it's not valid.
+				 */
+				for (; i < apw_state->prewarm_stop_idx; i++)
+				{
+					blk = block_info[i];
+					if (blk.database != database ||
+						blk.filenumber != filenumber ||
+						blk.forknum != forknum)
+						break;
+				}
+
+				/* Time to check if this newfound fork is valid */
+				continue;
+			}
 
-		/* Check whether blocknum is valid and within fork file size. */
-		if (blk->blocknum >= nblocks)
-		{
-			/* Move to next forknum. */
-			old_blk = blk;
-			continue;
-		}
+			nblocks = RelationGetNumberOfBlocksInFork(rel, blk.forknum);
 
-		/* Prewarm buffer. */
-		buf = ReadBufferExtended(rel, blk->forknum, blk->blocknum, RBM_NORMAL,
-								 NULL);
-		if (BufferIsValid(buf))
-		{
-			apw_state->prewarmed_blocks++;
-			ReleaseBuffer(buf);
-		}
+			/* Check whether blocknum is valid and within fork file size. */
+			if (blk.blocknum >= nblocks)
+			{
+				blk = block_info[++i];
+				continue;
+			}
 
-		old_blk = blk;
-	}
+			/* Prewarm buffers. */
+			while (i < apw_state->prewarm_stop_idx &&
+				   blk.database == database &&
+				   blk.filenumber == filenumber &&
+				   blk.forknum == forknum &&
+				   have_free_buffer())
+			{
+				CHECK_FOR_INTERRUPTS();
 
-	dsm_detach(seg);
+				buf = ReadBufferExtended(rel, blk.forknum, blk.blocknum, RBM_NORMAL,
+										 NULL);
+
+				blk = block_info[++i];
+				if (!BufferIsValid(buf))
+					break;
+
+				apw_state->prewarmed_blocks++;
+				ReleaseBuffer(buf);
+			}
+		}
 
-	/* Release lock on previous relation. */
-	if (rel)
-	{
 		relation_close(rel, AccessShareLock);
 		CommitTransactionCommand();
 	}
+
+	dsm_detach(seg);
 }
 
 /*
-- 
2.34.1

