Show index search count in EXPLAIN ANALYZE, take 2.

Expose the count of index searches/index descents in EXPLAIN ANALYZE's
output for index scan/index-only scan/bitmap index scan nodes.  This
information is particularly useful with scans that use ScalarArrayOp
quals, where the number of index searches can be unpredictable due to
implementation details that interact with physical index characteristics
(at least with nbtree SAOP scans, since Postgres 17 commit 5bf748b8).
The information shown also provides useful context when EXPLAIN ANALYZE
runs a plan with an index scan node that successfully applied the skip
scan optimization (set to be added to nbtree by an upcoming patch).

The instrumentation works by teaching all index AMs to increment a new
nsearches counter whenever a new index search begins.  The counter is
incremented at exactly the same point that index AMs already increment
the pg_stat_*_indexes.idx_scan counter (we're counting the same event,
but at the scan level rather than the relation level).  Parallel queries
have workers copy their local counter struct into shared memory when an
index scan node ends -- even when it isn't a parallel aware scan node.
An earlier version of this patch that only worked with parallel aware
scans became commit 5ead85fb (though that was quickly reverted by commit
d00107cd following "debug_parallel_query=regress" buildfarm failures).

Our approach doesn't match the approach used when tracking other index
scan related costs (e.g., "Rows Removed by Filter:").  It is comparable
to the approach used in similar cases involving costs that are only
readily accessible inside an access method, not from the executor proper
(e.g., "Heap Blocks:" output for a Bitmap Heap Scan, which was recently
enhanced to show per-worker costs by commit 5a1e6df3, using essentially
the same scheme as the one used here).  It is necessary for index AMs to
have direct responsibility for maintaining the new counter, since the
counter might need to be incremented multiple times per amgettuple call
(or per amgetbitmap call).  But it is also necessary for the executor
proper to manage the shared memory now used to transfer each worker's
counter struct to the leader.

Author: Peter Geoghegan <pg@bowt.ie>
Reviewed-By: Robert Haas <robertmhaas@gmail.com>
Reviewed-By: Tomas Vondra <tomas@vondra.me>
Reviewed-By: Masahiro Ikeda <ikedamsh@oss.nttdata.com>
Reviewed-By: Matthias van de Meent <boekewurm+postgres@gmail.com>
Discussion: https://postgr.es/m/CAH2-WzkRqvaqR2CTNqTZP0z6FuL4-3ED6eQB0yx38XBNj1v-4Q@mail.gmail.com
Discussion: https://postgr.es/m/CAH2-Wz=PKR6rB7qbx+Vnd7eqeB5VTcrW=iJvAsTsKbdG+kW_UA@mail.gmail.com

10 files changed, 80 insertions, 21 deletions

diff --git a/src/backend/access/brin/brin.c b/src/backend/access/brin/brin.c
index b01009c5d85..737ad638808 100644
--- a/src/backend/access/brin/brin.c
+++ b/src/backend/access/brin/brin.c
@@ -592,6 +592,8 @@ bringetbitmap(IndexScanDesc scan, TIDBitmap *tbm)
 	opaque = (BrinOpaque *) scan->opaque;
 	bdesc = opaque->bo_bdesc;
 	pgstat_count_index_scan(idxRel);
+	if (scan->instrument)
+		scan->instrument->nsearches++;
 
 	/*
 	 * We need to know the size of the table so that we know how long to
diff --git a/src/backend/access/gin/ginscan.c b/src/backend/access/gin/ginscan.c
index 84aa14594f8..f6cdd098a02 100644
--- a/src/backend/access/gin/ginscan.c
+++ b/src/backend/access/gin/ginscan.c
@@ -442,6 +442,8 @@ ginNewScanKey(IndexScanDesc scan)
 	MemoryContextSwitchTo(oldCtx);
 
 	pgstat_count_index_scan(scan->indexRelation);
+	if (scan->instrument)
+		scan->instrument->nsearches++;
 }
 
 void
diff --git a/src/backend/access/gist/gistget.c b/src/backend/access/gist/gistget.c
index cc40e928e0a..387d9972345 100644
--- a/src/backend/access/gist/gistget.c
+++ b/src/backend/access/gist/gistget.c
@@ -625,6 +625,8 @@ gistgettuple(IndexScanDesc scan, ScanDirection dir)
 		GISTSearchItem fakeItem;
 
 		pgstat_count_index_scan(scan->indexRelation);
+		if (scan->instrument)
+			scan->instrument->nsearches++;
 
 		so->firstCall = false;
 		so->curPageData = so->nPageData = 0;
@@ -750,6 +752,8 @@ gistgetbitmap(IndexScanDesc scan, TIDBitmap *tbm)
 		return 0;
 
 	pgstat_count_index_scan(scan->indexRelation);
+	if (scan->instrument)
+		scan->instrument->nsearches++;
 
 	/* Begin the scan by processing the root page */
 	so->curPageData = so->nPageData = 0;
