From 9972e4838b67ea63550505577d1447925d7f97e1 Mon Sep 17 00:00:00 2001
From: Melanie Plageman <melanieplageman@gmail.com>
Date: Fri, 22 Mar 2024 13:43:41 -0400
Subject: [PATCH v9 14/17] Unify parallel and serial BitmapHeapScan iterator
 interfaces

Introduce a new type, BitmapHeapIterator, which allows unified access to
both TBMIterator and TBMSharedIterators. This encapsulates the parallel
and serial iterators and their access and makes the bitmap heap scan
code a bit cleaner.
---
 src/backend/access/heap/heapam_handler.c  | 40 ++++++---------
 src/backend/executor/nodeBitmapHeapscan.c | 62 ++++++++++++++++++-----
 src/include/access/relscan.h              | 20 +++++---
 src/include/access/tableam.h              | 61 ++++++----------------
 src/tools/pgindent/typedefs.list          |  1 +
 5 files changed, 96 insertions(+), 88 deletions(-)

diff --git a/src/backend/access/heap/heapam_handler.c b/src/backend/access/heap/heapam_handler.c
index efd2784e03..1a9f7b02d1 100644
--- a/src/backend/access/heap/heapam_handler.c
+++ b/src/backend/access/heap/heapam_handler.c
@@ -55,6 +55,7 @@ static bool SampleHeapTupleVisible(TableScanDesc scan, Buffer buffer,
 								   OffsetNumber tupoffset);
 
 static BlockNumber heapam_scan_get_blocks_done(HeapScanDesc hscan);
+
 static inline void BitmapAdjustPrefetchIterator(HeapScanDesc scan);
 static inline void BitmapAdjustPrefetchTarget(HeapScanDesc scan);
 static inline void BitmapPrefetch(HeapScanDesc scan);
