From 0f5446c886e2750b8c1a4f99ca5f4203ee56a830 Mon Sep 17 00:00:00 2001
From: Melanie Plageman <melanieplageman@gmail.com>
Date: Tue, 21 May 2024 13:39:04 -0400
Subject: [PATCH v25 05/11] Bitmap Table Scans use unified TBMIterator

With the repurposing of TBMIterator as an interface for both parallel
and serial iteration through TIDBitmaps in commit [XXX: put commit
here], bitmap table scans may now use it.

Modify bitmap table scan code to use the TBMIterator. This requires
moving around a bit of code, so a few variables are initialized
elsewhere.

Author: Melanie Plageman
---
 src/backend/access/heap/heapam_handler.c  |   5 +-
 src/backend/executor/nodeBitmapHeapscan.c | 170 +++++++++-------------
 src/include/access/relscan.h              |  12 +-
 src/include/access/tableam.h              |   8 +-
 src/include/nodes/execnodes.h             |   4 +-
 src/include/nodes/tidbitmap.h             |  10 ++
 6 files changed, 86 insertions(+), 123 deletions(-)

diff --git a/src/backend/access/heap/heapam_handler.c b/src/backend/access/heap/heapam_handler.c
index a85edd60d78..c51f5e921f1 100644
--- a/src/backend/access/heap/heapam_handler.c
+++ b/src/backend/access/heap/heapam_handler.c
@@ -2135,10 +2135,7 @@ heapam_scan_bitmap_next_block(TableScanDesc scan,
 	{
 		CHECK_FOR_INTERRUPTS();
 
-		if (scan->st.bitmap.rs_shared_iterator)
-			tbmres = tbm_shared_iterate(scan->st.bitmap.rs_shared_iterator);
-		else
-			tbmres = tbm_private_iterate(scan->st.bitmap.rs_iterator);
+		tbmres = tbm_iterate(&scan->st.rs_tbmiterator);
 
 		if (tbmres == NULL)
 			return false;
diff --git a/src/backend/executor/nodeBitmapHeapscan.c b/src/backend/executor/nodeBitmapHeapscan.c
index 5fdbe1315b8..b35d09c2ef3 100644
--- a/src/backend/executor/nodeBitmapHeapscan.c
+++ b/src/backend/executor/nodeBitmapHeapscan.c
@@ -96,8 +96,7 @@ BitmapHeapNext(BitmapHeapScanState *node)
 	 */
 	if (!node->initialized)
 	{
-		TBMPrivateIterator *tbmiterator = NULL;
-		TBMSharedIterator *shared_tbmiterator = NULL;
+		TBMIterator tbmiterator;
 
 		if (!pstate)
 		{
@@ -107,70 +106,62 @@ BitmapHeapNext(BitmapHeapScanState *node)
 				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);
-				node->prefetch_pages = 0;
-				node->prefetch_target = -1;
-			}
+				node->prefetch_iterator =
+					tbm_begin_iterate(node->tbm, dsa,
+									  pstate ?
+									  pstate->prefetch_iterator :
+									  InvalidDsaPointer);
 #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
+			tbm = (TIDBitmap *) MultiExecProcNode(outerPlanState(node));
+			if (!tbm || !IsA(tbm, TIDBitmap))
+				elog(ERROR, "unrecognized result from subplan");
 
-				/* We have initialized the shared state so wake up others. */
-				BitmapDoneInitializingSharedState(pstate);
-			}
+			node->tbm = tbm;
 
-			/* Allocate a private iterator and attach the shared state to it */
-			shared_tbmiterator =
-				tbm_attach_shared_iterate(dsa, pstate->tbmiterator);
+			/*
+			 * 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)
 			{
-				node->shared_prefetch_iterator =
-					tbm_attach_shared_iterate(dsa, pstate->prefetch_iterator);
+				pstate->prefetch_iterator =
+					tbm_prepare_shared_iterate(tbm);
 			}
 #endif							/* USE_PREFETCH */
+
+			/* We have initialized the shared state so wake up others. */
+			BitmapDoneInitializingSharedState(pstate);
 		}
 
+		tbmiterator = tbm_begin_iterate(tbm, dsa,
+										pstate ?
+										pstate->tbmiterator :
+										InvalidDsaPointer);
+
+#ifdef USE_PREFETCH
+		if (node->prefetch_maximum > 0)
+			node->prefetch_iterator =
+				tbm_begin_iterate(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.
@@ -199,8 +190,7 @@ BitmapHeapNext(BitmapHeapScanState *node)
 			node->ss.ss_currentScanDesc = scan;
 		}
 