diff --git a/src/backend/access/hash/hashsearch.c b/src/backend/access/hash/hashsearch.c
index a3a1fccf33b..92c15a65be2 100644
--- a/src/backend/access/hash/hashsearch.c
+++ b/src/backend/access/hash/hashsearch.c
@@ -298,6 +298,8 @@ _hash_first(IndexScanDesc scan, ScanDirection dir)
 	HashScanPosItem *currItem;
 
 	pgstat_count_index_scan(rel);
+	if (scan->instrument)
+		scan->instrument->nsearches++;
 
 	/*
 	 * We do not support hash scans with no index qualification, because we
diff --git a/src/backend/access/heap/heapam_handler.c b/src/backend/access/heap/heapam_handler.c
index e78682c3cef..d74f0fbc5cd 100644
--- a/src/backend/access/heap/heapam_handler.c
+++ b/src/backend/access/heap/heapam_handler.c
@@ -749,7 +749,7 @@ heapam_relation_copy_for_cluster(Relation OldHeap, Relation NewHeap,
 
 		tableScan = NULL;
 		heapScan = NULL;
-		indexScan = index_beginscan(OldHeap, OldIndex, SnapshotAny, 0, 0);
+		indexScan = index_beginscan(OldHeap, OldIndex, SnapshotAny, NULL, 0, 0);
 		index_rescan(indexScan, NULL, 0, NULL, 0);
 	}
 	else
diff --git a/src/backend/access/index/genam.c b/src/backend/access/index/genam.c
index 07bae342e25..886c05655f4 100644
--- a/src/backend/access/index/genam.c
+++ b/src/backend/access/index/genam.c
@@ -119,6 +119,7 @@ RelationGetIndexScan(Relation indexRelation, int nkeys, int norderbys)
 	scan->ignore_killed_tuples = !scan->xactStartedInRecovery;
 
 	scan->opaque = NULL;
+	scan->instrument = NULL;
 
 	scan->xs_itup = NULL;
 	scan->xs_itupdesc = NULL;
@@ -446,7 +447,7 @@ systable_beginscan(Relation heapRelation,
 		}
 
 		sysscan->iscan = index_beginscan(heapRelation, irel,
-										 snapshot, nkeys, 0);
+										 snapshot, NULL, nkeys, 0);
 		index_rescan(sysscan->iscan, idxkey, nkeys, NULL, 0);
 		sysscan->scan = NULL;
 
@@ -711,7 +712,7 @@ systable_beginscan_ordered(Relation heapRelation,
 	}
 
 	sysscan->iscan = index_beginscan(heapRelation, indexRelation,
-									 snapshot, nkeys, 0);
+									 snapshot, NULL, nkeys, 0);
 	index_rescan(sysscan->iscan, idxkey, nkeys, NULL, 0);
 	sysscan->scan = NULL;
 
diff --git a/src/backend/access/index/indexam.c b/src/backend/access/index/indexam.c
index 8b1f555435b..55ec4c10352 100644
--- a/src/backend/access/index/indexam.c
+++ b/src/backend/access/index/indexam.c
@@ -256,6 +256,7 @@ IndexScanDesc
 index_beginscan(Relation heapRelation,
 				Relation indexRelation,
 				Snapshot snapshot,
+				IndexScanInstrumentation *instrument,
 				int nkeys, int norderbys)
 {
 	IndexScanDesc scan;
@@ -270,6 +271,7 @@ index_beginscan(Relation heapRelation,
 	 */
 	scan->heapRelation = heapRelation;
 	scan->xs_snapshot = snapshot;
+	scan->instrument = instrument;
 
 	/* prepare to fetch index matches from table */
 	scan->xs_heapfetch = table_index_fetch_begin(heapRelation);
