Fix incorrect reinit of run-time pruning for EPQ rechecks

For run-time partition pruning running during execution (rather than at
executor startup), EvalPlanQualStart() reuses the es_part_prune_states
from the parent EState.  This could cause an issue when the EPQ recheck is
finished, its EState is cleaned up and the es_query_cxt is deleted by
FreeExecutorState().  That cleanup could free memory that was allocated
again in InitExecPartitionPruneContexts() when it was called for the EPQ's
EState, and naturally, those fields got allocated in the EPQ's
es_query_cxt MemoryContext, which overwrote the ones for the parent
EState.  If pruning occurs again after the EPQ's EState was cleaned up,
then we could have accessed free'd memory.  This certainly caused issues in
debug builds because we clobber freed memory, but may not have in
production, depending on if the MemoryContext being deleted results in a
free().  It may not if the allocation was on the context's init block and
the context was pushed onto context_freelists[] for reuse.

Here we fix this by adding an "initialized" field to PartitionPruneState
and exit early in InitExecPartitionPruneContexts() when the state is
already initialized.

This was broken by bb3ec16e1, but the breakage there is slightly
different as that commit neglected to do what was fixed in 8741e48e5
(which did get backpatched to v18 in 9a82a64ed).

Reported-by: Vladimir Savin <vladimir@encord.com>
Author: David Rowley <dgrowleyml@gmail.com>
Reviewed-by: Andrey Rachitskiy <pl0h0yp1@gmail.com>
Discussion: https://postgr.es/m/CAP2G_gGpV=rKMLuET4RW-qO41GvBvi6Czsjsj16eT9hAwHxG6w@mail.gmail.com
Backpatch-through: 18

src/backend/executor/execPartition.c

@@ -2007,6 +2007,8 @@ CreatePartitionPruneState(EState *estate, PartitionPruneInfo *pruneinfo,
 	prunestate->other_subplans = bms_copy(pruneinfo->other_subplans);
 	prunestate->do_initial_prune = false;	/* may be set below */
 	prunestate->do_exec_prune = false;	/* may be set below */
+	prunestate->initialized = false;	/* set in
+										 * InitExecPartitionPruneContexts */
 	prunestate->num_partprunedata = n_part_hierarchies;
 
 	/*
@@ -2354,6 +2356,16 @@ InitExecPartitionPruneContexts(PartitionPruneState *prunestate,
 	Assert(parent_plan != NULL);
 	estate = parent_plan->state;
 
+	/*
+	 * PartitionPruneStates are sometimes shared and this one may have been
+	 * initialized already.  Sharing of states occurs for EPQ, for example.
+	 * See EvalPlanQualStart().
+	 */
+	if (prunestate->initialized)
+		return;
+
+	prunestate->initialized = true;
+
 	/*
 	 * No need to fix subplans maps if initial pruning didn't eliminate any
 	 * subplans.

src/include/executor/execPartition.h

@@ -113,6 +113,8 @@ typedef struct PartitionPruningData
  *						startup (at any hierarchy level).
  * do_exec_prune		true if pruning should be performed during
  *						executor run (at any hierarchy level).
+ * initialized			true if InitExecPartitionPruneContexts has been called
+ *						on this PartitionPruneState
  * num_partprunedata	Number of items in "partprunedata" array.
  * partprunedata		Array of PartitionPruningData pointers for the plan's
  *						partitioned relation(s), one for each partitioning
@@ -126,6 +128,7 @@ typedef struct PartitionPruneState
 	MemoryContext prune_context;
 	bool		do_initial_prune;
 	bool		do_exec_prune;
+	bool		initialized;
 	int			num_partprunedata;
 	PartitionPruningData *partprunedata[FLEXIBLE_ARRAY_MEMBER];
 } PartitionPruneState;

src/test/isolation/expected/eval-plan-qual-partition-prune.out

@@ -0,0 +1,30 @@
+Parsed test spec with 2 sessions
+
+starting permutation: s1_begin s1_update s2_update s1_commit s2_select
+step s1_begin: BEGIN;
+step s1_update: UPDATE lk SET n = n + 1 WHERE id = 1;
+step s2_update: 
+  WITH locked AS (
+    SELECT * FROM lk WHERE id = 1 FOR UPDATE
+  ), upd AS (
+    UPDATE lp SET b = lp.b + 1 + (SELECT count(*) FROM locked)
+    FROM lp lp2 WHERE lp2.a = lp.a
+    RETURNING lp.*
+  )
+  SELECT * FROM upd ORDER BY a;
+ <waiting ...>
+step s1_commit: COMMIT;
+step s2_update: <... completed>
+a|b
+-+-
+1|3
+2|3
+(2 rows)
+
+step s2_select: SELECT * FROM lp ORDER BY a;
+a|b
+-+-
+1|3
+2|3
+(2 rows)
+

src/test/isolation/isolation_schedule

@@ -38,6 +38,7 @@ test: fk-snapshot
 test: subxid-overflow
 test: eval-plan-qual
 test: eval-plan-qual-trigger
+test: eval-plan-qual-partition-prune
 test: inplace-inval
 test: intra-grant-inplace
 test: intra-grant-inplace-db

src/test/isolation/specs/eval-plan-qual-partition-prune.spec

@@ -0,0 +1,43 @@
+# Test run-time partition pruning after an EPQ recheck
+
+setup
+{
+  CREATE TABLE lp (a int PRIMARY KEY, b int) PARTITION BY LIST (a);
+  CREATE TABLE lp1 PARTITION OF lp FOR VALUES IN (1);
+  CREATE TABLE lp2 PARTITION OF lp FOR VALUES IN (2);
+  INSERT INTO lp VALUES (1, 1), (2, 1);
+  CREATE TABLE lk (id int PRIMARY KEY, n int);
+  INSERT INTO lk VALUES (1, 0);
+}
+
+teardown
+{
+  DROP TABLE lp, lk;
+}
+
+session s1
+step s1_begin	{ BEGIN; }
+step s1_update	{ UPDATE lk SET n = n + 1 WHERE id = 1; }
+step s1_commit	{ COMMIT; }
+
+session s2
+setup
+{
+  SET enable_hashjoin = off;
+  SET enable_mergejoin = off;
+  SET enable_seqscan = off;
+}
+step s2_update
+{
+  WITH locked AS (
+    SELECT * FROM lk WHERE id = 1 FOR UPDATE
+  ), upd AS (
+    UPDATE lp SET b = lp.b + 1 + (SELECT count(*) FROM locked)
+    FROM lp lp2 WHERE lp2.a = lp.a
+    RETURNING lp.*
+  )
+  SELECT * FROM upd ORDER BY a;
+}
+step s2_select	{ SELECT * FROM lp ORDER BY a; }
+
+permutation s1_begin s1_update s2_update s1_commit s2_select