1 files changed, 38 insertions, 10 deletions

diff --git a/src/backend/access/heap/visibilitymap.c b/src/backend/access/heap/visibilitymap.c
index 3209c87bb86..af64fe97e89 100644
--- a/src/backend/access/heap/visibilitymap.c
+++ b/src/backend/access/heap/visibilitymap.c
@@ -233,13 +233,18 @@ visibilitymap_pin_ok(BlockNumber heapBlk, Buffer buf)
  * marked all-visible; it is needed for Hot Standby, and can be
  * InvalidTransactionId if the page contains no tuples.
  *
+ * Caller is expected to set the heap page's PD_ALL_VISIBLE bit before calling
+ * this function. Except in recovery, caller should also pass the heap
+ * buffer. When checksums are enabled and we're not in recovery, we must add
+ * the heap buffer to the WAL chain to protect it from being torn.
+ *
  * You must pass a buffer containing the correct map page to this function.
  * Call visibilitymap_pin first to pin the right one. This function doesn't do
  * any I/O.
  */
 void
-visibilitymap_set(Relation rel, BlockNumber heapBlk, XLogRecPtr recptr,
-				  Buffer buf, TransactionId cutoff_xid)
+visibilitymap_set(Relation rel, BlockNumber heapBlk, Buffer heapBuf,
+				  XLogRecPtr recptr, Buffer vmBuf, TransactionId cutoff_xid)
 {
 	BlockNumber mapBlock = HEAPBLK_TO_MAPBLOCK(heapBlk);
 	uint32		mapByte = HEAPBLK_TO_MAPBYTE(heapBlk);
@@ -252,34 +257,55 @@ visibilitymap_set(Relation rel, BlockNumber heapBlk, XLogRecPtr recptr,
 #endif
 
 	Assert(InRecovery || XLogRecPtrIsInvalid(recptr));
+	Assert(InRecovery || BufferIsValid(heapBuf));
 
-	/* Check that we have the right page pinned */
-	if (!BufferIsValid(buf) || BufferGetBlockNumber(buf) != mapBlock)
-		elog(ERROR, "wrong buffer passed to visibilitymap_set");
+	/* Check that we have the right heap page pinned, if present */
+	if (BufferIsValid(heapBuf) && BufferGetBlockNumber(heapBuf) != heapBlk)
+		elog(ERROR, "wrong heap buffer passed to visibilitymap_set");
 
-	page = BufferGetPage(buf);
+	/* Check that we have the right VM page pinned */
+	if (!BufferIsValid(vmBuf) || BufferGetBlockNumber(vmBuf) != mapBlock)
+		elog(ERROR, "wrong VM buffer passed to visibilitymap_set");
+
+	page = BufferGetPage(vmBuf);
 	map = PageGetContents(page);
-	LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
+	LockBuffer(vmBuf, BUFFER_LOCK_EXCLUSIVE);
 
 	if (!(map[mapByte] & (1 << mapBit)))
 	{
 		START_CRIT_SECTION();
 
 		map[mapByte] |= (1 << mapBit);
-		MarkBufferDirty(buf);
+		MarkBufferDirty(vmBuf);
 
 		if (RelationNeedsWAL(rel))
 		{
 			if (XLogRecPtrIsInvalid(recptr))
-				recptr = log_heap_visible(rel->rd_node, heapBlk, buf,
+			{
+				Assert(!InRecovery);
+				recptr = log_heap_visible(rel->rd_node, heapBuf, vmBuf,
 										  cutoff_xid);
+
+				/*
+				 * If data checksums are enabled, we need to protect the heap
+				 * page from being torn.
+				 */
+				if (DataChecksumsEnabled())
+				{
+					Page heapPage = BufferGetPage(heapBuf);
+
+					/* caller is expected to set PD_ALL_VISIBLE first */
+					Assert(PageIsAllVisible(heapPage));
+					PageSetLSN(heapPage, recptr);
+				}
+			}
 			PageSetLSN(page, recptr);
 		}
 
 		END_CRIT_SECTION();
 	}
 
-	LockBuffer(buf, BUFFER_LOCK_UNLOCK);
+	LockBuffer(vmBuf, BUFFER_LOCK_UNLOCK);
 }
 
 /*
@@ -579,6 +605,8 @@ vm_extend(Relation rel, BlockNumber vm_nblocks)
 	/* Now extend the file */
 	while (vm_nblocks_now < vm_nblocks)
 	{
+		PageSetChecksumInplace(pg, vm_nblocks_now);
+
 		smgrextend(rel->rd_smgr, VISIBILITYMAP_FORKNUM, vm_nblocks_now,
 				   (char *) pg, false);
 		vm_nblocks_now++;