-		scan->st.bitmap.rs_iterator = tbmiterator;
-		scan->st.bitmap.rs_shared_iterator = shared_tbmiterator;
+		scan->st.rs_tbmiterator = tbmiterator;
 		node->initialized = true;
 
 		goto new_page;
@@ -286,7 +276,7 @@ new_page:
 		 * ahead of the current block.
 		 */
 		if (node->pstate == NULL &&
-			node->prefetch_iterator &&
+			!tbm_exhausted(&node->prefetch_iterator) &&
 			node->prefetch_blockno < node->blockno)
 			elog(ERROR,
 				 "prefetch and main iterators are out of sync. pfblockno: %d. blockno: %d",
@@ -333,16 +323,16 @@ 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 (!tbm_exhausted(prefetch_iterator))
 		{
-			tbmpre = tbm_private_iterate(prefetch_iterator);
+			tbmpre = tbm_iterate(prefetch_iterator);
 			node->prefetch_blockno = tbmpre ? tbmpre->blockno :
 				InvalidBlockNumber;
 		}
@@ -360,7 +350,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)
@@ -381,9 +371,9 @@ BitmapAdjustPrefetchIterator(BitmapHeapScanState *node)
 			 * we don't validate the blockno here as we do in non-parallel
 			 * case.
 			 */
-			if (prefetch_iterator)
+			if (!tbm_exhausted(prefetch_iterator))
 			{
-				tbmpre = tbm_shared_iterate(prefetch_iterator);
+				tbmpre = tbm_iterate(prefetch_iterator);
 				node->prefetch_blockno = tbmpre ? tbmpre->blockno :
 					InvalidBlockNumber;
 			}
@@ -447,21 +437,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 (!tbm_exhausted(prefetch_iterator))
 		{
 			while (node->prefetch_pages < node->prefetch_target)
 			{
-				TBMIterateResult *tbmpre;
+				TBMIterateResult *tbmpre = tbm_iterate(prefetch_iterator);
 				bool		skip_fetch;
 
-				tbmpre = tbm_private_iterate(prefetch_iterator);
 				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++;
@@ -489,9 +477,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 (!tbm_exhausted(prefetch_iterator))
 		{
 			while (1)
 			{
@@ -514,12 +502,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;
 				}
 
@@ -588,39 +575,30 @@ ExecReScanBitmapHeapScan(BitmapHeapScanState *node)
 		/*
 		 * End iteration on iterators saved in scan descriptor.
 		 */
-		if (scan->st.bitmap.rs_shared_iterator)
-		{
-			tbm_end_shared_iterate(scan->st.bitmap.rs_shared_iterator);
-			scan->st.bitmap.rs_shared_iterator = NULL;
-		}
-
-		if (scan->st.bitmap.rs_iterator)
-		{
-			tbm_end_private_iterate(scan->st.bitmap.rs_iterator);
-			scan->st.bitmap.rs_iterator = NULL;
-		}
+		tbm_end_iterate(&scan->st.rs_tbmiterator);
 
 		/* rescan to release any page pin */
 		table_rescan(node->ss.ss_currentScanDesc, NULL);
 	}
 
+	/* If we did not already clean up the prefetch iterator, do so now. */
+	if (!tbm_exhausted(&node->prefetch_iterator))
+		tbm_end_iterate(&node->prefetch_iterator);
+
 	/* 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);
 	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;
+	/* Only used for serial BHS */
 	node->blockno = InvalidBlockNumber;
 	node->prefetch_blockno = InvalidBlockNumber;
+	node->prefetch_pages = 0;
+	node->prefetch_target = -1;
 
 	ExecScanReScan(&node->ss);
 
@@ -679,17 +657,7 @@ ExecEndBitmapHeapScan(BitmapHeapScanState *node)
 		/*
 		 * End iteration on iterators saved in scan descriptor.
 		 */
-		if (scanDesc->st.bitmap.rs_shared_iterator)
-		{
-			tbm_end_shared_iterate(scanDesc->st.bitmap.rs_shared_iterator);
-			scanDesc->st.bitmap.rs_shared_iterator = NULL;
-		}
-
-		if (scanDesc->st.bitmap.rs_iterator)
-		{
-			tbm_end_private_iterate(scanDesc->st.bitmap.rs_iterator);
-			scanDesc->st.bitmap.rs_iterator = NULL;
-		}
+		tbm_end_iterate(&scanDesc->st.rs_tbmiterator);
 
 		/*
 		 * close table scan
@@ -697,15 +665,15 @@ ExecEndBitmapHeapScan(BitmapHeapScanState *node)
 		table_endscan(scanDesc);
 	}
 
+	/* If we did not already clean up the prefetch iterator, do so now. */
+	if (!tbm_exhausted(&node->prefetch_iterator))
+		tbm_end_iterate(&node->prefetch_iterator);
+
 	/*
 	 * release bitmaps and buffers if any
 	 */
-	if (node->prefetch_iterator)
-		tbm_end_private_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);
 }
