IOS/TableAM: Support AM-specific fast visibility tests

Previously, we assumed VM_ALL_VISIBLE is universal across all AMs. This
is probably not the case, so we introduce a new table method called
"table_index_vischeck_tuples" which allows anyone to ask the AM whether
a tuple is definitely visible to everyone or might be invisible to
someone.

The API is intended to replace direct calls to VM_ALL_VISIBLE and as such
doesn't include "definitely dead to everyone", as the Heap AM's VM doesn't
support *definitely dead* as output for its lookups; and thus it would be
too expensive for the Heap AM to produce such results.

A future commit will use this inside GIST and SP-GIST to fix a race
condition between IOS and VACUUM, which causes a bug with tuple
visibility, and a further patch will add support for this to nbtree.

Author: MMeent

src/backend/access/heap/heapam.c

@@ -101,11 +101,37 @@ static bool ConditionalMultiXactIdWait(MultiXactId multi, MultiXactStatus status
 									   uint16 infomask, Relation rel, int *remaining,
 									   bool logLockFailure);
 static void index_delete_sort(TM_IndexDeleteOp *delstate);
+static inline int heap_ivc_process_block(Relation rel, Buffer *vmbuf,
+										 TM_VisCheck *checks, int nchecks);
+static void heap_ivc_process_all(Relation rel, Buffer *vmbuf,
+								 TM_VisCheck *checks, int nchecks);
 static int	bottomup_sort_and_shrink(TM_IndexDeleteOp *delstate);
 static XLogRecPtr log_heap_new_cid(Relation relation, HeapTuple tup);
 static HeapTuple ExtractReplicaIdentity(Relation relation, HeapTuple tp, bool key_required,
 										bool *copy);
 
