diff --git a/src/backend/distributed/executor/citus_custom_scan.c b/src/backend/distributed/executor/citus_custom_scan.c index 87fca7422b1..2ca2ce8ccb3 100644 --- a/src/backend/distributed/executor/citus_custom_scan.c +++ b/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. diff --git a/src/backend/distributed/executor/local_executor.c b/src/backend/distributed/executor/local_executor.c index 5480a1d142d..7a827880049 100644 --- a/src/backend/distributed/executor/local_executor.c +++ b/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); diff --git a/src/backend/distributed/planner/local_plan_cache.c b/src/backend/distributed/planner/local_plan_cache.c index 1703635e090..85bf0d9f532 100644 --- a/src/backend/distributed/planner/local_plan_cache.c +++ b/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; } diff --git a/src/test/regress/bin/normalize.sed b/src/test/regress/bin/normalize.sed index 34762fe1538..b4a01d60b28 100644 --- a/src/test/regress/bin/normalize.sed +++ b/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/ diff --git a/src/test/regress/expected/local_plan_cache_reuse.out b/src/test/regress/expected/local_plan_cache_reuse.out new file mode 100644 index 00000000000..3aaa39c2c80 --- /dev/null +++ b/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; diff --git a/src/test/regress/expected/local_shard_execution.out b/src/test/regress/expected/local_shard_execution.out index f50db7245a6..8fc79f7bda3 100644 --- a/src/test/regress/expected/local_shard_execution.out +++ b/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 diff --git a/src/test/regress/expected/local_shard_execution_replicated.out b/src/test/regress/expected/local_shard_execution_replicated.out index 2ee728b66eb..512fa3d336d 100644 --- a/src/test/regress/expected/local_shard_execution_replicated.out +++ b/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 diff --git a/src/test/regress/expected/single_node.out b/src/test/regress/expected/single_node.out index bdd89459875..d565f730056 100644 --- a/src/test/regress/expected/single_node.out +++ b/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 diff --git a/src/test/regress/multi_mx_schedule b/src/test/regress/multi_mx_schedule index 48755266367..fea42f8dccb 100644 --- a/src/test/regress/multi_mx_schedule +++ b/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 diff --git a/src/test/regress/sql/local_plan_cache_reuse.sql b/src/test/regress/sql/local_plan_cache_reuse.sql new file mode 100644 index 00000000000..1dfa967419b --- /dev/null +++ b/src/test/regress/sql/local_plan_cache_reuse.sql @@ -0,0 +1,108 @@ +-- +-- 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'); + +-- 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; + +-- 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); +EXECUTE cached_count(1); +EXECUTE cached_count(1); +EXECUTE cached_count(1); +EXECUTE cached_count(1); +EXECUTE cached_count(1); +EXECUTE cached_count(1); +EXECUTE cached_count(1); +EXECUTE cached_count(1); +EXECUTE cached_count(1); + +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; + +-- ... and each one is reachable through the shard its key actually hashes to +SELECT key, value FROM dist_table WHERE key = 1; +SELECT key, value FROM dist_table WHERE key = 6; + +-- 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; + +\c - - - :master_port +SET client_min_messages TO WARNING; +DROP SCHEMA local_plan_cache_reuse CASCADE;