From faf85685fdd454d88e962e621f1b22ea686a8db2 Mon Sep 17 00:00:00 2001
From: Melanie Plageman <melanieplageman@gmail.com>
Date: Tue, 21 May 2024 13:39:04 -0400
Subject: [PATCH v22 06/19] BitmapHeapScan uses unified iterator

With the repurposing of TBMIterator as an interface for both parallel
and serial iteration through TIDBitmaps, BitmapHeapScan may now use it
since it has both parallel and serial iteration.
---
 src/backend/access/heap/heapam_handler.c  |   5 +-
 src/backend/executor/nodeBitmapHeapscan.c | 132 ++++++++--------------
 src/include/access/relscan.h              |   7 +-
 src/include/access/tableam.h              |  34 +-----
 src/include/nodes/execnodes.h             |   4 +-
 5 files changed, 54 insertions(+), 128 deletions(-)

diff --git a/src/backend/access/heap/heapam_handler.c b/src/backend/access/heap/heapam_handler.c
index b6ddf600a0a..c3788918777 100644
--- a/src/backend/access/heap/heapam_handler.c
+++ b/src/backend/access/heap/heapam_handler.c
@@ -2140,10 +2140,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_private_iterate(scan->tbmiterator);
+		tbmres = tbm_iterate(&scan->tbmiterator);
 
 		if (tbmres == NULL)
 			return false;