+/* sort template definitions for index */
+#define ST_SORT heap_ivc_sortby_tidheapblk
+#define ST_ELEMENT_TYPE TM_VisCheck
+#define ST_DECLARE
+#define ST_DEFINE
+#define ST_SCOPE static inline
+#define ST_COMPARE(a, b) ( \
+	a->tidblkno < b->tidblkno ? -1 : ( \
+		a->tidblkno > b->tidblkno ? 1 : 0 \
+	) \
+)
+
+#include "lib/sort_template.h"
+
+#define ST_SORT heap_ivc_sortby_idx
+#define ST_ELEMENT_TYPE TM_VisCheck
+#define ST_DECLARE
+#define ST_DEFINE
+#define ST_SCOPE static inline
+#define ST_COMPARE(a, b) (((int) a->idxoffnum) - ((int) b->idxoffnum))
+#include "lib/sort_template.h"
+
 
 /*
  * Each tuple lock mode has a corresponding heavyweight lock, and one or two
@@ -8779,6 +8805,157 @@ bottomup_sort_and_shrink(TM_IndexDeleteOp *delstate)
 	return nblocksfavorable;
 }
 
+/*
+ * heapam implementation of tableam's index_vischeck_tuples interface.
+ *
+ * This helper function is called by index AMs during index-only scans,
+ * to do VM-based visibility checks on individual tuples, so that the AM
+ * can hold the tuple in memory for e.g. reordering for extended periods of
+ * time while without holding thousands of pins to conflict with VACUUM.
+ *
+ * It's possible for this to generate a fair amount of I/O, since we may be
+ * checking hundreds of tuples from a single index block, but that is
+ * preferred over holding thousands of pins.
+ *
+ * We use heuristics to balance the costs of sorting TIDs with VM page
+ * lookups.
+ */
+void
+heap_index_vischeck_tuples(Relation rel, TM_IndexVisibilityCheckOp *checkop)
+{
+	Buffer			vmbuf = *checkop->vmbuf;
+	Buffer			storvmbuf = vmbuf;
+	TM_VisCheck	   *checks = checkop->checktids;
+	int				checkntids = checkop->checkntids;
+	int				upcomingvmbufchanges = 0;
+
+	/*
+	 * The first index scan will have to pin the VM buffer, and that first
+	 * change in the vm buffer shouldn't put us into the expensive VM page &
+	 * sort path; so we special-case this operation.
+	 */
+	if (!BufferIsValid(vmbuf))
+	{
+		int			processed;
+		processed = heap_ivc_process_block(rel, &vmbuf, checks,checkntids);
+		checkntids -= processed;
+		checks += processed;
+		storvmbuf = vmbuf;
+		Assert(processed > 0);
+	}
+
+	while (vmbuf == storvmbuf && checkntids > 0)
+	{
+		int			processed;
+
+		processed = heap_ivc_process_block(rel, &vmbuf, checks,checkntids);
+
+		Assert(processed <= checkntids);
+
+		checkntids -= processed;
+		checks += processed;
+	}
+
+	*checkop->vmbuf = vmbuf;
+
+	if (checkntids == 0)
+	{
+		return;
+	}
+
+	upcomingvmbufchanges = 0;
+
+	for (int i = 1; i < checkntids; i++)
+	{
+		/*
+		 * Instead of storing the previous iteration's result, we only match
+		 * the block numbers
+		 */
+		BlockNumber lastblkno = checks[i - 1].tidblkno;
+		BlockNumber newblkno = checks[i].tidblkno;
+		/*
+		 * divide-by-constant can be faster than BufferGetBlockNumber()
+		 */
+		BlockNumber lastvmblkno = HEAPBLK_TO_VMBLOCK(lastblkno);
+		BlockNumber newvmblkno = HEAPBLK_TO_VMBLOCK(newblkno);
+
+		if (lastvmblkno != newvmblkno)
+			upcomingvmbufchanges++;
+	}
+
+	if (upcomingvmbufchanges <= pg_ceil_log2_32(checkntids))
+	{
+		/*
+		 * No big amount of VM buf changes, so do all visibility checks
+		 * without sorting.
+		 */
+		heap_ivc_process_all(rel, checkop->vmbuf, checks, checkntids);
+
+		return;
+	}
+
+	/*
+	 * Order the TIDs to heap order, so that we will only need to visit every
+	 * VM page at most once.
+	 */
+	heap_ivc_sortby_tidheapblk(checks, checkntids);
+
+	/* do all visibility checks */
+	heap_ivc_process_all(rel, checkop->vmbuf, checks, checkntids);
+
+	/* put the checks back in index order */
+	heap_ivc_sortby_idx(checks, checkntids);
+}
+
+
+static inline int
+heap_ivc_process_block(Relation rel, Buffer *vmbuf, TM_VisCheck *checks,
+					   int nchecks)
+{
+	BlockNumber	blkno;
+	BlockNumber	prevblkno = blkno = checks->tidblkno;
+	TMVC_Result	result;
+	int			processed = 0;
+
+	if (VM_ALL_VISIBLE(rel, blkno, vmbuf))
+		result = TMVC_Visible;
+	else
+		result = TMVC_MaybeVisible;
+
+	do
+	{
+		checks->vischeckresult = result;
+
+		nchecks--;
+		processed++;
+		checks++;
+
+		if (nchecks <= 0)
+			return processed;
+
+		blkno = checks->tidblkno;
+	} while (blkno == prevblkno);
+
+	return processed;
+}
+
+static void
+heap_ivc_process_all(Relation rel, Buffer *vmbuf,
+					 TM_VisCheck *checks, int nchecks)
+{
+	while (nchecks > 0)
+	{
+		int			processed;
+
+		processed = heap_ivc_process_block(rel, vmbuf, checks, nchecks);
+
+		Assert(processed <= nchecks);
+
+		nchecks -= processed;
+		checks += processed;
+	}
+}
+
 /*
  * Perform XLogInsert for a heap-visible operation.  'block' is the block
  * being marked all-visible, and vm_buffer is the buffer containing the

src/backend/access/heap/heapam_handler.c

@@ -2648,6 +2648,7 @@ static const TableAmRoutine heapam_methods = {
 	.tuple_tid_valid = heapam_tuple_tid_valid,
 	.tuple_satisfies_snapshot = heapam_tuple_satisfies_snapshot,
 	.index_delete_tuples = heap_index_delete_tuples,
+	.index_vischeck_tuples = heap_index_vischeck_tuples,
 
 	.relation_set_new_filelocator = heapam_relation_set_new_filelocator,
 	.relation_nontransactional_truncate = heapam_relation_nontransactional_truncate,

src/backend/access/heap/visibilitymap.c

@@ -107,17 +107,6 @@
  */
 #define MAPSIZE (BLCKSZ - MAXALIGN(SizeOfPageHeaderData))
 
