Investigate unbounded growth of local plan cache

src/backend/distributed/executor/citus_custom_scan.c

@@ -638,7 +638,11 @@ ModifyJobNeedsEvaluation(Job *workerJob)
  *
  * We must not change the distributed plan since it may be reused across multiple
  * executions of a prepared statement. Instead we create a deep copy that we only
- * use for the current execution.
+ * use for the current execution. Begin functions then mutate the copy with any necessary
+ * per-execution, parameter-dependent data, such as the partition key value and the
+ * task list. ExecuteCoordinatorEvaluableExpressions() bakes in this particular
+ * execution's parameter values into the query tree, as well as functions such as
+ * nextval().
  *
  * We also exclude localPlannedStatements from the copyObject call for performance
  * reasons, as they are immutable, so no need to have a deep copy.

src/backend/distributed/executor/local_executor.c

@@ -302,6 +302,7 @@ ExecuteLocalTaskListExtended(List *taskList,
 													  ALLOCSET_DEFAULT_SIZES);
 
 	Task *task = NULL;
+
 	foreach_declared_ptr(task, taskList)
 	{
 		MemoryContext oldContext = MemoryContextSwitchTo(loopContext);

src/backend/distributed/planner/local_plan_cache.c

@@ -31,6 +31,8 @@ static char * DeparseLocalShardQuery(Query *jobQuery, List *relationShardList,
 									 Oid anchorDistributedTableId, int64 anchorShardId);
 static int ExtractParameterTypesForParamListInfo(ParamListInfo originalParamListInfo,
 												 Oid **parameterTypes);
+static PlannedStmt * lookupLocalPlan(uint64 shardId, int32 localGroupId, List *
+									 cachedPlanList);
 
 /*
  * CacheLocalPlanForShardQuery replaces the relation OIDs in the job query
@@ -41,16 +43,21 @@ void
 CacheLocalPlanForShardQuery(Task *task, DistributedPlan *originalDistributedPlan,
 							ParamListInfo paramListInfo)
 {
-	PlannedStmt *localPlan = GetCachedLocalPlan(task, originalDistributedPlan);
-	if (localPlan != NULL)
+	if (list_length(task->relationShardList) == 0)
 	{
-		/* we already have a local plan */
+		/* zero shard plan, no need to cache */
 		return;
 	}
 
-	if (list_length(task->relationShardList) == 0)
+	uint64 shardId = task->anchorShardId;
+	uint32 localGroupId = GetLocalGroupId();
+	PlannedStmt *localPlan = lookupLocalPlan(shardId, localGroupId,
+											 originalDistributedPlan->workerJob->
+											 localPlannedStatements);
+
+	if (localPlan != NULL)
 	{
-		/* zero shard plan, no need to cache */
+		/* Already have a plan, nothing to do here */
 		return;
 	}
 
@@ -88,16 +95,21 @@ CacheLocalPlanForShardQuery(Task *task, DistributedPlan *originalDistributedPlan
 
 	LockRelationOid(rangeTableEntry->relid, lockMode);
 
-	LocalPlannedStatement *localPlannedStatement = CitusMakeNode(LocalPlannedStatement);
 	localPlan = planner(localShardQuery, NULL, 0, NULL);
+	LocalPlannedStatement *localPlannedStatement = CitusMakeNode(LocalPlannedStatement);
 	localPlannedStatement->localPlan = localPlan;
-	localPlannedStatement->shardId = task->anchorShardId;
-	localPlannedStatement->localGroupId = GetLocalGroupId();
+	localPlannedStatement->shardId = shardId;
+	localPlannedStatement->localGroupId = localGroupId;
 
 	originalDistributedPlan->workerJob->localPlannedStatements =
 		lappend(originalDistributedPlan->workerJob->localPlannedStatements,
 				localPlannedStatement);
 
+	ereport(DEBUG2, (errmsg("Created and cached local plan for shard " UINT64_FORMAT
+							" and local group %d",
+							shardId,
+							(int) localGroupId)));
+
 	MemoryContextSwitchTo(oldContext);
 }
 
