From d57d42e876ab87085593738066f70cbad546aa73 Mon Sep 17 00:00:00 2001
From: Siddharth Kothari <sidkot@google.com>
Date: Tue, 17 Feb 2026 12:23:56 +0000
Subject: [PATCH v3] Add parallel instrumentation hooks for CustomScan
 providers

CustomScan providers currently lack a standard method to aggregate
instrumentation data from parallel workers back to the leader process
before the Dynamic Shared Memory segment is destroyed. Furthermore,
unlike built-in scans, they are missing the callbacks required to allocate
instrumentation DSM when they are in the parallel part of a plan but
not necessarily parallel-aware.

This patch introduces four optional callbacks to CustomExecMethods:
- EstimateDSMCustomScanInstrumentation
- InitializeDSMCustomScanInstrumentation
- InitializeWorkerCustomScanInstrumentation
- RetrieveInstrumentationCustomScan

This allows custom scan providers to fully support parallel EXPLAIN ANALYZE
metrics and consolidate instrumentation from shared memory similarly to
built-in scan nodes.
---
 doc/src/sgml/custom-scan.sgml       | 43 +++++++++++++++++++++++++++++
 src/backend/executor/execParallel.c | 12 ++++++++
 src/backend/executor/nodeCustom.c   | 37 +++++++++++++++++++++++++
 src/include/executor/nodeCustom.h   |  7 +++++
 src/include/nodes/extensible.h      | 13 +++++++++
 5 files changed, 112 insertions(+)

diff --git a/doc/src/sgml/custom-scan.sgml b/doc/src/sgml/custom-scan.sgml
index a200d502cdd..8c2a4f82041 100644
--- a/doc/src/sgml/custom-scan.sgml
+++ b/doc/src/sgml/custom-scan.sgml
@@ -409,6 +409,49 @@ void (*ShutdownCustomScan) (CustomScanState *node);
 
    <para>
 <programlisting>
+void (*EstimateDSMCustomScanInstrumentation) (CustomScanState *node,
+                                              ParallelContext *pcxt);
+</programlisting>
+    Estimate the amount of dynamic shared memory that will be required
+    for parallel operation's instrumentation (for <literal>EXPLAIN ANALYZE</literal>).
+    This callback is optional, and need only be supplied if this custom
+    scan provider supports parallel execution and needs custom instrumentation.
+   </para>
+
+   <para>
+<programlisting>
+void (*InitializeDSMCustomScanInstrumentation) (CustomScanState *node,
+                                                ParallelContext *pcxt);
+</programlisting>
+    Initialize the dynamic shared memory that will be required for parallel
+    operation's instrumentation.
+    This callback is optional, and need only be supplied if this custom
+    scan provider supports parallel execution and needs custom instrumentation.
+   </para>
+
+   <para>
+<programlisting>
+void (*InitializeWorkerCustomScanInstrumentation) (CustomScanState *node,
+                                                   ParallelWorkerContext *pwcxt);
+</programlisting>
+    Initialize a parallel worker's local state based on the shared state
+    set up by the leader during <function>InitializeDSMCustomScanInstrumentation</function>.
+    This callback is optional, and need only be supplied if this custom
+    scan provider supports parallel execution and needs custom instrumentation.
+   </para>
+
+   <para>
+<programlisting>
+void (*RetrieveInstrumentationCustomScan) (CustomScanState *node);
+</programlisting>
+    Retrieve parallel instrumentation data from dynamic shared memory into local
+    memory when the custom-scan plan node finishes execution.
+    This callback is optional, and need only be supplied if this custom
+    scan provider supports parallel execution and needs custom instrumentation.
+   </para>
+
+   <para>
+<programlisting>
 void (*ExplainCustomScan) (CustomScanState *node,
                            List *ancestors,
                            ExplainState *es);
diff --git a/src/backend/executor/execParallel.c b/src/backend/executor/execParallel.c
index 81b87d82fab..7983ac1024e 100644
--- a/src/backend/executor/execParallel.c
+++ b/src/backend/executor/execParallel.c
@@ -304,6 +304,9 @@ ExecParallelEstimate(PlanState *planstate, ExecParallelEstimateContext *e)
 			if (planstate->plan->parallel_aware)
 				ExecCustomScanEstimate((CustomScanState *) planstate,
 									   e->pcxt);
+			/* even when not parallel-aware, for EXPLAIN ANALYZE */
+			ExecCustomScanInstrumentEstimate((CustomScanState *) planstate,
+											 e->pcxt);
 			break;
 		case T_BitmapHeapScanState:
 			if (planstate->plan->parallel_aware)