-/* Number of heap blocks we can represent in one byte */
-#define HEAPBLOCKS_PER_BYTE (BITS_PER_BYTE / BITS_PER_HEAPBLOCK)
-
-/* Number of heap blocks we can represent in one visibility map page. */
-#define HEAPBLOCKS_PER_PAGE (MAPSIZE * HEAPBLOCKS_PER_BYTE)
-
-/* Mapping from heap block number to the right bit in the visibility map */
-#define HEAPBLK_TO_MAPBLOCK(x) ((x) / HEAPBLOCKS_PER_PAGE)
-#define HEAPBLK_TO_MAPBYTE(x) (((x) % HEAPBLOCKS_PER_PAGE) / HEAPBLOCKS_PER_BYTE)
-#define HEAPBLK_TO_OFFSET(x) (((x) % HEAPBLOCKS_PER_BYTE) * BITS_PER_HEAPBLOCK)
-
 /* Masks for counting subsets of bits in the visibility map. */
 #define VISIBLE_MASK8	(0x55)	/* The lower bit of each bit pair */
 #define FROZEN_MASK8	(0xaa)	/* The upper bit of each bit pair */
@@ -137,9 +126,9 @@ static Buffer vm_extend(Relation rel, BlockNumber vm_nblocks);
 bool
 visibilitymap_clear(Relation rel, BlockNumber heapBlk, Buffer vmbuf, uint8 flags)
 {
-	BlockNumber mapBlock = HEAPBLK_TO_MAPBLOCK(heapBlk);
-	int			mapByte = HEAPBLK_TO_MAPBYTE(heapBlk);
-	int			mapOffset = HEAPBLK_TO_OFFSET(heapBlk);
+	BlockNumber mapBlock = HEAPBLK_TO_VMBLOCK(heapBlk);
+	int			mapByte = HEAPBLK_TO_VMBYTE(heapBlk);
+	int			mapOffset = HEAPBLK_TO_VMOFFSET(heapBlk);
 	uint8		mask = flags << mapOffset;
 	char	   *map;
 	bool		cleared = false;
@@ -190,7 +179,7 @@ visibilitymap_clear(Relation rel, BlockNumber heapBlk, Buffer vmbuf, uint8 flags
 void
 visibilitymap_pin(Relation rel, BlockNumber heapBlk, Buffer *vmbuf)
 {
-	BlockNumber mapBlock = HEAPBLK_TO_MAPBLOCK(heapBlk);
+	BlockNumber mapBlock = HEAPBLK_TO_VMBLOCK(heapBlk);
 
 	/* Reuse the old pinned buffer if possible */
 	if (BufferIsValid(*vmbuf))
@@ -214,7 +203,7 @@ visibilitymap_pin(Relation rel, BlockNumber heapBlk, Buffer *vmbuf)
 bool
 visibilitymap_pin_ok(BlockNumber heapBlk, Buffer vmbuf)
 {
-	BlockNumber mapBlock = HEAPBLK_TO_MAPBLOCK(heapBlk);
+	BlockNumber mapBlock = HEAPBLK_TO_VMBLOCK(heapBlk);
 
 	return BufferIsValid(vmbuf) && BufferGetBlockNumber(vmbuf) == mapBlock;
 }
@@ -247,9 +236,9 @@ visibilitymap_set(Relation rel, BlockNumber heapBlk, Buffer heapBuf,
 				  XLogRecPtr recptr, Buffer vmBuf, TransactionId cutoff_xid,
 				  uint8 flags)
 {
-	BlockNumber mapBlock = HEAPBLK_TO_MAPBLOCK(heapBlk);
-	uint32		mapByte = HEAPBLK_TO_MAPBYTE(heapBlk);
-	uint8		mapOffset = HEAPBLK_TO_OFFSET(heapBlk);
+	BlockNumber mapBlock = HEAPBLK_TO_VMBLOCK(heapBlk);
+	uint32		mapByte = HEAPBLK_TO_VMBYTE(heapBlk);
+	uint8		mapOffset = HEAPBLK_TO_VMOFFSET(heapBlk);
 	Page		page;
 	uint8	   *map;
 	uint8		status;
@@ -340,9 +329,9 @@ visibilitymap_set(Relation rel, BlockNumber heapBlk, Buffer heapBuf,
 uint8
 visibilitymap_get_status(Relation rel, BlockNumber heapBlk, Buffer *vmbuf)
 {
-	BlockNumber mapBlock = HEAPBLK_TO_MAPBLOCK(heapBlk);
-	uint32		mapByte = HEAPBLK_TO_MAPBYTE(heapBlk);
-	uint8		mapOffset = HEAPBLK_TO_OFFSET(heapBlk);
+	BlockNumber mapBlock = HEAPBLK_TO_VMBLOCK(heapBlk);
+	uint32		mapByte = HEAPBLK_TO_VMBYTE(heapBlk);
+	uint8		mapOffset = HEAPBLK_TO_VMOFFSET(heapBlk);
 	char	   *map;
 	uint8		result;
 
@@ -445,9 +434,9 @@ visibilitymap_prepare_truncate(Relation rel, BlockNumber nheapblocks)
 	BlockNumber newnblocks;
 
 	/* last remaining block, byte, and bit */
-	BlockNumber truncBlock = HEAPBLK_TO_MAPBLOCK(nheapblocks);
-	uint32		truncByte = HEAPBLK_TO_MAPBYTE(nheapblocks);
-	uint8		truncOffset = HEAPBLK_TO_OFFSET(nheapblocks);
+	BlockNumber truncBlock = HEAPBLK_TO_VMBLOCK(nheapblocks);
+	uint32		truncByte = HEAPBLK_TO_VMBYTE(nheapblocks);
+	uint8		truncOffset = HEAPBLK_TO_VMOFFSET(nheapblocks);
 
 #ifdef TRACE_VISIBILITYMAP
 	elog(DEBUG1, "vm_truncate %s %d", RelationGetRelationName(rel), nheapblocks);

src/backend/access/index/indexam.c

@@ -628,6 +628,12 @@ index_getnext_tid(IndexScanDesc scan, ScanDirection direction)
 	/* XXX: we should assert that a snapshot is pushed or registered */
 	Assert(TransactionIdIsValid(RecentXmin));
 
+	/*
+	 * Reset xs_visrecheck, so we don't confuse the next tuple's visibility
+	 * state with that of the previous.
+	 */
+	scan->xs_visrecheck = TMVC_Unchecked;
+
 	/*
 	 * The AM's amgettuple proc finds the next index entry matching the scan
 	 * keys, and puts the TID into scan->xs_heaptid.  It should also set

src/backend/access/table/tableamapi.c

@@ -61,6 +61,7 @@ GetTableAmRoutine(Oid amhandler)
 	Assert(routine->tuple_get_latest_tid != NULL);
 	Assert(routine->tuple_satisfies_snapshot != NULL);
 	Assert(routine->index_delete_tuples != NULL);
+	Assert(routine->index_vischeck_tuples != NULL);
 
 	Assert(routine->tuple_insert != NULL);