@@ -231,6 +243,11 @@ ExtractParameterTypesForParamListInfo(ParamListInfo originalParamListInfo,
  * plan in the distributedPlan for the given task if exists.
  *
  * Otherwise, the function returns NULL.
+ *
+ * Callers must ensure that distributedPlan and its workerJob are non-NULL, and
+ * that the job has only one task. IsLocalPlanCachingSupported() establishes
+ * both for the Citus Begin functions, and ExecuteLocalTaskListExtended() in
+ * local_executor.c does so via its isSingleTask check.
  */
 PlannedStmt *
 GetCachedLocalPlan(Task *task, DistributedPlan *distributedPlan)
@@ -240,22 +257,31 @@ GetCachedLocalPlan(Task *task, DistributedPlan *distributedPlan)
 		return NULL;
 	}
 
-	if (list_length(distributedPlan->workerJob->taskList) != 1)
+	if (list_length(distributedPlan->workerJob->taskList) > 1)
 	{
 		/* we only support plan caching for single shard queries */
 		return NULL;
 	}
 
-	List *cachedPlanList = distributedPlan->workerJob->localPlannedStatements;
-	LocalPlannedStatement *localPlannedStatement = NULL;
+	return lookupLocalPlan(task->anchorShardId, GetLocalGroupId(), distributedPlan->
+						   workerJob->localPlannedStatements);
+}
 
