From 1e2491a448642e6086e85b8f5e04502f8cdf752c Mon Sep 17 00:00:00 2001
From: Thomas Munro <thomas.munro@gmail.com>
Date: Wed, 11 Dec 2024 21:22:08 +1300
Subject: [PATCH v5 2/2] Fix pg_truncate_visibility_map() protocol.

pg_truncate_visibility_map() is supposed to be similar to
RelationTruncate(), but it is missing several developments.  Reorder
operations to match, and apply changes equivalent to commits 412ad7a5,
75818b3a, TODO:patch-0001 to create atomicity at do and redo time.

XXX The comments about atomicity are untrue without XLogFlush().
---
 contrib/pg_visibility/pg_visibility.c | 30 +++++++++++++++++++++------
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/contrib/pg_visibility/pg_visibility.c b/contrib/pg_visibility/pg_visibility.c
index 5d0deaba61e..32999ac9a97 100644
--- a/contrib/pg_visibility/pg_visibility.c
+++ b/contrib/pg_visibility/pg_visibility.c
@@ -19,6 +19,7 @@
 #include "funcapi.h"
 #include "miscadmin.h"
 #include "storage/bufmgr.h"
+#include "storage/proc.h"
 #include "storage/procarray.h"
 #include "storage/read_stream.h"
 #include "storage/smgr.h"
@@ -390,6 +391,7 @@ pg_truncate_visibility_map(PG_FUNCTION_ARGS)
 	Relation	rel;
 	ForkNumber	fork;
 	BlockNumber block;
+	BlockNumber old_block;
 
 	rel = relation_open(relid, AccessExclusiveLock);
 
@@ -399,12 +401,22 @@ pg_truncate_visibility_map(PG_FUNCTION_ARGS)
 	/* Forcibly reset cached file size */
 	RelationGetSmgr(rel)->smgr_cached_nblocks[VISIBILITYMAP_FORKNUM] = InvalidBlockNumber;
 
+	/* Compute new and old size before entering critical section. */
+	fork = VISIBILITYMAP_FORKNUM;
 	block = visibilitymap_prepare_truncate(rel, 0);
-	if (BlockNumberIsValid(block))
-	{
-		fork = VISIBILITYMAP_FORKNUM;
-		smgrtruncate(RelationGetSmgr(rel), &fork, 1, &block);
-	}
+	old_block = BlockNumberIsValid(block) ? smgrnblocks(RelationGetSmgr(rel), fork) : 0;
+
+	/*
+	 * Buffer dropping, file truncation and WAL logging must be effectively
+	 * atomic.  Interrupts are suppressed (buffer I/O waits can't be canceled),
+	 * and we either succeed in all steps or panic.  Also prevent checkpoint
+	 * start/complete, so we can't crash and recover from a redo point after
+	 * this record if we haven't finished truncating and fsync'ing the file(s).
+	 * See RelationTruncate() for discussion.
+	 */
+	Assert((MyProc->delayChkptFlags & (DELAY_CHKPT_START | DELAY_CHKPT_COMPLETE)) == 0);
+	MyProc->delayChkptFlags |= DELAY_CHKPT_START | DELAY_CHKPT_COMPLETE;
+	START_CRIT_SECTION();
 
 	if (RelationNeedsWAL(rel))
 	{
@@ -420,6 +432,12 @@ pg_truncate_visibility_map(PG_FUNCTION_ARGS)
 		XLogInsert(RM_SMGR_ID, XLOG_SMGR_TRUNCATE | XLR_SPECIAL_REL_UPDATE);
 	}
 
+	if (BlockNumberIsValid(block))
+		smgrtruncatefrom(RelationGetSmgr(rel), &fork, 1, &old_block, &block);
+
+	END_CRIT_SECTION();
+	MyProc->delayChkptFlags &= ~(DELAY_CHKPT_START | DELAY_CHKPT_COMPLETE);
+
 	/*
 	 * Release the lock right away, not at commit time.
 	 *
@@ -430,7 +448,7 @@ pg_truncate_visibility_map(PG_FUNCTION_ARGS)
 	 * here and when we sent the messages at our eventual commit.  However,
 	 * we're currently only sending a non-transactional smgr invalidation,
 	 * which will have been posted to shared memory immediately from within
-	 * smgr_truncate.  Therefore, there should be no race here.
+	 * smgrtruncatefrom.  Therefore, there should be no race here.
 	 *
 	 * The reason why it's desirable to release the lock early here is because
 	 * of the possibility that someone will need to use this to blow away many
-- 
2.39.5