@@ -552,6 +555,9 @@ ExecParallelInitializeDSM(PlanState *planstate,
 			if (planstate->plan->parallel_aware)
 				ExecCustomScanInitializeDSM((CustomScanState *) planstate,
 											d->pcxt);
+			/* even when not parallel-aware, for EXPLAIN ANALYZE */
+			ExecCustomScanInstrumentInitDSM((CustomScanState *) planstate,
+											d->pcxt);
 			break;
 		case T_BitmapHeapScanState:
 			if (planstate->plan->parallel_aware)
@@ -1166,6 +1172,9 @@ ExecParallelRetrieveInstrumentation(PlanState *planstate,
 		case T_TidRangeScanState:
 			ExecTidRangeScanRetrieveInstrumentation((TidRangeScanState *) planstate);
 			break;
+		case T_CustomScanState:
+			ExecCustomScanRetrieveInstrumentation((CustomScanState *) planstate);
+			break;
 		default:
 			break;
 	}
@@ -1451,6 +1460,9 @@ ExecParallelInitializeWorker(PlanState *planstate, ParallelWorkerContext *pwcxt)
 			if (planstate->plan->parallel_aware)
 				ExecCustomScanInitializeWorker((CustomScanState *) planstate,
 											   pwcxt);
+			/* even when not parallel-aware, for EXPLAIN ANALYZE */
+			ExecCustomScanInstrumentInitWorker((CustomScanState *) planstate,
+											   pwcxt);
 			break;
 		case T_BitmapHeapScanState:
 			if (planstate->plan->parallel_aware)
diff --git a/src/backend/executor/nodeCustom.c b/src/backend/executor/nodeCustom.c
index b7cc890cd20..04a043543d1 100644
--- a/src/backend/executor/nodeCustom.c
+++ b/src/backend/executor/nodeCustom.c
@@ -217,6 +217,43 @@ ExecCustomScanInitializeWorker(CustomScanState *node,
 	}
 }
 
+void
+ExecCustomScanInstrumentEstimate(CustomScanState *node, ParallelContext *pcxt)
+{
+	const CustomExecMethods *methods = node->methods;
+
+	if (methods->EstimateDSMCustomScanInstrumentation)
+		methods->EstimateDSMCustomScanInstrumentation(node, pcxt);
+}
+
+void
+ExecCustomScanInstrumentInitDSM(CustomScanState *node, ParallelContext *pcxt)
+{
+	const CustomExecMethods *methods = node->methods;
+
+	if (methods->InitializeDSMCustomScanInstrumentation)
+		methods->InitializeDSMCustomScanInstrumentation(node, pcxt);
+}
+
+void
+ExecCustomScanInstrumentInitWorker(CustomScanState *node,
+								   ParallelWorkerContext *pwcxt)
+{
+	const CustomExecMethods *methods = node->methods;
+
+	if (methods->InitializeWorkerCustomScanInstrumentation)
+		methods->InitializeWorkerCustomScanInstrumentation(node, pwcxt);
+}
+
+void
+ExecCustomScanRetrieveInstrumentation(CustomScanState *node)
+{
+	const CustomExecMethods *methods = node->methods;
+
+	if (methods->RetrieveInstrumentationCustomScan)
+		methods->RetrieveInstrumentationCustomScan(node);
+}
+
 void
 ExecShutdownCustomScan(CustomScanState *node)
 {
diff --git a/src/include/executor/nodeCustom.h b/src/include/executor/nodeCustom.h
index fb0acc6e414..f702ef01709 100644
--- a/src/include/executor/nodeCustom.h
+++ b/src/include/executor/nodeCustom.h
@@ -37,6 +37,13 @@ extern void ExecCustomScanReInitializeDSM(CustomScanState *node,
 										  ParallelContext *pcxt);
 extern void ExecCustomScanInitializeWorker(CustomScanState *node,
 										   ParallelWorkerContext *pwcxt);
+extern void ExecCustomScanInstrumentEstimate(CustomScanState *node,
+											 ParallelContext *pcxt);
+extern void ExecCustomScanInstrumentInitDSM(CustomScanState *node,
+											ParallelContext *pcxt);
+extern void ExecCustomScanInstrumentInitWorker(CustomScanState *node,
+											   ParallelWorkerContext *pwcxt);
+extern void ExecCustomScanRetrieveInstrumentation(CustomScanState *node);
 extern void ExecShutdownCustomScan(CustomScanState *node);
 
 #endif							/* NODECUSTOM_H */
diff --git a/src/include/nodes/extensible.h b/src/include/nodes/extensible.h
index 517db95c4a3..69f5a627039 100644
--- a/src/include/nodes/extensible.h
+++ b/src/include/nodes/extensible.h
@@ -151,6 +151,19 @@ typedef struct CustomExecMethods
 											   void *coordinate);
 	void		(*ShutdownCustomScan) (CustomScanState *node);
 
+	/* Optional: estimate parallel instrumentation size */
+	void		(*EstimateDSMCustomScanInstrumentation) (CustomScanState *node,
+														 ParallelContext *pcxt);
+	/* Optional: initialize parallel instrumentation DSM */
+	void		(*InitializeDSMCustomScanInstrumentation) (CustomScanState *node,
+														   ParallelContext *pcxt);
+	/* Optional: initialize parallel instrumentation worker */
+	void		(*InitializeWorkerCustomScanInstrumentation) (CustomScanState *node,
+															  ParallelWorkerContext *pwcxt);
+
+	/* Optional: retrieve parallel instrumentation */
+	void		(*RetrieveInstrumentationCustomScan) (CustomScanState *node);
+
 	/* Optional: print additional information in EXPLAIN */
 	void		(*ExplainCustomScan) (CustomScanState *node,
 									  List *ancestors,
-- 
2.55.0.795.g602f6c329a-goog