-	int32 localGroupId = GetLocalGroupId();
+
+static PlannedStmt *
+lookupLocalPlan(uint64 shardId, int32 localGroupId, List *cachedPlanList)
+{
+	LocalPlannedStatement *localPlannedStatement = NULL;
 
 	foreach_declared_ptr(localPlannedStatement, cachedPlanList)
 	{
-		if (localPlannedStatement->shardId == task->anchorShardId &&
+		if (localPlannedStatement->shardId == shardId &&
 			localPlannedStatement->localGroupId == localGroupId)
 		{
+			ereport(DEBUG2, (errmsg("Found cached local plan for shard " UINT64_FORMAT
+									" and local group %d",
+									shardId, (int) localGroupId)));
+
 			/* already have a cached plan, no need to continue */
 			return localPlannedStatement->localPlan;
 		}

src/test/regress/bin/normalize.sed

@@ -12,6 +12,7 @@ s/ port=[0-9]+ / port=xxxxx /g
 s/placement [0-9]+/placement xxxxx/g
 s/shard [0-9]+/shard xxxxx/g
 s/Shard [0-9]+/Shard xxxxx/g
+s/and local group [0-9]+/and local group xxxxx/g
 s/assigned task [0-9]+ to node/assigned task to node/
 s/node group [12] (but|does)/node group \1/
 

src/test/regress/expected/local_plan_cache_reuse.out

@@ -0,0 +1,192 @@
+--
+-- LOCAL_PLAN_CACHE_REUSE
+--
+-- Regression coverage for the local shard plan cache in local_plan_cache.c.
+--
+-- (a) The shard query behind a prepared statement should be deparsed and
+--     planned once, then reused. The dedup lookup in
+--     CacheLocalPlanForShardQuery() used to consult the generic plan's job,
+--     whose task list is always empty when pruning is deferred, so the lookup
+--     always missed: every execution re-planned the shard query and appended
+--     another entry to localPlannedStatements.
+--
+-- (b) A plan cached while the job had a single task must not be reused once the
+--     same prepared statement routes rows to more than one shard. A cached
+--     INSERT plan carries every VALUES row, so reusing it for a multi-task job
+--     would write rows into the wrong shard.
+CREATE SCHEMA local_plan_cache_reuse;
+SET search_path TO local_plan_cache_reuse;
+SET citus.shard_count TO 4;
+SET citus.shard_replication_factor TO 1;
+SET citus.next_shard_id TO 1490000;
+CREATE TABLE dist_table (key int PRIMARY KEY, value text);
+SELECT create_distributed_table('dist_table', 'key');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- The tests below rely on keys 1 and 5 sharing a shard while key 6 lives on a
+-- different one. Assert that up front so a change in the hash function turns
+-- into an obvious failure here rather than a confusing one further down.
+SELECT get_shard_id_for_distribution_column('dist_table', 1) =
+       get_shard_id_for_distribution_column('dist_table', 5) AS keys_1_and_5_share_a_shard,
+       get_shard_id_for_distribution_column('dist_table', 1) =
+       get_shard_id_for_distribution_column('dist_table', 6) AS keys_1_and_6_share_a_shard;
+ keys_1_and_5_share_a_shard | keys_1_and_6_share_a_shard
+---------------------------------------------------------------------
+ t                          | f
+(1 row)
+
+-- All three keys land on shards local to worker_1, so the statements below take
+-- the local execution path.
+\c - - - :worker_1_port
+SET search_path TO local_plan_cache_reuse;
+--
+-- (a) the shard query is deparsed and planned once, not once per execution
+--
+PREPARE cached_count(int) AS SELECT count(*) FROM dist_table WHERE key = $1;
+SET client_min_messages TO DEBUG2;
+-- Postgres uses a custom plan for the first executions, so no caching happens
+-- yet. Once it switches to the generic plan we should see exactly one
+-- "Created and cached local plan" message, and every execution after that
+-- should report a cache hit instead of planning the shard query again.
+EXECUTE cached_count(1);
+DEBUG:  Deferred pruning for a fast-path router query
+DEBUG:  Creating router plan
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+EXECUTE cached_count(1);
+DEBUG:  Deferred pruning for a fast-path router query
+DEBUG:  Creating router plan
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+EXECUTE cached_count(1);
+DEBUG:  Deferred pruning for a fast-path router query
+DEBUG:  Creating router plan
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+EXECUTE cached_count(1);
+DEBUG:  Deferred pruning for a fast-path router query
+DEBUG:  Creating router plan
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+EXECUTE cached_count(1);
+DEBUG:  Deferred pruning for a fast-path router query
+DEBUG:  Creating router plan
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+EXECUTE cached_count(1);
+DEBUG:  Deferred pruning for a fast-path router query
+DEBUG:  Creating router plan
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+EXECUTE cached_count(1);
+DEBUG:  Created and cached local plan for shard xxxxx and local group xxxxx
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+EXECUTE cached_count(1);
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+EXECUTE cached_count(1);
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+EXECUTE cached_count(1);
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+RESET client_min_messages;
+--
+-- (b) a plan cached for a single-task multi-row INSERT is not reused once the
+--     same statement spreads its rows over several shards
+--
+PREPARE multi_row_insert(int, int) AS
+	INSERT INTO dist_table VALUES ($1, 'first'), ($2, 'second')
+	ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value;
+-- Keys 1 and 5 share a shard, so each of these executions builds a single task
+-- and eventually caches a plan for that shard. That cached plan carries both
+-- VALUES rows.
+EXECUTE multi_row_insert(1, 5);
+EXECUTE multi_row_insert(1, 5);
+EXECUTE multi_row_insert(1, 5);
+EXECUTE multi_row_insert(1, 5);
+EXECUTE multi_row_insert(1, 5);
+EXECUTE multi_row_insert(1, 5);
+EXECUTE multi_row_insert(1, 5);
+EXECUTE multi_row_insert(1, 5);
+DELETE FROM dist_table;
+-- Same prepared statement, but now the rows route to two different shards, so
+-- the job has two tasks and the plan cached above must not be used. If it were,
+-- both rows would be written to the shard holding key 1 and the row for key 6
+-- would be lost.
+EXECUTE multi_row_insert(1, 6);
+-- every row is present exactly once ...
+SELECT key, value FROM dist_table ORDER BY key;
+ key | value
+---------------------------------------------------------------------
+   1 | first
+   6 | second
+(2 rows)
+
+-- ... and each one is reachable through the shard its key actually hashes to
+SELECT key, value FROM dist_table WHERE key = 1;
+ key | value
+---------------------------------------------------------------------
+   1 | first
+(1 row)
+
+SELECT key, value FROM dist_table WHERE key = 6;
+ key | value
+---------------------------------------------------------------------
+   6 | second
+(1 row)
+
+-- inspect the shard placements directly: one row in each of the two shards
+SELECT 1490000 AS shardid, key, value FROM dist_table_1490000
+UNION ALL
+SELECT 1490002 AS shardid, key, value FROM dist_table_1490002
+ORDER BY shardid, key;
+ shardid | key | value
+---------------------------------------------------------------------
+ 1490000 |   1 | first
+ 1490002 |   6 | second
+(2 rows)
+
+\c - - - :master_port
+SET client_min_messages TO WARNING;
+DROP SCHEMA local_plan_cache_reuse CASCADE;

src/test/regress/expected/local_shard_execution.out

@@ -2404,10 +2404,16 @@ SET citus.log_local_commands TO ON;
 SET client_min_messages TO DEBUG2;
 CALL register_for_event(19, 1, 'yes');
 DEBUG:  not pushing down procedure to the same node
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
 NOTICE:  executing the command locally: INSERT INTO local_shard_execution.event_responses_1480001 AS citus_table_alias (event_id, user_id, response) VALUES (19, 1, 'yes'::local_shard_execution.invite_resp) ON CONFLICT(event_id, user_id) DO UPDATE SET response = excluded.response
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
 NOTICE:  executing the command locally: SELECT count(*) AS count FROM local_shard_execution.event_responses_1480001 event_responses WHERE (event_id OPERATOR(pg_catalog.=) 19)
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
 NOTICE:  executing the command locally: SELECT count(*) AS count FROM (SELECT NULL::integer AS event_id, NULL::integer AS user_id, NULL::local_shard_execution.invite_resp AS response WHERE false) event_responses(event_id, user_id, response) WHERE ((event_id OPERATOR(pg_catalog.=) 19) AND false)
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
 NOTICE:  executing the command locally: UPDATE local_shard_execution.event_responses_1480001 event_responses SET response = 'yes'::local_shard_execution.invite_resp WHERE (event_id OPERATOR(pg_catalog.=) 19)
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
 -- should be fine even if no parameters exists in the query
 SELECT count(*) FROM event_responses WHERE event_id = 16;
 DEBUG:  Distributed planning for a fast-path router query

src/test/regress/expected/local_shard_execution_replicated.out

@@ -2342,10 +2342,16 @@ SET citus.log_local_commands TO ON;
 SET client_min_messages TO DEBUG2;
 CALL register_for_event(19, 1, 'yes');
 DEBUG:  stored procedure does not have co-located tables
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
 NOTICE:  executing the command locally: INSERT INTO local_shard_execution_replicated.event_responses_1501001 AS citus_table_alias (event_id, user_id, response) VALUES (19, 1, 'yes'::local_shard_execution_replicated.invite_resp) ON CONFLICT(event_id, user_id) DO UPDATE SET response = excluded.response
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
 NOTICE:  executing the command locally: SELECT count(*) AS count FROM local_shard_execution_replicated.event_responses_1501001 event_responses WHERE (event_id OPERATOR(pg_catalog.=) 19)
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
 NOTICE:  executing the command locally: SELECT count(*) AS count FROM (SELECT NULL::integer AS event_id, NULL::integer AS user_id, NULL::local_shard_execution_replicated.invite_resp AS response WHERE false) event_responses(event_id, user_id, response) WHERE ((event_id OPERATOR(pg_catalog.=) 19) AND false)
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
 NOTICE:  executing the command locally: UPDATE local_shard_execution_replicated.event_responses_1501001 event_responses SET response = 'yes'::local_shard_execution_replicated.invite_resp WHERE (event_id OPERATOR(pg_catalog.=) 19)
+DEBUG:  Found cached local plan for shard xxxxx and local group xxxxx
 -- should be fine even if no parameters exists in the query
 SELECT count(*) FROM event_responses WHERE event_id = 16;
 DEBUG:  Distributed planning for a fast-path router query

src/test/regress/expected/single_node.out

@@ -1225,6 +1225,7 @@ DEBUG:  Creating router plan
 (1 row)
 
 EXECUTE use_local_query_cache(1);
+DEBUG:  Created and cached local plan for shard xxxxx and local group xxxxx
  count
 ---------------------------------------------------------------------
      1

src/test/regress/multi_mx_schedule

@@ -48,6 +48,7 @@ test: multi_mx_modifications
 test: local_shard_execution_replicated
 # the following test has to be run sequentially
 test: local_shard_execution
+test: local_plan_cache_reuse
 test: multi_mx_repartition_udt_w1 multi_mx_repartition_udt_w2
 test: local_shard_copy
 test: undistribute_table_cascade_mx