From 455604a6fb9fcfd5d87a48ddbcc6dd90f3ea0192 Mon Sep 17 00:00:00 2001
From: Melanie Plageman <melanieplageman@gmail.com>
Date: Sat, 6 Apr 2024 12:57:10 -0400
Subject: [PATCH v19 12/21] BitmapHeapScan: Use UnifiedTBMIterator

BitmapHeapScan can be simplified by using the recently introduced
UnifiedTBMIterator interface for accessing shared or non-shared
TBMIterators.

This required a bit of reorganization of the !node->initialized path in
BitmapHeapNext().

Author: Melanie Plageman
Reviewed-by: Tomas Vondra
Discussion: https://postgr.es/m/CAAKRu_ZwCwWFeL_H3ia26bP2e7HiKLWt0ZmGXPVwPO6uXq0vaA%40mail.gmail.com
---
 src/backend/access/heap/heapam_handler.c  |   5 +-
 src/backend/executor/nodeBitmapHeapscan.c | 116 +++++++++-------------
 src/include/access/relscan.h              |   7 +-
 src/include/access/tableam.h              |  40 ++------
 src/include/nodes/execnodes.h             |   6 +-
 5 files changed, 57 insertions(+), 117 deletions(-)

diff --git a/src/backend/access/heap/heapam_handler.c b/src/backend/access/heap/heapam_handler.c
index 903cb80157e..ec6128e1d65 100644
--- a/src/backend/access/heap/heapam_handler.c
+++ b/src/backend/access/heap/heapam_handler.c
@@ -2208,10 +2208,7 @@ heapam_scan_bitmap_next_block(TableScanDesc scan,
 	{
 		CHECK_FOR_INTERRUPTS();
 
-		if (scan->shared_tbmiterator)
-			tbmres = tbm_shared_iterate(scan->shared_tbmiterator);
-		else
-			tbmres = tbm_iterate(scan->tbmiterator);
+		tbmres = unified_tbm_iterate(&scan->rs_bhs_iterator);
 
 		if (tbmres == NULL)
 		{
diff --git a/src/backend/executor/nodeBitmapHeapscan.c b/src/backend/executor/nodeBitmapHeapscan.c
index 23b21247dea..cd76d5e46ad 100644
--- a/src/backend/executor/nodeBitmapHeapscan.c
+++ b/src/backend/executor/nodeBitmapHeapscan.c
@@ -96,39 +96,28 @@ BitmapHeapNext(BitmapHeapScanState *node)
 	 */
 	if (!node->initialized)
 	{
-		TBMIterator *tbmiterator = NULL;
-		TBMSharedIterator *shared_tbmiterator = NULL;
+		/*
+		 * The leader will immediately come out of the function, but others
+		 * will be blocked until leader populates the TBM and wakes them up.
+		 * Serial bitmap table scan will simply build the bitmap.
+		 */
+		bool		init_shared_state = node->pstate ?
+			BitmapShouldInitializeSharedState(node->pstate) : false;
 
-		if (!pstate)
+		/*
+		 * Only serial bitmap table scans and the parallel leader in a
+		 * parallel bitmap table scan should build the bitmap.
+		 */
+		if (!pstate || init_shared_state)
 		{
 			tbm = (TIDBitmap *) MultiExecProcNode(outerPlanState(node));
 
 			if (!tbm || !IsA(tbm, TIDBitmap))
 				elog(ERROR, "unrecognized result from subplan");
-
 			node->tbm = tbm;
-			tbmiterator = tbm_begin_iterate(tbm);
 
-#ifdef USE_PREFETCH
-			if (node->prefetch_maximum > 0)
-				node->prefetch_iterator = tbm_begin_iterate(tbm);
-#endif							/* USE_PREFETCH */
-		}
-		else
-		{
-			/*
-			 * The leader will immediately come out of the function, but
-			 * others will be blocked until leader populates the TBM and wakes
-			 * them up.
-			 */
-			if (BitmapShouldInitializeSharedState(pstate))
+			if (init_shared_state)
 			{
-				tbm = (TIDBitmap *) MultiExecProcNode(outerPlanState(node));
-				if (!tbm || !IsA(tbm, TIDBitmap))
-					elog(ERROR, "unrecognized result from subplan");
-
-				node->tbm = tbm;
-
 				/*
 				 * Prepare to iterate over the TBM. This will return the
 				 * dsa_pointer of the iterator state which will be used by
@@ -142,21 +131,9 @@ BitmapHeapNext(BitmapHeapScanState *node)
 						tbm_prepare_shared_iterate(tbm);
 				}
 #endif
-
 				/* We have initialized the shared state so wake up others. */
 				BitmapDoneInitializingSharedState(pstate);
 			}
-
-			/* Allocate a private iterator and attach the shared state to it */
-			shared_tbmiterator = tbm_attach_shared_iterate(dsa, pstate->tbmiterator);
-
-#ifdef USE_PREFETCH
-			if (node->prefetch_maximum > 0)
-			{
-				node->shared_prefetch_iterator =
-					tbm_attach_shared_iterate(dsa, pstate->prefetch_iterator);
-			}
-#endif							/* USE_PREFETCH */
 		}
 
 		/*
@@ -187,8 +164,21 @@ BitmapHeapNext(BitmapHeapScanState *node)
 			node->ss.ss_currentScanDesc = scan;
 		}
 
-		scan->tbmiterator = tbmiterator;
-		scan->shared_tbmiterator = shared_tbmiterator;
+		unified_tbm_begin_iterate(&scan->rs_bhs_iterator, tbm, dsa,
+								  pstate ?
+								  pstate->tbmiterator :
+								  InvalidDsaPointer);
+
+#ifdef USE_PREFETCH
+		if (node->prefetch_maximum > 0)
+		{
+			unified_tbm_begin_iterate(&node->prefetch_iterator, tbm, dsa,
+									  pstate ?
+									  pstate->prefetch_iterator :
+									  InvalidDsaPointer);
+		}
+#endif							/* USE_PREFETCH */
+
 		node->initialized = true;
 
 		goto new_page;
@@ -268,7 +258,7 @@ new_page:
 		 * ahead of the current block.
 		 */
 		if (node->pstate == NULL &&
-			node->prefetch_iterator &&
+			!node->prefetch_iterator.exhausted &&
 			node->pfblockno < node->blockno)
 			elog(ERROR, "prefetch and main iterators are out of sync");
 
@@ -309,21 +299,20 @@ BitmapAdjustPrefetchIterator(BitmapHeapScanState *node)
 {
 #ifdef USE_PREFETCH
 	ParallelBitmapHeapState *pstate = node->pstate;
+	UnifiedTBMIterator *prefetch_iterator = &node->prefetch_iterator;
 	TBMIterateResult *tbmpre;
 
 	if (pstate == NULL)
 	{
-		TBMIterator *prefetch_iterator = node->prefetch_iterator;
-
 		if (node->prefetch_pages > 0)
 		{
 			/* The main iterator has closed the distance by one page */
 			node->prefetch_pages--;
 		}
-		else if (prefetch_iterator)
+		else if (!prefetch_iterator->exhausted)
 		{
 			/* Do not let the prefetch iterator get behind the main one */
-			tbmpre = tbm_iterate(prefetch_iterator);
+			tbmpre = unified_tbm_iterate(prefetch_iterator);
 			node->pfblockno = tbmpre ? tbmpre->blockno : InvalidBlockNumber;
 		}
 		return;
@@ -340,8 +329,6 @@ BitmapAdjustPrefetchIterator(BitmapHeapScanState *node)
 	 */
 	if (node->prefetch_maximum > 0)
 	{
-		TBMSharedIterator *prefetch_iterator = node->shared_prefetch_iterator;
-
 		SpinLockAcquire(&pstate->mutex);
 		if (pstate->prefetch_pages > 0)
 		{
@@ -361,9 +348,9 @@ BitmapAdjustPrefetchIterator(BitmapHeapScanState *node)
 			 * we don't validate the blockno here as we do in non-parallel
 			 * case.
 			 */
-			if (prefetch_iterator)
+			if (!prefetch_iterator->exhausted)
 			{
-				tbmpre = tbm_shared_iterate(prefetch_iterator);
+				tbmpre = unified_tbm_iterate(prefetch_iterator);
 				node->pfblockno = tbmpre ? tbmpre->blockno : InvalidBlockNumber;
 			}
 		}
@@ -423,23 +410,21 @@ BitmapPrefetch(BitmapHeapScanState *node, TableScanDesc scan)
 {
 #ifdef USE_PREFETCH
 	ParallelBitmapHeapState *pstate = node->pstate;
+	UnifiedTBMIterator *prefetch_iterator = &node->prefetch_iterator;
 
 	if (pstate == NULL)
 	{
-		TBMIterator *prefetch_iterator = node->prefetch_iterator;
-
-		if (prefetch_iterator)
+		if (!prefetch_iterator->exhausted)
 		{
 			while (node->prefetch_pages < node->prefetch_target)
 			{
-				TBMIterateResult *tbmpre = tbm_iterate(prefetch_iterator);
+				TBMIterateResult *tbmpre = unified_tbm_iterate(prefetch_iterator);
 				bool		skip_fetch;
 
 				if (tbmpre == NULL)
 				{
 					/* No more pages to prefetch */
-					tbm_end_iterate(prefetch_iterator);
-					node->prefetch_iterator = NULL;
+					unified_tbm_end_iterate(prefetch_iterator);
 					break;
 				}
 				node->prefetch_pages++;
@@ -467,9 +452,7 @@ BitmapPrefetch(BitmapHeapScanState *node, TableScanDesc scan)
 
 	if (pstate->prefetch_pages < pstate->prefetch_target)
 	{
-		TBMSharedIterator *prefetch_iterator = node->shared_prefetch_iterator;
-
-		if (prefetch_iterator)
+		if (!prefetch_iterator->exhausted)
 		{
 			while (1)
 			{
@@ -492,12 +475,11 @@ BitmapPrefetch(BitmapHeapScanState *node, TableScanDesc scan)
 				if (!do_prefetch)
 					return;
 
-				tbmpre = tbm_shared_iterate(prefetch_iterator);
+				tbmpre = unified_tbm_iterate(prefetch_iterator);
 				if (tbmpre == NULL)
 				{
 					/* No more pages to prefetch */
-					tbm_end_shared_iterate(prefetch_iterator);
-					node->shared_prefetch_iterator = NULL;
+					unified_tbm_end_iterate(prefetch_iterator);
 					break;
 				}
 
@@ -564,18 +546,14 @@ ExecReScanBitmapHeapScan(BitmapHeapScanState *node)
 		table_rescan(node->ss.ss_currentScanDesc, NULL);
 
 	/* release bitmaps and buffers if any */
-	if (node->prefetch_iterator)
-		tbm_end_iterate(node->prefetch_iterator);
-	if (node->shared_prefetch_iterator)
-		tbm_end_shared_iterate(node->shared_prefetch_iterator);
+	if (node->prefetch_iterator.exhausted)
+		unified_tbm_end_iterate(&node->prefetch_iterator);
 	if (node->tbm)
 		tbm_free(node->tbm);
 	if (node->pvmbuffer != InvalidBuffer)
 		ReleaseBuffer(node->pvmbuffer);
 	node->tbm = NULL;
-	node->prefetch_iterator = NULL;
 	node->initialized = false;
-	node->shared_prefetch_iterator = NULL;
 	node->pvmbuffer = InvalidBuffer;
 	node->recheck = true;
 	node->blockno = InvalidBlockNumber;
@@ -623,12 +601,10 @@ ExecEndBitmapHeapScan(BitmapHeapScanState *node)
 	/*
 	 * release bitmaps and buffers if any
 	 */
-	if (node->prefetch_iterator)
-		tbm_end_iterate(node->prefetch_iterator);
+	if (!node->prefetch_iterator.exhausted)
+		unified_tbm_end_iterate(&node->prefetch_iterator);
 	if (node->tbm)
 		tbm_free(node->tbm);
-	if (node->shared_prefetch_iterator)
-		tbm_end_shared_iterate(node->shared_prefetch_iterator);
 	if (node->pvmbuffer != InvalidBuffer)
 		ReleaseBuffer(node->pvmbuffer);
 }
@@ -666,11 +642,9 @@ ExecInitBitmapHeapScan(BitmapHeapScan *node, EState *estate, int eflags)
 	scanstate->pvmbuffer = InvalidBuffer;
 	scanstate->exact_pages = 0;
 	scanstate->lossy_pages = 0;
-	scanstate->prefetch_iterator = NULL;
 	scanstate->prefetch_pages = 0;
 	scanstate->prefetch_target = -1;
 	scanstate->initialized = false;
-	scanstate->shared_prefetch_iterator = NULL;
 	scanstate->pstate = NULL;
 	scanstate->recheck = true;
 	scanstate->blockno = InvalidBlockNumber;
diff --git a/src/include/access/relscan.h b/src/include/access/relscan.h
index 92b829cebc7..7fec4ff8acb 100644
--- a/src/include/access/relscan.h
+++ b/src/include/access/relscan.h
@@ -20,13 +20,11 @@
 #include "storage/buf.h"
 #include "storage/spin.h"
 #include "utils/relcache.h"
+#include "nodes/tidbitmap.h"
 
 
 struct ParallelTableScanDescData;
 
-struct TBMIterator;
-struct TBMSharedIterator;
-
 /*
  * Generic descriptor for table scans. This is the base-class for table scans,
  * which needs to be embedded in the scans of individual AMs.
@@ -44,8 +42,7 @@ typedef struct TableScanDescData
 	ItemPointerData rs_maxtid;
 
 	/* Only used for Bitmap table scans */
-	struct TBMIterator *tbmiterator;
-	struct TBMSharedIterator *shared_tbmiterator;
+	UnifiedTBMIterator rs_bhs_iterator;
 
 	/*
 	 * Information about type and behaviour of the scan, a bitmask of members
diff --git a/src/include/access/tableam.h b/src/include/access/tableam.h
index 68a479496d7..b4b7fa13948 100644
--- a/src/include/access/tableam.h
+++ b/src/include/access/tableam.h
@@ -938,16 +938,12 @@ static inline TableScanDesc
 table_beginscan_bm(Relation rel, Snapshot snapshot,
 				   int nkeys, struct ScanKeyData *key, bool need_tuple)
 {
-	TableScanDesc result;
 	uint32		flags = SO_TYPE_BITMAPSCAN | SO_ALLOW_PAGEMODE;
 
 	if (need_tuple)
 		flags |= SO_NEED_TUPLES;
 
-	result = rel->rd_tableam->scan_begin(rel, snapshot, nkeys, key, NULL, flags);
-	result->shared_tbmiterator = NULL;
-	result->tbmiterator = NULL;
-	return result;
+	return rel->rd_tableam->scan_begin(rel, snapshot, nkeys, key, NULL, flags);
 }
 
 /*
@@ -994,20 +990,9 @@ table_beginscan_tid(Relation rel, Snapshot snapshot)
 static inline void
 table_endscan(TableScanDesc scan)
 {
-	if (scan->rs_flags & SO_TYPE_BITMAPSCAN)
-	{
-		if (scan->shared_tbmiterator)
-		{
-			tbm_end_shared_iterate(scan->shared_tbmiterator);
-			scan->shared_tbmiterator = NULL;
-		}
-
-		if (scan->tbmiterator)
-		{
-			tbm_end_iterate(scan->tbmiterator);
-			scan->tbmiterator = NULL;
-		}
-	}
+	if (scan->rs_flags & SO_TYPE_BITMAPSCAN &&
+		!scan->rs_bhs_iterator.exhausted)
+		unified_tbm_end_iterate(&scan->rs_bhs_iterator);
 
 	scan->rs_rd->rd_tableam->scan_end(scan);
 }
@@ -1019,20 +1004,9 @@ static inline void
 table_rescan(TableScanDesc scan,
 			 struct ScanKeyData *key)
 {
-	if (scan->rs_flags & SO_TYPE_BITMAPSCAN)
-	{
-		if (scan->shared_tbmiterator)
-		{
-			tbm_end_shared_iterate(scan->shared_tbmiterator);
-			scan->shared_tbmiterator = NULL;
-		}
-
-		if (scan->tbmiterator)
-		{
-			tbm_end_iterate(scan->tbmiterator);
-			scan->tbmiterator = NULL;
-		}
-	}
+	if (scan->rs_flags & SO_TYPE_BITMAPSCAN &&
+		!scan->rs_bhs_iterator.exhausted)
+		unified_tbm_end_iterate(&scan->rs_bhs_iterator);
 
 	scan->rs_rd->rd_tableam->scan_rescan(scan, key, false, false, false, false);
 }
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
index d96703b04d4..6693cf66c77 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -1795,12 +1795,11 @@ typedef struct ParallelBitmapHeapState
  *		pvmbuffer		   buffer for visibility-map lookups of prefetched pages
  *		exact_pages		   total number of exact pages retrieved
  *		lossy_pages		   total number of lossy pages retrieved
- *		prefetch_iterator  iterator for prefetching ahead of current page
+ *		prefetch_iterator  for prefetching ahead of current page
  *		prefetch_pages	   # pages prefetch iterator is ahead of current
  *		prefetch_target    current target prefetch distance
  *		prefetch_maximum   maximum value for prefetch_target
  *		initialized		   is node is ready to iterate
- *		shared_prefetch_iterator shared iterator for prefetching
  *		pstate			   shared state for parallel bitmap scan
  *		recheck			   do current page's tuples need recheck
  *		blockno			   used to validate pf and current block in sync
@@ -1815,12 +1814,11 @@ typedef struct BitmapHeapScanState
 	Buffer		pvmbuffer;
 	long		exact_pages;
 	long		lossy_pages;
-	TBMIterator *prefetch_iterator;
 	int			prefetch_pages;
 	int			prefetch_target;
 	int			prefetch_maximum;
 	bool		initialized;
-	TBMSharedIterator *shared_prefetch_iterator;
+	UnifiedTBMIterator prefetch_iterator;
 	ParallelBitmapHeapState *pstate;
 	bool		recheck;
 	BlockNumber blockno;
-- 
2.40.1