diff --git a/src/backend/executor/nodeBitmapHeapscan.c b/src/backend/executor/nodeBitmapHeapscan.c
index 5eb54f62f20..9c156787770 100644
--- a/src/backend/executor/nodeBitmapHeapscan.c
+++ b/src/backend/executor/nodeBitmapHeapscan.c
@@ -69,7 +69,6 @@ BitmapHeapNext(BitmapHeapScanState *node)
 {
 	ExprContext *econtext;
 	TableScanDesc scan;
-	TIDBitmap  *tbm;
 	TupleTableSlot *slot;
 	ParallelBitmapHeapState *pstate = node->pstate;
 	dsa_area   *dsa = node->ss.ps.state->es_query_dsa;
@@ -80,7 +79,6 @@ BitmapHeapNext(BitmapHeapScanState *node)
 	econtext = node->ss.ps.ps_ExprContext;
 	slot = node->ss.ss_ScanTupleSlot;
 	scan = node->ss.ss_currentScanDesc;
-	tbm = node->tbm;
 
 	/*
 	 * If we haven't yet performed the underlying index scan, do it, and begin
@@ -96,76 +94,49 @@ BitmapHeapNext(BitmapHeapScanState *node)
 	 */
 	if (!node->initialized)
 	{
-		TBMPrivateIterator *tbmiterator = NULL;
-		TBMSharedIterator *shared_tbmiterator = NULL;
-
 		if (!pstate)
 		{
-			tbm = (TIDBitmap *) MultiExecProcNode(outerPlanState(node));
+			node->tbm = (TIDBitmap *) MultiExecProcNode(outerPlanState(node));
 
-			if (!tbm || !IsA(tbm, TIDBitmap))
+			if (!node->tbm || !IsA(node->tbm, TIDBitmap))
 				elog(ERROR, "unrecognized result from subplan");
-
-			node->tbm = tbm;
-			tbmiterator = tbm_begin_private_iterate(tbm);
-
-#ifdef USE_PREFETCH
-			if (node->prefetch_maximum > 0)
-				node->prefetch_iterator = tbm_begin_private_iterate(tbm);
-#endif							/* USE_PREFETCH */
 		}
-		else
+		else if (BitmapShouldInitializeSharedState(pstate))
 		{
 			/*
 			 * 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))
-			{
-				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
-				 * multiple processes to iterate jointly.
-				 */
-				pstate->tbmiterator = tbm_prepare_shared_iterate(tbm);
-#ifdef USE_PREFETCH
-				if (node->prefetch_maximum > 0)
-				{
-					pstate->prefetch_iterator =
-						tbm_prepare_shared_iterate(tbm);
-
-					/*
-					 * We don't need the mutex here as we haven't yet woke up
-					 * others.
-					 */
-					pstate->prefetch_pages = 0;
-					pstate->prefetch_target = -1;
-				}
-#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);
+			node->tbm = (TIDBitmap *) MultiExecProcNode(outerPlanState(node));
+			if (!node->tbm || !IsA(node->tbm, TIDBitmap))
+				elog(ERROR, "unrecognized result from subplan");
 
+			/*
+			 * Prepare to iterate over the TBM. This will return the
+			 * dsa_pointer of the iterator state which will be used by
+			 * multiple processes to iterate jointly.
+			 */
+			pstate->tbmiterator = tbm_prepare_shared_iterate(node->tbm);
 #ifdef USE_PREFETCH
 			if (node->prefetch_maximum > 0)
 			{
-				node->shared_prefetch_iterator =
-					tbm_attach_shared_iterate(dsa, pstate->prefetch_iterator);
+				pstate->prefetch_iterator =
+					tbm_prepare_shared_iterate(node->tbm);
 			}
-#endif							/* USE_PREFETCH */
+#endif
+			/* We have initialized the shared state so wake up others. */
+			BitmapDoneInitializingSharedState(pstate);
 		}
 
+#ifdef USE_PREFETCH
+		if (node->prefetch_maximum > 0)
+			tbm_begin_iterate(&node->prefetch_iterator, node->tbm, dsa,
+							  pstate ?
+							  pstate->prefetch_iterator :
+							  InvalidDsaPointer);
+#endif							/* USE_PREFETCH */
+
 		/*
 		 * If this is the first scan of the underlying table, create the table
 		 * scan descriptor and begin the scan.
@@ -194,8 +165,11 @@ BitmapHeapNext(BitmapHeapScanState *node)
 			node->ss.ss_currentScanDesc = scan;
 		}
 
-		scan->tbmiterator = tbmiterator;
-		scan->shared_tbmiterator = shared_tbmiterator;
+		tbm_begin_iterate(&scan->tbmiterator, node->tbm, dsa,
+						  pstate ?
+						  pstate->tbmiterator :
+						  InvalidDsaPointer);
+
 		node->initialized = true;
 
 		goto new_page;
@@ -275,7 +249,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");
 
@@ -320,17 +294,17 @@ BitmapAdjustPrefetchIterator(BitmapHeapScanState *node)
 
 	if (pstate == NULL)
 	{
-		TBMPrivateIterator *prefetch_iterator = node->prefetch_iterator;
+		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_private_iterate(prefetch_iterator);
+			tbmpre = tbm_iterate(prefetch_iterator);
 			node->pfblockno = tbmpre ? tbmpre->blockno : InvalidBlockNumber;
 		}
 		return;
@@ -347,7 +321,7 @@ BitmapAdjustPrefetchIterator(BitmapHeapScanState *node)
 	 */
 	if (node->prefetch_maximum > 0)
 	{
-		TBMSharedIterator *prefetch_iterator = node->shared_prefetch_iterator;
+		TBMIterator *prefetch_iterator = &node->prefetch_iterator;
 
 		SpinLockAcquire(&pstate->mutex);
 		if (pstate->prefetch_pages > 0)
@@ -368,9 +342,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 = tbm_iterate(prefetch_iterator);
 				node->pfblockno = tbmpre ? tbmpre->blockno : InvalidBlockNumber;
 			}
 		}
@@ -433,20 +407,19 @@ BitmapPrefetch(BitmapHeapScanState *node, TableScanDesc scan)
 
 	if (pstate == NULL)
 	{
-		TBMPrivateIterator *prefetch_iterator = node->prefetch_iterator;
+		TBMIterator *prefetch_iterator = &node->prefetch_iterator;
 
-		if (prefetch_iterator)
+		if (!prefetch_iterator->exhausted)
 		{
 			while (node->prefetch_pages < node->prefetch_target)
 			{
-				TBMIterateResult *tbmpre = tbm_private_iterate(prefetch_iterator);
+				TBMIterateResult *tbmpre = tbm_iterate(prefetch_iterator);
 				bool		skip_fetch;
 
 				if (tbmpre == NULL)
 				{
 					/* No more pages to prefetch */
-					tbm_end_private_iterate(prefetch_iterator);
-					node->prefetch_iterator = NULL;
+					tbm_end_iterate(prefetch_iterator);
 					break;
 				}
 				node->prefetch_pages++;
@@ -474,9 +447,9 @@ BitmapPrefetch(BitmapHeapScanState *node, TableScanDesc scan)
 
 	if (pstate->prefetch_pages < pstate->prefetch_target)
 	{
-		TBMSharedIterator *prefetch_iterator = node->shared_prefetch_iterator;
+		TBMIterator *prefetch_iterator = &node->prefetch_iterator;
 
-		if (prefetch_iterator)
+		if (!prefetch_iterator->exhausted)
 		{
 			while (1)
 			{
@@ -499,12 +472,11 @@ BitmapPrefetch(BitmapHeapScanState *node, TableScanDesc scan)
 				if (!do_prefetch)
 					return;
 
-				tbmpre = tbm_shared_iterate(prefetch_iterator);
+				tbmpre = tbm_iterate(prefetch_iterator);
 				if (tbmpre == NULL)
 				{
 					/* No more pages to prefetch */
-					tbm_end_shared_iterate(prefetch_iterator);
-					node->shared_prefetch_iterator = NULL;
+					tbm_end_iterate(prefetch_iterator);
 					break;
 				}
 
@@ -571,18 +543,13 @@ ExecReScanBitmapHeapScan(BitmapHeapScanState *node)
 		table_rescan(node->ss.ss_currentScanDesc, NULL);
 
 	/* release bitmaps and buffers if any */
-	if (node->prefetch_iterator)
-		tbm_end_private_iterate(node->prefetch_iterator);
-	if (node->shared_prefetch_iterator)
-		tbm_end_shared_iterate(node->shared_prefetch_iterator);
+	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;
@@ -630,12 +597,9 @@ ExecEndBitmapHeapScan(BitmapHeapScanState *node)
 	/*
 	 * release bitmaps and buffers if any
 	 */
-	if (node->prefetch_iterator)
-		tbm_end_private_iterate(node->prefetch_iterator);
+	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);
 }
@@ -673,11 +637,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 159921ffb50..b4511677b1d 100644
--- a/src/include/access/relscan.h
+++ b/src/include/access/relscan.h
@@ -16,6 +16,7 @@
 
 #include "access/htup_details.h"
 #include "access/itup.h"
+#include "nodes/tidbitmap.h"
 #include "port/atomics.h"
 #include "storage/buf.h"
 #include "storage/spin.h"
@@ -24,9 +25,6 @@
 
 struct ParallelTableScanDescData;
 
-struct TBMPrivateIterator;
-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.
@@ -40,8 +38,7 @@ typedef struct TableScanDescData
 	struct ScanKeyData *rs_key; /* array of scan key descriptors */
 
 	/* Iterators for Bitmap Table Scans */
-	struct TBMPrivateIterator *tbmiterator;
-	struct TBMSharedIterator *shared_tbmiterator;
+	TBMIterator tbmiterator;
 
 	/* Range of ItemPointers for table_scan_getnextslot_tidrange() to scan. */
 	ItemPointerData rs_mintid;
diff --git a/src/include/access/tableam.h b/src/include/access/tableam.h
index 9802d42df9d..7034b79c2de 100644
--- a/src/include/access/tableam.h
+++ b/src/include/access/tableam.h
@@ -954,16 +954,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);
 }
 
 /*
@@ -1024,19 +1020,7 @@ 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_private_iterate(scan->tbmiterator);
-			scan->tbmiterator = NULL;
-		}
-	}
+		tbm_end_iterate(&scan->tbmiterator);
 
 	scan->rs_rd->rd_tableam->scan_end(scan);
 }
@@ -1049,19 +1033,7 @@ 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_private_iterate(scan->tbmiterator);
-			scan->tbmiterator = NULL;
-		}
-	}
+		tbm_end_iterate(&scan->tbmiterator);
 
 	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 81447fea8c5..1f262db776c 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -1803,7 +1803,6 @@ typedef struct ParallelBitmapHeapState
  *		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
@@ -1818,12 +1817,11 @@ typedef struct BitmapHeapScanState
 	Buffer		pvmbuffer;
 	long		exact_pages;
 	long		lossy_pages;
-	TBMPrivateIterator *prefetch_iterator;
 	int			prefetch_pages;
 	int			prefetch_target;
 	int			prefetch_maximum;
 	bool		initialized;
-	TBMSharedIterator *shared_prefetch_iterator;
+	TBMIterator prefetch_iterator;
 	ParallelBitmapHeapState *pstate;
 	bool		recheck;
 	BlockNumber blockno;
-- 
2.34.1