@@ -286,6 +288,7 @@ index_beginscan(Relation heapRelation,
 IndexScanDesc
 index_beginscan_bitmap(Relation indexRelation,
 					   Snapshot snapshot,
+					   IndexScanInstrumentation *instrument,
 					   int nkeys)
 {
 	IndexScanDesc scan;
@@ -299,6 +302,7 @@ index_beginscan_bitmap(Relation indexRelation,
 	 * up by RelationGetIndexScan.
 	 */
 	scan->xs_snapshot = snapshot;
+	scan->instrument = instrument;
 
 	return scan;
 }
@@ -448,14 +452,19 @@ index_restrpos(IndexScanDesc scan)
 
 /*
  * index_parallelscan_estimate - estimate shared memory for parallel scan
+ *
+ * When instrument=true, estimate includes SharedIndexScanInstrumentation
+ * space.  When parallel_aware=true, estimate includes whatever space the
+ * index AM's amestimateparallelscan routine requested when called.
  */
 Size
 index_parallelscan_estimate(Relation indexRelation, int nkeys, int norderbys,
-							Snapshot snapshot)
+							Snapshot snapshot, bool instrument,
+							bool parallel_aware, int nworkers)
 {
 	Size		nbytes;
 
-	Assert(snapshot != InvalidSnapshot);
+	Assert(instrument || parallel_aware);
 
 	RELATION_CHECKS;
 
@@ -463,12 +472,22 @@ index_parallelscan_estimate(Relation indexRelation, int nkeys, int norderbys,
 	nbytes = add_size(nbytes, EstimateSnapshotSpace(snapshot));
 	nbytes = MAXALIGN(nbytes);
 
+	if (instrument)
+	{
+		Size		sharedinfosz;
+
+		sharedinfosz = offsetof(SharedIndexScanInstrumentation, winstrument) +
+			nworkers * sizeof(IndexScanInstrumentation);
+		nbytes = add_size(nbytes, sharedinfosz);
+		nbytes = MAXALIGN(nbytes);
+	}
+
 	/*
-	 * If amestimateparallelscan is not provided, assume there is no
-	 * AM-specific data needed.  (It's hard to believe that could work, but
-	 * it's easy enough to cater to it here.)
+	 * If parallel scan index AM interface can't be used (or index AM provides
+	 * no such interface), assume there is no AM-specific data needed
 	 */
-	if (indexRelation->rd_indam->amestimateparallelscan != NULL)
+	if (parallel_aware &&
+		indexRelation->rd_indam->amestimateparallelscan != NULL)
 		nbytes = add_size(nbytes,
 						  indexRelation->rd_indam->amestimateparallelscan(nkeys,
 																		  norderbys));
@@ -488,11 +507,14 @@ index_parallelscan_estimate(Relation indexRelation, int nkeys, int norderbys,
  */
 void
 index_parallelscan_initialize(Relation heapRelation, Relation indexRelation,
-							  Snapshot snapshot, ParallelIndexScanDesc target)
+							  Snapshot snapshot, bool instrument,
+							  bool parallel_aware, int nworkers,
+							  SharedIndexScanInstrumentation **sharedinfo,
+							  ParallelIndexScanDesc target)
 {
 	Size		offset;
 
-	Assert(snapshot != InvalidSnapshot);
+	Assert(instrument || parallel_aware);
 
 	RELATION_CHECKS;
 
@@ -502,15 +524,34 @@ index_parallelscan_initialize(Relation heapRelation, Relation indexRelation,
 
 	target->ps_locator = heapRelation->rd_locator;
 	target->ps_indexlocator = indexRelation->rd_locator;
-	target->ps_offset = offset;
+	target->ps_offset_ins = 0;
+	target->ps_offset_am = 0;
 	SerializeSnapshot(snapshot, target->ps_snapshot_data);
 
+	if (instrument)
+	{
+		Size		sharedinfosz;
+
+		target->ps_offset_ins = offset;
+		sharedinfosz = offsetof(SharedIndexScanInstrumentation, winstrument) +
+			nworkers * sizeof(IndexScanInstrumentation);
+		offset = add_size(offset, sharedinfosz);
+		offset = MAXALIGN(offset);
+
+		/* Set leader's *sharedinfo pointer, and initialize stats */
+		*sharedinfo = (SharedIndexScanInstrumentation *)
+			OffsetToPointer(target, target->ps_offset_ins);
+		memset(*sharedinfo, 0, sharedinfosz);
+		(*sharedinfo)->num_workers = nworkers;
+	}
+
 	/* aminitparallelscan is optional; assume no-op if not provided by AM */
-	if (indexRelation->rd_indam->aminitparallelscan != NULL)
+	if (parallel_aware && indexRelation->rd_indam->aminitparallelscan != NULL)
 	{
 		void	   *amtarget;
 
-		amtarget = OffsetToPointer(target, offset);
+		target->ps_offset_am = offset;
+		amtarget = OffsetToPointer(target, target->ps_offset_am);
 		indexRelation->rd_indam->aminitparallelscan(amtarget);
 	}
 }
@@ -538,8 +579,10 @@ index_parallelrescan(IndexScanDesc scan)
  * Caller must be holding suitable locks on the heap and the index.
  */
 IndexScanDesc
-index_beginscan_parallel(Relation heaprel, Relation indexrel, int nkeys,
-						 int norderbys, ParallelIndexScanDesc pscan)
+index_beginscan_parallel(Relation heaprel, Relation indexrel,
+						 IndexScanInstrumentation *instrument,
+						 int nkeys, int norderbys,
+						 ParallelIndexScanDesc pscan)
 {
 	Snapshot	snapshot;
 	IndexScanDesc scan;
@@ -558,6 +601,7 @@ index_beginscan_parallel(Relation heaprel, Relation indexrel, int nkeys,
 	 */
 	scan->heapRelation = heaprel;
 	scan->xs_snapshot = snapshot;
+	scan->instrument = instrument;
 
 	/* prepare to fetch index matches from table */
 	scan->xs_heapfetch = table_index_fetch_begin(heaprel);
diff --git a/src/backend/access/nbtree/nbtree.c b/src/backend/access/nbtree/nbtree.c
index 25188a644ef..c0a8833e068 100644
--- a/src/backend/access/nbtree/nbtree.c
+++ b/src/backend/access/nbtree/nbtree.c
@@ -574,7 +574,7 @@ btparallelrescan(IndexScanDesc scan)
 	Assert(parallel_scan);
 
 	btscan = (BTParallelScanDesc) OffsetToPointer(parallel_scan,
-												  parallel_scan->ps_offset);
+												  parallel_scan->ps_offset_am);
 
 	/*
 	 * In theory, we don't need to acquire the LWLock here, because there
@@ -652,7 +652,7 @@ _bt_parallel_seize(IndexScanDesc scan, BlockNumber *next_scan_page,
 	}
 
 	btscan = (BTParallelScanDesc) OffsetToPointer(parallel_scan,
-												  parallel_scan->ps_offset);
+												  parallel_scan->ps_offset_am);
 
 	while (1)
 	{
@@ -760,7 +760,7 @@ _bt_parallel_release(IndexScanDesc scan, BlockNumber next_scan_page,
 	Assert(BlockNumberIsValid(next_scan_page));
 
 	btscan = (BTParallelScanDesc) OffsetToPointer(parallel_scan,
-												  parallel_scan->ps_offset);
+												  parallel_scan->ps_offset_am);
 
 	LWLockAcquire(&btscan->btps_lock, LW_EXCLUSIVE);
 	btscan->btps_nextScanPage = next_scan_page;
@@ -799,7 +799,7 @@ _bt_parallel_done(IndexScanDesc scan)
 		return;
 
 	btscan = (BTParallelScanDesc) OffsetToPointer(parallel_scan,
-												  parallel_scan->ps_offset);
+												  parallel_scan->ps_offset_am);
 
 	/*
 	 * Mark the parallel scan as done, unless some other process did so
@@ -837,7 +837,7 @@ _bt_parallel_primscan_schedule(IndexScanDesc scan, BlockNumber curr_page)
 	Assert(so->numArrayKeys);
 
 	btscan = (BTParallelScanDesc) OffsetToPointer(parallel_scan,
-												  parallel_scan->ps_offset);
+												  parallel_scan->ps_offset_am);
 
 	LWLockAcquire(&btscan->btps_lock, LW_EXCLUSIVE);
 	if (btscan->btps_lastCurrPage == curr_page &&
diff --git a/src/backend/access/nbtree/nbtsearch.c b/src/backend/access/nbtree/nbtsearch.c
index 6b2f464aa04..22b27d01d00 100644
--- a/src/backend/access/nbtree/nbtsearch.c
+++ b/src/backend/access/nbtree/nbtsearch.c
@@ -950,6 +950,8 @@ _bt_first(IndexScanDesc scan, ScanDirection dir)
 	 * _bt_search/_bt_endpoint below
 	 */
 	pgstat_count_index_scan(rel);
+	if (scan->instrument)
+		scan->instrument->nsearches++;
 
 	/*----------
 	 * Examine the scan keys to discover where we need to start the scan.
diff --git a/src/backend/access/spgist/spgscan.c b/src/backend/access/spgist/spgscan.c
index 53f910e9d89..25893050c58 100644
--- a/src/backend/access/spgist/spgscan.c
+++ b/src/backend/access/spgist/spgscan.c
@@ -421,6 +421,8 @@ spgrescan(IndexScanDesc scan, ScanKey scankey, int nscankeys,
 
 	/* count an indexscan for stats */
 	pgstat_count_index_scan(scan->indexRelation);
+	if (scan->instrument)
+		scan->instrument->nsearches++;
 }
 
 void