From 74dc075dc8f844e036fc38e005fc512b6dd54bc9 Mon Sep 17 00:00:00 2001
From: Amit Langote <amitlan@postgresql.org>
Date: Tue, 11 Nov 2025 22:30:52 +0900
Subject: [PATCH v2 4/5] Use pruning-aware locking in cached plans

Extend GetCachedPlan() to perform ExecutorPrep() on each planned
statement, capturing unpruned relids and initial pruning results.
Use this data to acquire execution locks only on surviving partitions,
avoiding unnecessary locking of pruned tables even when using cached
plans.

Introduce CachedPlanPrepData to carry ExecutorPrep results
through the plan caching layer. Adjust call sites in SPI,
functions, portals, and EXPLAIN to propagate this data.

This ensures pruning decisions made during initial pruning are
consistently reused without redoing pruning logic in executor paths
like parallel workers. It also lays the groundwork for
pruning-dependent lock behavior during plan reuse.

To maintain correctness when all target partitions are pruned, also
reinstate the firstResultRel locking behavior lost in commit
28317de72. That commit required the first ModifyTable target to
remain initialized for executor assumptions to hold. We now
explicitly track these relids in PlannerGlobal and PlannedStmt so they
are locked even if pruned, preserving that rule across cached plan
reuse.
---
 src/backend/commands/prepare.c         |  19 +-
 src/backend/executor/nodeModifyTable.c |   4 +-
 src/backend/executor/spi.c             |  26 ++-
 src/backend/optimizer/plan/planner.c   |   1 +
 src/backend/optimizer/plan/setrefs.c   |   3 +
 src/backend/tcop/postgres.c            |   9 +-
 src/backend/utils/cache/plancache.c    | 234 ++++++++++++++++++++++++-
 src/include/nodes/pathnodes.h          |   3 +
 src/include/nodes/plannodes.h          |  10 ++
 src/include/utils/plancache.h          |  24 ++-
 10 files changed, 312 insertions(+), 21 deletions(-)

diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c
index afd449c73ba..23332d19b37 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -154,6 +154,7 @@ ExecuteQuery(ParseState *pstate,
 {
 	PreparedStatement *entry;
 	CachedPlan *cplan;
+	CachedPlanPrepData cprep = {0};
 	List	   *plan_list;
 	ParamListInfo paramLI = NULL;
 	EState	   *estate = NULL;
@@ -193,7 +194,10 @@ ExecuteQuery(ParseState *pstate,
 									   entry->plansource->query_string);
 
 	/* Replan if needed, and increment plan refcount for portal */
-	cplan = GetCachedPlan(entry->plansource, paramLI, NULL, NULL);
+	/* Keep ExecutorPrep state with the portal and its resowner. */
+	cprep.context = portal->portalContext;
+	cprep.owner = portal->resowner;
+	cplan = GetCachedPlan(entry->plansource, paramLI, NULL, NULL, &cprep);
 	plan_list = cplan->stmt_list;
 
 	/*
@@ -205,7 +209,7 @@ ExecuteQuery(ParseState *pstate,
 					  query_string,
 					  entry->plansource->commandTag,
 					  plan_list,
-					  NIL,
+					  cprep.prep_list,
 					  cplan);
 
 	/*
@@ -575,6 +579,7 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	PreparedStatement *entry;
 	const char *query_string;
 	CachedPlan *cplan;
+	CachedPlanPrepData cprep = {0};
 	List	   *plan_list;
 	List	   *prep_list;
 	ListCell   *p;
@@ -633,8 +638,14 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	}
 
 	/* Replan if needed, and acquire a transient refcount */
+	/* ExecutorPrep state is local to this EXPLAIN EXECUTE call. */
+	cprep.context = CurrentMemoryContext;
+	cprep.owner = CurrentResourceOwner;
+	if (es->generic)
+		cprep.eflags = EXEC_FLAG_EXPLAIN_GENERIC;
 	cplan = GetCachedPlan(entry->plansource, paramLI,
-						  CurrentResourceOwner, pstate->p_queryEnv);
+						  CurrentResourceOwner, pstate->p_queryEnv,
+						  &cprep);
 
 	INSTR_TIME_SET_CURRENT(planduration);
 	INSTR_TIME_SUBTRACT(planduration, planstart);
@@ -653,7 +664,7 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	}
 
 	plan_list = cplan->stmt_list;
