1 files changed, 57 insertions, 13 deletions

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);