@@ -745,11 +713,9 @@ ExecInitBitmapHeapScan(BitmapHeapScan *node, EState *estate, int eflags)
 	/* Zero the statistics counters */
 	memset(&scanstate->stats, 0, sizeof(BitmapHeapScanInstrumentation));
 
-	scanstate->prefetch_iterator = NULL;
 	scanstate->prefetch_pages = 0;
-	scanstate->prefetch_target = 0;
+	scanstate->prefetch_target = -1;
 	scanstate->initialized = false;
-	scanstate->shared_prefetch_iterator = NULL;
 	scanstate->pstate = NULL;
 	scanstate->recheck = true;
 	scanstate->blockno = InvalidBlockNumber;
@@ -908,7 +874,7 @@ ExecBitmapHeapInitializeDSM(BitmapHeapScanState *node,
 	/* Initialize the mutex */
 	SpinLockInit(&pstate->mutex);
 	pstate->prefetch_pages = 0;
-	pstate->prefetch_target = 0;
+	pstate->prefetch_target = -1;
 	pstate->state = BM_INITIAL;
 
 	ConditionVariableInit(&pstate->cv);
@@ -945,6 +911,8 @@ ExecBitmapHeapReInitializeDSM(BitmapHeapScanState *node,
 		return;
 
 	pstate->state = BM_INITIAL;
+	pstate->prefetch_pages = 0;
+	pstate->prefetch_target = -1;
 
 	if (DsaPointerIsValid(pstate->tbmiterator))
 		tbm_free_shared_area(dsa, pstate->tbmiterator);
diff --git a/src/include/access/relscan.h b/src/include/access/relscan.h
index 153561addc0..8ca8f789617 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/relfilelocator.h"
@@ -25,9 +26,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.
@@ -45,12 +43,8 @@ typedef struct TableScanDescData
 	 */
 	union
 	{
-		/* Iterators for Bitmap Table Scans */
-		struct
-		{
-			struct TBMPrivateIterator *rs_iterator;
-			struct TBMSharedIterator *rs_shared_iterator;
-		}			bitmap;
+		/* Iterator for Bitmap Table Scans */
+		TBMIterator rs_tbmiterator;
 
 		/*
 		 * Range of ItemPointers for table_scan_getnextslot_tidrange() to
diff --git a/src/include/access/tableam.h b/src/include/access/tableam.h
index ea4c21114ad..dbcea659fbc 100644
--- a/src/include/access/tableam.h
+++ b/src/include/access/tableam.h
@@ -958,17 +958,13 @@ 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->st.bitmap.rs_shared_iterator = NULL;
-	result->st.bitmap.rs_iterator = NULL;
-	return result;
+	return rel->rd_tableam->scan_begin(rel, snapshot, nkeys, key,
+									   NULL, flags);
 }
 
 /*
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
index 73266920f3a..7cd4b57509e 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -1840,7 +1840,6 @@ typedef struct SharedBitmapHeapInstrumentation
  *		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
  *		sinstrument		   statistics for parallel workers
  *		recheck			   do current page's tuples need recheck
@@ -1855,12 +1854,11 @@ typedef struct BitmapHeapScanState
 	TIDBitmap  *tbm;
 	Buffer		pvmbuffer;
 	BitmapHeapScanInstrumentation stats;
-	TBMPrivateIterator *prefetch_iterator;
+	TBMIterator prefetch_iterator;
 	int			prefetch_pages;
 	int			prefetch_target;
 	int			prefetch_maximum;
 	bool		initialized;
-	TBMSharedIterator *shared_prefetch_iterator;
 	ParallelBitmapHeapState *pstate;
 	SharedBitmapHeapInstrumentation *sinstrument;
 	bool		recheck;
diff --git a/src/include/nodes/tidbitmap.h b/src/include/nodes/tidbitmap.h
index 65ecc5ba045..db678d541fc 100644
--- a/src/include/nodes/tidbitmap.h
+++ b/src/include/nodes/tidbitmap.h
@@ -92,4 +92,14 @@ extern void tbm_end_iterate(TBMIterator *iterator);
 
 extern TBMIterateResult *tbm_iterate(TBMIterator *iterator);
 
+static inline bool
+tbm_exhausted(TBMIterator *iterator)
+{
+	/*
+	 * It doesn't matter if we check the private or shared iterator here. If
+	 * tbm_end_iterate() was called, they will be NULL
+	 */
+	return !iterator->private_iterator;
+}
+
 #endif							/* TIDBITMAP_H */
-- 
2.45.2