-	prep_list = NIL;
+	prep_list = cprep.prep_list;
 
 	/* Explain each query */
 	i = 0;
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index 4c5647ac38a..c5812612f8d 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -4648,8 +4648,8 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 	 * as a reference for building the ResultRelInfo of the target partition.
 	 * In either case, it doesn't matter which result relation is kept, so we
 	 * just keep the first one, if all others have been pruned.  See also,
-	 * ExecDoInitialPruning(), which ensures that this first result relation
-	 * has been locked.
+	 * AcquireExecutorLocksUnpruned(), which ensures that this first result
+	 * relation has been locked.
 	 */
 	i = 0;
 	foreach(l, node->resultRelations)
diff --git a/src/backend/executor/spi.c b/src/backend/executor/spi.c
index 7a3cb944d6f..d580f1e0425 100644
--- a/src/backend/executor/spi.c
+++ b/src/backend/executor/spi.c
@@ -1579,6 +1579,7 @@ SPI_cursor_open_internal(const char *name, SPIPlanPtr plan,
 {
 	CachedPlanSource *plansource;
 	CachedPlan *cplan;
+	CachedPlanPrepData cprep = {0};
 	List	   *stmt_list;
 	char	   *query_string;
 	Snapshot	snapshot;
@@ -1659,7 +1660,11 @@ SPI_cursor_open_internal(const char *name, SPIPlanPtr plan,
 	 */
 
 	/* Replan if needed, and increment plan refcount for portal */
-	cplan = GetCachedPlan(plansource, paramLI, NULL, _SPI_current->queryEnv);
+	/* ExecutorPrep state lives in this portal's context. */
+	cprep.context = portal->portalContext;
+	cprep.owner = portal->resowner;
+	cplan = GetCachedPlan(plansource, paramLI, NULL, _SPI_current->queryEnv,
+						  &cprep);
 	stmt_list = cplan->stmt_list;
 
 	if (!plan->saved)
@@ -1685,7 +1690,7 @@ SPI_cursor_open_internal(const char *name, SPIPlanPtr plan,
 					  query_string,
 					  plansource->commandTag,
 					  stmt_list,
-					  NIL,
+					  cprep.prep_list,	/* lives in portalContext */
 					  cplan);
 
 	/*
@@ -2078,6 +2083,7 @@ SPI_plan_get_cached_plan(SPIPlanPtr plan)
 {
 	CachedPlanSource *plansource;
 	CachedPlan *cplan;
+	CachedPlanPrepData cprep = {0};
 	SPICallbackArg spicallbackarg;
 	ErrorContextCallback spierrcontext;
 
@@ -2101,9 +2107,13 @@ SPI_plan_get_cached_plan(SPIPlanPtr plan)
 	error_context_stack = &spierrcontext;
 
 	/* Get the generic plan for the query */
+	/* ExecutorPrep() state lives in caller's active context. */
+	cprep.context = CurrentMemoryContext;
+	cprep.owner = CurrentResourceOwner;
 	cplan = GetCachedPlan(plansource, NULL,
 						  plan->saved ? CurrentResourceOwner : NULL,
-						  _SPI_current->queryEnv);
+						  _SPI_current->queryEnv,
+						  &cprep);
 	Assert(cplan == plansource->gplan);
 
 	/* Pop the error context stack */
@@ -2501,6 +2511,7 @@ _SPI_execute_plan(SPIPlanPtr plan, const SPIExecuteOptions *options,
 		CachedPlanSource *plansource = (CachedPlanSource *) lfirst(lc1);
 		List	   *stmt_list;
 		ListCell   *lc2;
+		CachedPlanPrepData cprep = {0};
 		List	   *prep_list;
 		int			i;
 
@@ -2577,11 +2588,16 @@ _SPI_execute_plan(SPIPlanPtr plan, const SPIExecuteOptions *options,
 		 * Replan if needed, and increment plan refcount.  If it's a saved
 		 * plan, the refcount must be backed by the plan_owner.
 		 */
+
+		/* ExecutorPrep state is per _SPI_execute_plan call. */
+		cprep.context = CurrentMemoryContext;
+		cprep.owner = CurrentResourceOwner;
 		cplan = GetCachedPlan(plansource, options->params,
-							  plan_owner, _SPI_current->queryEnv);
+							  plan_owner, _SPI_current->queryEnv,
+							  &cprep);
 
 		stmt_list = cplan->stmt_list;
-		prep_list = NIL;
+		prep_list = cprep.prep_list;
 
 		/*
 		 * If we weren't given a specific snapshot to use, and the statement
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index c4fd646b999..4c76e78c1da 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -608,6 +608,7 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
 											  glob->prunableRelids);
 	result->permInfos = glob->finalrteperminfos;
 	result->resultRelations = glob->resultRelations;
+	result->firstResultRels = glob->firstResultRels;
 	result->appendRelations = glob->appendRelations;
 	result->subplans = glob->subplans;
 	result->rewindPlanIDs = glob->rewindPlanIDs;
diff --git a/src/backend/optimizer/plan/setrefs.c b/src/backend/optimizer/plan/setrefs.c
index ccdc9bc264a..229b39060ae 100644
--- a/src/backend/optimizer/plan/setrefs.c
+++ b/src/backend/optimizer/plan/setrefs.c
@@ -1274,6 +1274,9 @@ set_plan_refs(PlannerInfo *root, Plan *plan, int rtoffset)
 						lappend_int(root->glob->resultRelations,
 									splan->rootRelation);
 				}
+				root->glob->firstResultRels =
+					lappend_int(root->glob->firstResultRels,
+								linitial_int(splan->resultRelations));
 			}
 			break;
 		case T_Append:
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index d3964a12a14..249829f59a0 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -1639,6 +1639,7 @@ exec_bind_message(StringInfo input_message)
 	int16	   *rformats = NULL;
 	CachedPlanSource *psrc;
 	CachedPlan *cplan;
+	CachedPlanPrepData cprep = {0};
 	Portal		portal;
 	char	   *query_string;
 	char	   *saved_stmt_name;
@@ -2021,7 +2022,11 @@ exec_bind_message(StringInfo input_message)
 	 * will be generated in MessageContext.  The plan refcount will be
 	 * assigned to the Portal, so it will be released at portal destruction.
 	 */
-	cplan = GetCachedPlan(psrc, params, NULL, NULL);
+
+	/* ExecutorPrep() state lives in portal context. */
+	cprep.context = portal->portalContext;
+	cprep.owner = portal->resowner;
+	cplan = GetCachedPlan(psrc, params, NULL, NULL, &cprep);
 
 	/*
 	 * Now we can define the portal.
@@ -2034,7 +2039,7 @@ exec_bind_message(StringInfo input_message)
 					  query_string,
 					  psrc->commandTag,
 					  cplan->stmt_list,
-					  NIL,
+					  cprep.prep_list,
 					  cplan);
 
 	/* Portal is defined, set the plan ID based on its contents. */
diff --git a/src/backend/utils/cache/plancache.c b/src/backend/utils/cache/plancache.c
index 6661d2c6b73..c1cfd47422c 100644
--- a/src/backend/utils/cache/plancache.c
+++ b/src/backend/utils/cache/plancache.c
@@ -93,7 +93,7 @@ static bool StmtPlanRequiresRevalidation(CachedPlanSource *plansource);
 static bool BuildingPlanRequiresSnapshot(CachedPlanSource *plansource);
 static List *RevalidateCachedQuery(CachedPlanSource *plansource,
 								   QueryEnvironment *queryEnv);
-static bool CheckCachedPlan(CachedPlanSource *plansource);
+static bool PrepAndCheckCachedPlan(CachedPlanSource *plansource, CachedPlanPrepData *cprep);
 static CachedPlan *BuildCachedPlan(CachedPlanSource *plansource, List *qlist,
 								   ParamListInfo boundParams, QueryEnvironment *queryEnv);
 static bool choose_custom_plan(CachedPlanSource *plansource,
@@ -101,6 +101,8 @@ static bool choose_custom_plan(CachedPlanSource *plansource,
 static double cached_plan_cost(CachedPlan *plan, bool include_planner);
 static Query *QueryListGetPrimaryStmt(List *stmts);
 static void AcquireExecutorLocks(List *stmt_list, bool acquire);
+static void AcquireExecutorLocksUnpruned(List *stmt_list, bool acquire,
+										 CachedPlanPrepData *cprep);
 static void AcquirePlannerLocks(List *stmt_list, bool acquire);
 static void ScanQueryForLocks(Query *parsetree, bool acquire);
 static bool ScanQueryWalker(Node *node, bool *acquire);
@@ -137,6 +139,26 @@ ResourceOwnerForgetPlanCacheRef(ResourceOwner owner, CachedPlan *plan)
 /* GUC parameter */
 int			plan_cache_mode = PLAN_CACHE_MODE_AUTO;
 
+/*
+ * Lock acquisition policy for execution locks.
+ *
+ * LOCK_ALL acquires locks on all relations mentioned in the plan,
+ * reproducing the behavior of AcquireExecutorLocks().
+ *
+ * LOCK_UNPRUNED restricts locking to only the unpruned relations. That
+ * includes those mentioned in PlannedStmt.unprunableRelids and the leaf
+ * partitions remaining after performing initial pruning.
+ */
+typedef enum LockPolicy
+{
+	LOCK_ALL,
+	LOCK_UNPRUNED,
+} LockPolicy;
+
+static void AcquireExecutorLocksWithPolicy(List *stmt_list,
+										   LockPolicy policy, bool acquire,
+										   CachedPlanPrepData *cprep);
+
 /*
  * InitPlanCache: initialize module during InitPostgres.
  *
@@ -938,7 +960,12 @@ RevalidateCachedQuery(CachedPlanSource *plansource,
 }
 
 /*
- * CheckCachedPlan: see if the CachedPlanSource's generic plan is valid.
+ * PrepAndCheckCachedPlan: see if the CachedPlanSource's generic plan is valid.
+ *
+ * If 'cprep' is not NULL, ExecutorPrep() is applied to each PlannedStmt to
+ * compute the set of partitions that survive initial runtime pruning in order
+ * to only lock them.  The resulting ExecPrep structures are saved in cprep for
+ * later reuse by ExecutorStart().
  *
  * Caller must have already called RevalidateCachedQuery to verify that the
  * querytree is up to date.
@@ -947,7 +974,7 @@ RevalidateCachedQuery(CachedPlanSource *plansource,
  * (We must do this for the "true" result to be race-condition-free.)
  */
 static bool
-CheckCachedPlan(CachedPlanSource *plansource)
+PrepAndCheckCachedPlan(CachedPlanSource *plansource, CachedPlanPrepData *cprep)
 {
 	CachedPlan *plan = plansource->gplan;
 
@@ -975,13 +1002,15 @@ CheckCachedPlan(CachedPlanSource *plansource)
 	 */
 	if (plan->is_valid)
 	{
+		LockPolicy policy = !cprep ? LOCK_ALL : LOCK_UNPRUNED;
+
 		/*
 		 * Plan must have positive refcount because it is referenced by
 		 * plansource; so no need to fear it disappears under us here.
 		 */
 		Assert(plan->refcount > 0);
 
-		AcquireExecutorLocks(plan->stmt_list, true);
+		AcquireExecutorLocksWithPolicy(plan->stmt_list, policy, true, cprep);
 
 		/*
 		 * If plan was transient, check to see if TransactionXmin has
@@ -1003,7 +1032,7 @@ CheckCachedPlan(CachedPlanSource *plansource)
 		}
 
 		/* Oops, the race case happened.  Release useless locks. */
-		AcquireExecutorLocks(plan->stmt_list, false);
+		AcquireExecutorLocksWithPolicy(plan->stmt_list, policy, false, cprep);
 	}
 
 	/*
@@ -1283,6 +1312,10 @@ cached_plan_cost(CachedPlan *plan, bool include_planner)
  * On return, the plan is valid and we have sufficient locks to begin
  * execution.
  *
+ * If 'cprep' is not NULL and a generic plan is reused, the function prepares
+ * each PlannedStmt via ExecutorPrep() and stores the results in
+ * cprep->prep_list.  These are intended to be passed later to ExecutorStart().
+ *
  * On return, the refcount of the plan has been incremented; a later
  * ReleaseCachedPlan() call is expected.  If "owner" is not NULL then
  * the refcount has been reported to that ResourceOwner (note that this
@@ -1293,7 +1326,8 @@ cached_plan_cost(CachedPlan *plan, bool include_planner)
  */
 CachedPlan *
 GetCachedPlan(CachedPlanSource *plansource, ParamListInfo boundParams,
-			  ResourceOwner owner, QueryEnvironment *queryEnv)
+			  ResourceOwner owner, QueryEnvironment *queryEnv,
+			  CachedPlanPrepData *cprep)
 {
 	CachedPlan *plan = NULL;
 	List	   *qlist;
@@ -1315,7 +1349,9 @@ GetCachedPlan(CachedPlanSource *plansource, ParamListInfo boundParams,
 
 	if (!customplan)
 	{
-		if (CheckCachedPlan(plansource))
+		if (cprep)
+			cprep->params = boundParams;
+		if (PrepAndCheckCachedPlan(plansource, cprep))
 		{
 			/* We want a generic plan, and we already have a valid one */
 			plan = plansource->gplan;
@@ -1902,6 +1938,38 @@ QueryListGetPrimaryStmt(List *stmts)
 	return NULL;
 }
 
+/*
+ * AcquireExecutorLocksWithPolicy
+ *		Acquire or release execution locks for a cached plan according to
+ *		the specified policy.
+ *
+ * LOCK_ALL reproduces AcquireExecutorLocks(), locking every relation in
+ * each PlannedStmt's rtable.  LOCK_UNPRUNED restricts locking to the
+ * unprunable rels and partitions that survive initial runtime pruning.
+ *
+ * When LOCK_UNPRUNED is used on acquire, ExecutorPrep() is invoked for
+ * each PlannedStmt and the resulting ExecPrep pointers are appended to
+ * cprep->prep_list in cprep->context.  On release, the same ExecPrep
+ * list is consulted to determine which relations to unlock and is then
+ * cleaned up with ExecPrepCleanup().
+ */
+static void
+AcquireExecutorLocksWithPolicy(List *stmt_list, LockPolicy policy, bool acquire,
+							   CachedPlanPrepData *cprep)
+{
+	switch (policy)
+	{
+		case LOCK_ALL:
+			AcquireExecutorLocks(stmt_list, acquire);
+			break;
+		case LOCK_UNPRUNED:
+			AcquireExecutorLocksUnpruned(stmt_list, acquire, cprep);
+			break;
+		default:
+			elog(ERROR, "invalid LockPolicy");
+	}
+}
+
 /*
  * AcquireExecutorLocks: acquire locks needed for execution of a cached plan;
  * or release them if acquire is false.
@@ -1954,6 +2022,158 @@ AcquireExecutorLocks(List *stmt_list, bool acquire)
 	}
 }
 
+/*
+ * LockRelids
+ * 		Acquire or release locks on the specified relids, which reference
+ * 		entries in the provided range table.
+ *
+ * Helper for AcquireExecutorLocksUnpruned().
+ */
+static void
+LockRelids(List *rtable, Bitmapset *relids, bool acquire)
+{
+	int	rtindex = -1;
+
+	while ((rtindex = bms_next_member(relids, rtindex)) >= 0)
+	{
+		RangeTblEntry *rte = list_nth_node(RangeTblEntry, rtable, rtindex - 1);
+
+		Assert(rte->rtekind == RTE_RELATION ||
+			   (rte->rtekind == RTE_SUBQUERY && OidIsValid(rte->relid)));
+
+		/*
+		 * Acquire the appropriate type of lock on each relation OID. Note
+		 * that we don't actually try to open the rel, and hence will not
+		 * fail if it's been dropped entirely --- we'll just transiently
+		 * acquire a non-conflicting lock.
+		 */
+		if (acquire)
+			LockRelationOid(rte->relid, rte->rellockmode);
+		else
+			UnlockRelationOid(rte->relid, rte->rellockmode);
+	}
+}
+
+/*
+ * AcquireExecutorLocksUnpruned
+ *		Acquire or release execution locks for only unpruned relations
+ *		referenced by the given PlannedStmts.
+ *
+ * On acquire, this:
+ *	- locks unprunable rels listed in PlannedStmt.unprunableRelids
+ *	- runs ExecutorPrep() to perform initial runtime pruning
+ *	- locks the surviving partitions reported in the prep estate
+ *	- appends the ExecPrep pointer for each PlannedStmt to cprep->prep_list
+ *
+ * On release, it:
+ *	- looks up the ExecPrep object for each PlannedStmt from cprep->prep_list
+ *	  (which must already be populated)
+ *	- unlocks the same relations identified during acquire
+ *	- calls ExecPrepCleanup() on each ExecPrep
+ *
+ * prep_list is extended during acquire and must match stmt_list one-to-one
+ * when releasing locks.  Memory allocation for ExecPrep happens in
+ * cprep->context.  Locks are acquired using cprep->owner.
+ */
+
+static void
+AcquireExecutorLocksUnpruned(List *stmt_list, bool acquire,
+							 CachedPlanPrepData *cprep)
+{
+	MemoryContext oldcontext = MemoryContextSwitchTo(cprep->context);
+	ListCell   *lc1;
+	List	   *prep_list;
+	int			i;
+
+	Assert(cprep);
+
+	/*
+	 * When releasing locks, use the ExecPrep list (if any) created during
+	 * acquisition to determine which relids to unlock. The list must match
+	 * the PlannedStmt list one-to-one.
+	 */
+	prep_list = cprep->prep_list;
+	Assert(acquire || list_length(prep_list) == list_length(stmt_list));
+
+	i = 0;
+	foreach(lc1, stmt_list)
+	{
+		PlannedStmt *plannedstmt = lfirst_node(PlannedStmt, lc1);
+		ExecPrep *prep;
+
+		if (plannedstmt->commandType == CMD_UTILITY)
+		{
+			/* Same as AcquireExecutorLocks(). */
+			Query	   *query = UtilityContainsQuery(plannedstmt->utilityStmt);
+
+			if (query)
+				ScanQueryForLocks(query, acquire);
+
+			/* Keep the list one-to-one with stmt_list. */
+			if (acquire)
+				cprep->prep_list = lappend(cprep->prep_list, NULL);
+			continue;
+		}
+
+		/*
+		 * Lock tables mentioned in the original query and other unprunable
+		 * relations that were added to the plan via inheritance expansion.
+		 */
+		LockRelids(plannedstmt->rtable, plannedstmt->unprunableRelids, acquire);
+
+		/* Lock partitions surviving runtime initial pruning. */
+		if (acquire)
+		{
+			prep = ExecutorPrep(plannedstmt, cprep->params, cprep->owner, true,
+								cprep->eflags);
+			Assert(prep || plannedstmt->partPruneInfos == NULL);
+			cprep->prep_list = lappend(cprep->prep_list, prep);
+		}
+		else
+			prep = list_nth(prep_list, i++);
+
+		Assert(prep == NULL || prep->prep_estate);
+		if (prep)
+		{
+			EState *prep_estate = prep->prep_estate;
+
+			/*
+			 * es_unpruned_relids includes plannedstmt->unprunableRelids,
+			 * which we've already locked. Filter them out to avoid double-locking.
+			 */
+			Bitmapset *lock_relids = bms_difference(prep_estate->es_unpruned_relids,
+													plannedstmt->unprunableRelids);
+
+			/*
+			 * firstResultRels may contain pruned partitions that must still be
+			 * locked to satisfy executor assumptions (see comments in
+			 * ExecInitModifyTable(). Ensure they’re included here.
+			 */
+			if (plannedstmt->resultRelations)
+			{
+				ListCell *lc2;
+
+				foreach(lc2, plannedstmt->firstResultRels)
+				{
+					Index       firstResultRel = lfirst_int(lc2);
+
+					if (!bms_is_member(firstResultRel, lock_relids))
+						lock_relids = bms_add_member(lock_relids, firstResultRel);
+				}
+			}
+
+			LockRelids(plannedstmt->rtable, lock_relids, acquire);
+			bms_free(lock_relids);
+		}
+
+		/* Clean up prep if releasing locks. */
+		if (!acquire)
+			ExecPrepCleanup(prep);
+	}
+
+	MemoryContextSwitchTo(oldcontext);
+}
+
 /*
  * AcquirePlannerLocks: acquire locks needed for planning of a querytree list;
  * or release them if acquire is false.
diff --git a/src/include/nodes/pathnodes.h b/src/include/nodes/pathnodes.h
index 30d889b54c5..6fb86dc05f6 100644
--- a/src/include/nodes/pathnodes.h
+++ b/src/include/nodes/pathnodes.h
@@ -141,6 +141,9 @@ typedef struct PlannerGlobal
 	/* "flat" list of integer RT indexes */
 	List	   *resultRelations;
 
+	/* "flat" list of integer RT indexes (one per ModifyTable node) */
+	List	   *firstResultRels;
+
 	/* "flat" list of AppendRelInfos */
 	List	   *appendRelations;
 
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index c4393a94321..eb211f1ba56 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -123,6 +123,16 @@ typedef struct PlannedStmt
 	/* integer list of RT indexes, or NIL */
 	List	   *resultRelations;
 
+	/*
+	 * rtable indexes of first target relation in each ModifyTable node in the
+	 * plan for INSERT/UPDATE/DELETE/MERGE.  NIL if resultRelations is NIL.
+	 *
+	 * These are used by AcquireExecutorLocksUnpruned() to ensure that the
+	 * first result rel for each ModifyTable remains locked even if pruned;
+	 * see ExecInitModifyTable() for the executor side assumptions.
+	 */
+	List	   *firstResultRels;
+
 	/* list of AppendRelInfo nodes */
 	List	   *appendRelations;
 
diff --git a/src/include/utils/plancache.h b/src/include/utils/plancache.h
index a82b66d4bc2..c7b8ec4be39 100644
--- a/src/include/utils/plancache.h
+++ b/src/include/utils/plancache.h
@@ -197,6 +197,27 @@ typedef struct CachedExpression
 } CachedExpression;
 
 
+/*
+ * CachedPlanPrepData
+ *      Carries ExecutorPrep results for each PlannedStmt in a CachedPlan,
+ *      along with context and owner information needed to allocate them.
+ *
+ * prep_list is indexed one-to-one with CachedPlan->stmt_list, and is
+ * populated when GetCachedPlan() prepares a reused generic plan.  The
+ * same list is later used to determine which relations to unlock when
+ * releasing execution locks.
+ *
+ * ExecutorPrep state is allocated in 'context' and owned by 'owner'.
+ */
+typedef struct CachedPlanPrepData
+{
+	List   *prep_list;		/* one ExecPrep per PlannedStmt, or NULL */
+	ParamListInfo params;	/* params visible to ExecutorPrep */
+	MemoryContext context;	/* where to allocate ExecPrep objects */
+	ResourceOwner owner;	/* ResourceOwner for ExecutorPrep state */
+	int		eflags;			/* executor flags to pass to ExecutorPrep */
+} CachedPlanPrepData;
+
 extern void InitPlanCache(void);
 extern void ResetPlanCache(void);
 
@@ -240,7 +261,8 @@ extern List *CachedPlanGetTargetList(CachedPlanSource *plansource,
 extern CachedPlan *GetCachedPlan(CachedPlanSource *plansource,
 								 ParamListInfo boundParams,
 								 ResourceOwner owner,
-								 QueryEnvironment *queryEnv);
+								 QueryEnvironment *queryEnv,
+								 CachedPlanPrepData *cprep);
 extern void ReleaseCachedPlan(CachedPlan *plan, ResourceOwner owner);
 
 extern bool CachedPlanAllowsSimpleValidityCheck(CachedPlanSource *plansource,
-- 
2.47.3