@@ -2131,17 +2132,15 @@ BitmapAdjustPrefetchIterator(HeapScanDesc scan)
 
 	if (pstate == NULL)
 	{
-		TBMIterator *prefetch_iterator = scan->rs_base.pf_tbmiterator;
-
 		if (scan->prefetch_pages > 0)
 		{
 			/* The main iterator has closed the distance by one page */
 			scan->prefetch_pages--;
 		}
-		else if (prefetch_iterator)
+		else if (scan->rs_base.rs_pf_tbmiterator)
 		{
 			/* Do not let the prefetch iterator get behind the main one */
-			tbm_iterate(prefetch_iterator, &tbmpre);
+			bhs_iterate(scan->rs_base.rs_pf_tbmiterator, &tbmpre);
 			scan->pfblockno = tbmpre.blockno;
 		}
 		return;
@@ -2154,8 +2153,6 @@ BitmapAdjustPrefetchIterator(HeapScanDesc scan)
 	 */
 	if (scan->rs_base.prefetch_maximum > 0)
 	{
-		TBMSharedIterator *prefetch_iterator = scan->rs_base.pf_shared_tbmiterator;
-
 		SpinLockAcquire(&pstate->mutex);
 		if (pstate->prefetch_pages > 0)
 		{
@@ -2175,9 +2172,9 @@ BitmapAdjustPrefetchIterator(HeapScanDesc scan)
 			 * we don't validate the blockno here as we do in non-parallel
 			 * case.
 			 */
-			if (prefetch_iterator)
+			if (scan->rs_base.rs_pf_tbmiterator)
 			{
-				tbm_shared_iterate(prefetch_iterator, &tbmpre);
+				bhs_iterate(scan->rs_base.rs_pf_tbmiterator, &tbmpre);
 				scan->pfblockno = tbmpre.blockno;
 			}
 		}
@@ -2209,10 +2206,7 @@ heapam_scan_bitmap_next_block(TableScanDesc scan,
 	{
 		CHECK_FOR_INTERRUPTS();
 
-		if (scan->shared_tbmiterator)
-			tbm_shared_iterate(scan->shared_tbmiterator, &tbmres);
-		else
-			tbm_iterate(scan->tbmiterator, &tbmres);
+		bhs_iterate(scan->rs_tbmiterator, &tbmres);
 
 		if (!BlockNumberIsValid(tbmres.blockno))
 		{
@@ -2353,7 +2347,7 @@ heapam_scan_bitmap_next_block(TableScanDesc scan,
 	 * ahead of the current block.
 	 */
 	if (scan->bm_parallel == NULL &&
-		scan->pf_tbmiterator &&
+		scan->rs_pf_tbmiterator &&
 		hscan->pfblockno > hscan->rs_base.blockno)
 		elog(ERROR, "prefetch and main iterators are out of sync");
 
@@ -2427,22 +2421,20 @@ BitmapPrefetch(HeapScanDesc scan)
 
 	if (pstate == NULL)
 	{
-		TBMIterator *prefetch_iterator = scan->rs_base.pf_tbmiterator;
-
-		if (prefetch_iterator)
+		if (scan->rs_base.rs_pf_tbmiterator)
 		{
 			while (scan->prefetch_pages < scan->prefetch_target)
 			{
 				TBMIterateResult tbmpre;
 				bool		skip_fetch;
 
-				tbm_iterate(prefetch_iterator, &tbmpre);
+				bhs_iterate(scan->rs_base.rs_pf_tbmiterator, &tbmpre);
 
 				if (!BlockNumberIsValid(tbmpre.blockno))
 				{
 					/* No more pages to prefetch */
-					tbm_end_iterate(prefetch_iterator);
-					scan->rs_base.pf_tbmiterator = NULL;
+					bhs_end_iterate(scan->rs_base.rs_pf_tbmiterator);
+					scan->rs_base.rs_pf_tbmiterator = NULL;
 					break;
 				}
 				scan->prefetch_pages++;
@@ -2470,9 +2462,7 @@ BitmapPrefetch(HeapScanDesc scan)
 
 	if (pstate->prefetch_pages < pstate->prefetch_target)
 	{
-		TBMSharedIterator *prefetch_iterator = scan->rs_base.pf_shared_tbmiterator;
-
-		if (prefetch_iterator)
+		if (scan->rs_base.rs_pf_tbmiterator)
 		{
 			while (1)
 			{
@@ -2495,12 +2485,12 @@ BitmapPrefetch(HeapScanDesc scan)
 				if (!do_prefetch)
 					return;
 
-				tbm_shared_iterate(prefetch_iterator, &tbmpre);
+				bhs_iterate(scan->rs_base.rs_pf_tbmiterator, &tbmpre);
 				if (!BlockNumberIsValid(tbmpre.blockno))
 				{
 					/* No more pages to prefetch */
-					tbm_end_shared_iterate(prefetch_iterator);
-					scan->rs_base.pf_shared_tbmiterator = NULL;
+					bhs_end_iterate(scan->rs_base.rs_pf_tbmiterator);
+					scan->rs_base.rs_pf_tbmiterator = NULL;
 					break;
 				}
 
diff --git a/src/backend/executor/nodeBitmapHeapscan.c b/src/backend/executor/nodeBitmapHeapscan.c
index f241f4cb2c..b548642088 100644
--- a/src/backend/executor/nodeBitmapHeapscan.c
+++ b/src/backend/executor/nodeBitmapHeapscan.c
@@ -52,6 +52,54 @@
 static TupleTableSlot *BitmapHeapNext(BitmapHeapScanState *node);
 static inline void BitmapDoneInitializingSharedState(ParallelBitmapHeapState *pstate);
 static bool BitmapShouldInitializeSharedState(ParallelBitmapHeapState *pstate);
+static BitmapHeapIterator *bhs_begin_iterate(TIDBitmap *tbm,
+											 ParallelBitmapHeapState *pstate, dsa_area *personal_area);
+
+BitmapHeapIterator *
+bhs_begin_iterate(TIDBitmap *tbm, ParallelBitmapHeapState *pstate, dsa_area *personal_area)
+{
+	BitmapHeapIterator *result = palloc(sizeof(BitmapHeapIterator));
+
+	result->serial = NULL;
+	result->parallel = NULL;
+
+	if (pstate)
+		result->parallel = tbm_attach_shared_iterate(personal_area, pstate->tbmiterator);
+	else
+		result->serial = tbm_begin_iterate(tbm);
+
+	return result;
+}
+
+void
+bhs_iterate(BitmapHeapIterator *iterator, TBMIterateResult *tbmres)
+{
+	Assert(iterator);
+
+	if (iterator->serial)
+		tbm_iterate(iterator->serial, tbmres);
+	else
+		tbm_shared_iterate(iterator->parallel, tbmres);
+}
+
+void
+bhs_end_iterate(BitmapHeapIterator *iterator)
+{
+	Assert(iterator);
+
+	if (iterator->serial)
+	{
+		tbm_end_iterate(iterator->serial);
+		iterator->serial = NULL;
+	}
+	else
+	{
+		tbm_end_shared_iterate(iterator->parallel);
+		iterator->parallel = NULL;
+	}
+
+	pfree(iterator);
+}
 
 
 /* ----------------------------------------------------------------
@@ -165,21 +213,11 @@ BitmapHeapNext(BitmapHeapScanState *node)
 		scan->prefetch_maximum = pf_maximum;
 		scan->bm_parallel = node->pstate;
 
-		if (!scan->bm_parallel)
-			scan->tbmiterator = tbm_begin_iterate(tbm);
-		else
-			/* Allocate a private iterator and attach the shared state to it */
-			scan->shared_tbmiterator = tbm_attach_shared_iterate(dsa, scan->bm_parallel->tbmiterator);
+		scan->rs_tbmiterator = bhs_begin_iterate(tbm, scan->bm_parallel, dsa);
 
 #ifdef USE_PREFETCH
 		if (scan->prefetch_maximum > 0)
-		{
-			if (!scan->bm_parallel)
-				scan->pf_tbmiterator = tbm_begin_iterate(tbm);
-			else
-				scan->pf_shared_tbmiterator =
-					tbm_attach_shared_iterate(dsa, scan->bm_parallel->prefetch_iterator);
-		}
+			bhs_begin_iterate(tbm, scan->bm_parallel, dsa);
 #endif							/* USE_PREFETCH */
 
 
diff --git a/src/include/access/relscan.h b/src/include/access/relscan.h
index 93168bd350..65092d7226 100644
--- a/src/include/access/relscan.h
+++ b/src/include/access/relscan.h
@@ -27,6 +27,17 @@ struct ParallelTableScanDescData;
 struct TBMIterator;
 struct TBMSharedIterator;
 struct ParallelBitmapHeapState;
+struct TBMIterateResult;
+
+typedef struct BitmapHeapIterator
+{
+	struct TBMIterator *serial;
+	struct TBMSharedIterator *parallel;
+} BitmapHeapIterator;
+
+extern void bhs_iterate(BitmapHeapIterator *tbmiterator, struct TBMIterateResult *tbmres);
+
+extern void bhs_end_iterate(BitmapHeapIterator *iterator);
 
 /*
  * Generic descriptor for table scans. This is the base-class for table scans,
@@ -45,12 +56,9 @@ typedef struct TableScanDescData
 	ItemPointerData rs_maxtid;
 
 	/* Only used for Bitmap table scans */
-	struct TBMIterator *tbmiterator;
-	struct TBMSharedIterator *shared_tbmiterator;
-	/* Prefetch iterators */
-	struct TBMIterator *pf_tbmiterator;
-	struct TBMSharedIterator *pf_shared_tbmiterator;
-
+	BitmapHeapIterator *rs_tbmiterator;
+	/* Prefetch iterator */
+	BitmapHeapIterator *rs_pf_tbmiterator;
 	/* maximum value for prefetch_target */
 	int			prefetch_maximum;
 	struct ParallelBitmapHeapState *bm_parallel;
diff --git a/src/include/access/tableam.h b/src/include/access/tableam.h
index 9cab4462d6..2ded1a124b 100644
--- a/src/include/access/tableam.h
+++ b/src/include/access/tableam.h
@@ -948,10 +948,8 @@ table_beginscan_bm(Relation rel, Snapshot snapshot,
 	uint32		flags = SO_TYPE_BITMAPSCAN | SO_ALLOW_PAGEMODE | extra_flags;
 
 	result = rel->rd_tableam->scan_begin(rel, snapshot, nkeys, key, NULL, flags);
-	result->shared_tbmiterator = NULL;
-	result->tbmiterator = NULL;
-	result->pf_shared_tbmiterator = NULL;
-	result->pf_tbmiterator = NULL;
+	result->rs_tbmiterator = NULL;
+	result->rs_pf_tbmiterator = NULL;
 	result->bm_parallel = NULL;
 	return result;
 }
@@ -1015,29 +1013,15 @@ table_endscan(TableScanDesc scan)
 {
 	if (scan->rs_flags & SO_TYPE_BITMAPSCAN)
 	{
-		if (scan->shared_tbmiterator)
+		bhs_end_iterate(scan->rs_tbmiterator);
+		scan->rs_tbmiterator = NULL;
+#ifdef USE_PREFETCH
+		if (scan->rs_pf_tbmiterator)
 		{
-			tbm_end_shared_iterate(scan->shared_tbmiterator);
-			scan->shared_tbmiterator = NULL;
-		}
-
-		if (scan->pf_shared_tbmiterator)
-		{
-			tbm_end_shared_iterate(scan->pf_shared_tbmiterator);
-			scan->pf_shared_tbmiterator = NULL;
-		}
-
-		if (scan->tbmiterator)
-		{
-			tbm_end_iterate(scan->tbmiterator);
-			scan->tbmiterator = NULL;
-		}
-
-		if (scan->pf_tbmiterator)
-		{
-			tbm_end_iterate(scan->pf_tbmiterator);
-			scan->pf_tbmiterator = NULL;
+			bhs_end_iterate(scan->rs_pf_tbmiterator);
+			scan->rs_pf_tbmiterator = NULL;
 		}
+#endif
 	}
 
 	scan->rs_rd->rd_tableam->scan_end(scan);
@@ -1052,29 +1036,16 @@ table_rescan(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->pf_shared_tbmiterator)
-		{
-			tbm_end_shared_iterate(scan->pf_shared_tbmiterator);
-			scan->pf_shared_tbmiterator = NULL;
-		}
-
-		if (scan->tbmiterator)
-		{
-			tbm_end_iterate(scan->tbmiterator);
-			scan->tbmiterator = NULL;
-		}
+		bhs_end_iterate(scan->rs_tbmiterator);
+		scan->rs_tbmiterator = NULL;
 
-		if (scan->pf_tbmiterator)
+#ifdef USE_PREFETCH
+		if (scan->rs_pf_tbmiterator)
 		{
-			tbm_end_iterate(scan->pf_tbmiterator);
-			scan->pf_tbmiterator = NULL;
+			bhs_end_iterate(scan->rs_pf_tbmiterator);
+			scan->rs_pf_tbmiterator = NULL;
 		}
+#endif
 	}
 
 	scan->rs_rd->rd_tableam->scan_rescan(scan, key, false, false, false, false);
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index e2a0525dd4..58317eac44 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -259,6 +259,7 @@ BitString
 BitmapAnd
 BitmapAndPath
 BitmapAndState
+BitmapHeapIterator
 BitmapHeapPath
 BitmapHeapScan
 BitmapHeapScanState
-- 
2.40.1

