From ac6049e8b0af21666314edb73018fd99940d0731 Mon Sep 17 00:00:00 2001
From: Siddharth Kothari <sidkot@google.com>
Date: Tue, 17 Feb 2026 12:23:56 +0000
Subject: [PATCH] Add RetrieveInstrumentationCustomScan hook 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.

This patch introduces an optional RetrieveInstrumentationCustomScan
callback to the CustomExecMethods struct. This allows custom scan
providers to implement logic to collect and consolidate instrumentation
from shared memory.

The new hook is called in ExecRetrieveInstrumentation for CustomScanState
nodes during the parallel query cleanup phase.
---
 src/backend/executor/execParallel.c | 3 +++
 src/backend/executor/nodeCustom.c   | 9 +++++++++
 src/include/executor/nodeCustom.h   | 1 +
 src/include/nodes/extensible.h      | 3 +++
 4 files changed, 16 insertions(+)

diff --git a/src/backend/executor/execParallel.c b/src/backend/executor/execParallel.c
index f87978c137e..c2416a98170 100644
--- a/src/backend/executor/execParallel.c
+++ b/src/backend/executor/execParallel.c
@@ -1118,6 +1118,9 @@ ExecParallelRetrieveInstrumentation(PlanState *planstate,
 		case T_BitmapHeapScanState:
 			ExecBitmapHeapRetrieveInstrumentation((BitmapHeapScanState *) planstate);
 			break;
+		case T_CustomScanState:
+			ExecCustomScanRetrieveInstrumentation((CustomScanState *) planstate);
+			break;
 		default:
 			break;
 	}
diff --git a/src/backend/executor/nodeCustom.c b/src/backend/executor/nodeCustom.c
index a9ad5af6a98..f027d24993f 100644
--- a/src/backend/executor/nodeCustom.c
+++ b/src/backend/executor/nodeCustom.c
@@ -217,6 +217,15 @@ ExecCustomScanInitializeWorker(CustomScanState *node,
 	}
 }
 
+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..0a9cfb40381 100644
--- a/src/include/executor/nodeCustom.h
+++ b/src/include/executor/nodeCustom.h
@@ -37,6 +37,7 @@ extern void ExecCustomScanReInitializeDSM(CustomScanState *node,
 										  ParallelContext *pcxt);
 extern void ExecCustomScanInitializeWorker(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..cda478b538f 100644
--- a/src/include/nodes/extensible.h
+++ b/src/include/nodes/extensible.h
@@ -151,6 +151,9 @@ typedef struct CustomExecMethods
 											   void *coordinate);
 	void		(*ShutdownCustomScan) (CustomScanState *node);
 
+	/* Optional: retrieve parallel instrumentation */
+	void		(*RetrieveInstrumentationCustomScan) (CustomScanState *node);
+
 	/* Optional: print additional information in EXPLAIN */
 	void		(*ExplainCustomScan) (CustomScanState *node,
 									  List *ancestors,
--