diff --git a/src/backend/distributed/executor/citus_custom_scan.c b/src/backend/distributed/executor/citus_custom_scan.c index 87fca7422b1..187276771f3 100644 --- a/src/backend/distributed/executor/citus_custom_scan.c +++ b/src/backend/distributed/executor/citus_custom_scan.c @@ -393,20 +393,15 @@ CitusBeginReadOnlyScan(CustomScanState *node, EState *estate, int eflags) /* parameters are filled in, so we can generate a task for this execution */ RegenerateTaskForFasthPathQuery(workerJob); - if (IsLocalPlanCachingSupported(workerJob, originalDistributedPlan)) - { - Task *task = linitial(workerJob->taskList); + bool planAddedToCache = false; + CacheLocalPlanForShardQuery(workerJob, originalDistributedPlan, + estate->es_param_list_info, &planAddedToCache); - /* - * We are going to execute this task locally. If it's not already in - * the cache, create a local plan now and add it to the cache. During - * execution, we will get the plan from the cache. - * - * The plan will be cached across executions when originalDistributedPlan - * represents a prepared statement. - */ - CacheLocalPlanForShardQuery(task, originalDistributedPlan, - estate->es_param_list_info); + /* Use a newly cached plan for the current execution as well. */ + if (planAddedToCache) + { + currentPlan->workerJob->localPlannedStatements = + originalDistributedPlan->workerJob->localPlannedStatements; } } @@ -520,29 +515,15 @@ CitusBeginModifyScan(CustomScanState *node, EState *estate, int eflags) } - /* - * Now that we have populated the task placements we can determine whether - * any of them are local to this node and cache a plan if needed. - */ - if (IsLocalPlanCachingSupported(workerJob, originalDistributedPlan)) - { - Task *task = linitial(workerJob->taskList); + bool planAddedToCache = false; + CacheLocalPlanForShardQuery(workerJob, originalDistributedPlan, + estate->es_param_list_info, &planAddedToCache); - /* - * We are going to execute this task locally. If it's not already in - * the cache, create a local plan now and add it to the cache. During - * execution, we will get the plan from the cache. - * - * WARNING: In this function we'll use the original plan with the original - * query tree, meaning parameters and function calls are back and we'll - * redo evaluation in the local (Postgres) executor. The reason we do this - * is that we only need to cache one generic plan per shard. - * - * The plan will be cached across executions when originalDistributedPlan - * represents a prepared statement. - */ - CacheLocalPlanForShardQuery(task, originalDistributedPlan, - estate->es_param_list_info); + /* Use a newly cached plan for the current execution as well. */ + if (planAddedToCache) + { + currentPlan->workerJob->localPlannedStatements = + originalDistributedPlan->workerJob->localPlannedStatements; } MemoryContextSwitchTo(oldContext); diff --git a/src/backend/distributed/executor/local_executor.c b/src/backend/distributed/executor/local_executor.c index 5480a1d142d..f43894371f8 100644 --- a/src/backend/distributed/executor/local_executor.c +++ b/src/backend/distributed/executor/local_executor.c @@ -341,7 +341,10 @@ ExecuteLocalTaskListExtended(List *taskList, continue; } - PlannedStmt *localPlan = GetCachedLocalPlan(task, distributedPlan); + LocalPlannedStatement *localPlannedStatement = + GetCachedLocalPlan(task, distributedPlan); + PlannedStmt *localPlan = localPlannedStatement ? + localPlannedStatement->localPlan : NULL; /* * If the plan is already cached, don't need to re-plan, just diff --git a/src/backend/distributed/planner/local_plan_cache.c b/src/backend/distributed/planner/local_plan_cache.c index 1703635e090..d786a666fa0 100644 --- a/src/backend/distributed/planner/local_plan_cache.c +++ b/src/backend/distributed/planner/local_plan_cache.c @@ -24,34 +24,70 @@ #include "distributed/multi_executor.h" #include "distributed/version_compat.h" - +static bool IsLocalPlanCachingSupported(Job *currentJob, + DistributedPlan *originalDistributedPlan); static Query * GetLocalShardQueryForCache(Query *jobQuery, Task *task, ParamListInfo paramListInfo); static char * DeparseLocalShardQuery(Query *jobQuery, List *relationShardList, Oid anchorDistributedTableId, int64 anchorShardId); static int ExtractParameterTypesForParamListInfo(ParamListInfo originalParamListInfo, Oid **parameterTypes); +static bool IsJobEligibleForPlanCache(Job *job); + +static LocalPlannedStatement * FindCachedLocalPlannedStatement(Task *task, + DistributedPlan * + distributedPlan); +static List * AddLocalPlanToCache(List **localPlannedStatementsList, + LocalPlannedStatement *localPlannedStatement); +static LocalPlannedStatement * CreatePlanForLocalCache(Job *currentJob, + DistributedPlan * + originalDistributedPlan, + ParamListInfo paramListInfo); + +/* Check if the job has a single task */ +static inline bool +JobHasSingleTask(const Job *job) +{ + return list_length(job->taskList) == 1; +} + + +/* Check if the job has multiple tasks */ +static inline bool +JobHasMultipleTasks(const Job *job) +{ + return list_length(job->taskList) > 1; +} + /* - * CacheLocalPlanForShardQuery replaces the relation OIDs in the job query - * with shard relation OIDs and then plans the query and caches the result - * in the originalDistributedPlan (which may be preserved across executions). + * CacheLocalPlanForShardQuery returns an existing local plan for the current + * shard, or, when local plan caching is supported, creates and caches one in + * originalDistributedPlan (which may be preserved across executions). + * Returns NULL when local plan caching is unsupported or no plan can be created. */ -void -CacheLocalPlanForShardQuery(Task *task, DistributedPlan *originalDistributedPlan, - ParamListInfo paramListInfo) +LocalPlannedStatement * +CacheLocalPlanForShardQuery(Job *currentJob, DistributedPlan *originalDistributedPlan, + ParamListInfo paramListInfo, bool *planAddedToCache) { - PlannedStmt *localPlan = GetCachedLocalPlan(task, originalDistributedPlan); - if (localPlan != NULL) + Assert(planAddedToCache); + + if (!IsLocalPlanCachingSupported(currentJob, originalDistributedPlan)) { - /* we already have a local plan */ - return; + /* Local plan caching is not supported for this job and plan */ + *planAddedToCache = false; + return NULL; } - if (list_length(task->relationShardList) == 0) + Task *task = linitial(currentJob->taskList); + LocalPlannedStatement *localPlannedStatement = FindCachedLocalPlannedStatement(task, + originalDistributedPlan); + + if (localPlannedStatement != NULL) { - /* zero shard plan, no need to cache */ - return; + /* we already have a local plan */ + *planAddedToCache = false; + return localPlannedStatement; } /* @@ -61,11 +97,35 @@ CacheLocalPlanForShardQuery(Task *task, DistributedPlan *originalDistributedPlan MemoryContext oldContext = MemoryContextSwitchTo(GetMemoryChunkContext(originalDistributedPlan)); + localPlannedStatement = CreatePlanForLocalCache(currentJob, originalDistributedPlan, + paramListInfo); + + if (localPlannedStatement) + { + *planAddedToCache = true; + AddLocalPlanToCache(&originalDistributedPlan->workerJob->localPlannedStatements, + localPlannedStatement); + } + else + { + *planAddedToCache = false; + } + + MemoryContextSwitchTo(oldContext); + return localPlannedStatement; +} + + +static LocalPlannedStatement * +CreatePlanForLocalCache(Job *currentJob, DistributedPlan *originalDistributedPlan, + ParamListInfo paramListInfo) +{ /* * We prefer to use jobQuery (over task->query) because we don't want any * functions/params to have been evaluated in the cached plan. */ Query *jobQuery = copyObject(originalDistributedPlan->workerJob->jobQuery); + Task *task = linitial(currentJob->taskList); Query *localShardQuery = GetLocalShardQueryForCache(jobQuery, task, paramListInfo); @@ -82,23 +142,20 @@ CacheLocalPlanForShardQuery(Task *task, DistributedPlan *originalDistributedPlan { pfree(jobQuery); pfree(localShardQuery); - MemoryContextSwitchTo(oldContext); - return; + + return NULL; } LockRelationOid(rangeTableEntry->relid, lockMode); LocalPlannedStatement *localPlannedStatement = CitusMakeNode(LocalPlannedStatement); - localPlan = planner(localShardQuery, NULL, 0, NULL); + + PlannedStmt *localPlan = planner(localShardQuery, NULL, 0, NULL); localPlannedStatement->localPlan = localPlan; localPlannedStatement->shardId = task->anchorShardId; localPlannedStatement->localGroupId = GetLocalGroupId(); - originalDistributedPlan->workerJob->localPlannedStatements = - lappend(originalDistributedPlan->workerJob->localPlannedStatements, - localPlannedStatement); - - MemoryContextSwitchTo(oldContext); + return localPlannedStatement; } @@ -226,13 +283,7 @@ ExtractParameterTypesForParamListInfo(ParamListInfo originalParamListInfo, } -/* - * GetCachedLocalPlan is a helper function which return the cached - * plan in the distributedPlan for the given task if exists. - * - * Otherwise, the function returns NULL. - */ -PlannedStmt * +LocalPlannedStatement * GetCachedLocalPlan(Task *task, DistributedPlan *distributedPlan) { if (distributedPlan == NULL || distributedPlan->workerJob == NULL) @@ -240,12 +291,32 @@ GetCachedLocalPlan(Task *task, DistributedPlan *distributedPlan) return NULL; } - if (list_length(distributedPlan->workerJob->taskList) != 1) + if (JobHasMultipleTasks(distributedPlan->workerJob)) { - /* we only support plan caching for single shard queries */ + /* + * If there are multiple tasks in the job; i.e. multishard query, + * we do not cache the local plan. + */ return NULL; } + return FindCachedLocalPlannedStatement(task, distributedPlan); +} + + +/* + * FindCachedLocalPlannedStatement is a helper function which return the + * cached plan in the distributedPlan for the given task if exists. + * + * It's the caller's duty to ensure if the job is eligible for plan + * caching before calling this function. This function is simply + * scanning and matching the plan in the cache list. + * + * If found, returns the plan otherwise returns NULL. + */ +static LocalPlannedStatement * +FindCachedLocalPlannedStatement(Task *task, DistributedPlan *distributedPlan) +{ List *cachedPlanList = distributedPlan->workerJob->localPlannedStatements; LocalPlannedStatement *localPlannedStatement = NULL; @@ -257,7 +328,7 @@ GetCachedLocalPlan(Task *task, DistributedPlan *distributedPlan) localPlannedStatement->localGroupId == localGroupId) { /* already have a cached plan, no need to continue */ - return localPlannedStatement->localPlan; + return localPlannedStatement; } } @@ -265,23 +336,28 @@ GetCachedLocalPlan(Task *task, DistributedPlan *distributedPlan) } +static List * +AddLocalPlanToCache(List **localPlannedStatementsList, + LocalPlannedStatement *localPlannedStatement) +{ + *localPlannedStatementsList = + lappend(*localPlannedStatementsList, + localPlannedStatement); + + return *localPlannedStatementsList; +} + + /* - * IsLocalPlanCachingSupported returns whether (part of) the task can be planned - * and executed locally and whether caching is supported (single shard, no volatile - * functions). + * IsJobEligibleForPlanCache checks whether the given job is eligible for local + * plan caching. */ -bool -IsLocalPlanCachingSupported(Job *currentJob, DistributedPlan *originalDistributedPlan) +static bool +IsJobEligibleForPlanCache(Job *job) { - if (originalDistributedPlan->numberOfTimesExecuted < 1) - { - /* - * Only cache if a plan is being reused (via a prepared statement). - */ - return false; - } + Assert(job); - if (!currentJob->deferredPruning) + if (!job->deferredPruning) { /* * When not using deferred pruning we may have already replaced distributed @@ -296,26 +372,81 @@ IsLocalPlanCachingSupported(Job *currentJob, DistributedPlan *originalDistribute return false; } - List *taskList = currentJob->taskList; - if (list_length(taskList) != 1) + if (!JobHasSingleTask(job)) { /* we only support plan caching for single shard queries */ return false; } - Task *task = linitial(taskList); + Task *task = linitial(job->taskList); if (!TaskAccessesLocalNode(task)) { /* not a local task */ return false; } + if (list_length(task->relationShardList) == 0) + { + /* zero shard plan, no need to cache */ + return false; + } + + return true; +} + + +/* + * IsLocalPlanCachingSupported returns whether (part of) the task can be planned + * and executed locally and whether caching is supported (single shard, no volatile + * functions). + */ +static bool +IsLocalPlanCachingSupported(Job *currentJob, DistributedPlan *originalDistributedPlan) +{ + Assert(originalDistributedPlan); + Assert(originalDistributedPlan->workerJob); + Assert(originalDistributedPlan->workerJob->jobQuery); + + /* + * It's sensible to make this the fist validation criteria because if + * local execution is disabled, there's no point in checking further. + */ if (!EnableLocalExecution) { /* user requested not to use local execution */ return false; } + /* Checking Plan Job */ + if (JobHasMultipleTasks(originalDistributedPlan->workerJob)) + { + /* + * More than one task means a genuinely multi-shard execution; do not + * reuse a single-shard cached plan for it (citusdata/citus#8330). A + * length of 0 must still fall through to the lookup below: for + * deferred-pruning jobs, the persistent originalDistributedPlan's + * workerJob->taskList is never populated (only the per-execution + * copy's is regenerated), so 0 here says nothing about the number + * of tasks in the current execution. + */ + return false; + } + + /* Checking current job - not the same as the original distributed plan's worker job */ + if (!IsJobEligibleForPlanCache(currentJob)) + { + /* Job is not eligible for local plan caching */ + return false; + } + + if (originalDistributedPlan->numberOfTimesExecuted < 1) + { + /* + * Only cache if a plan is being reused (via a prepared statement). + */ + return false; + } + if (GetCurrentLocalExecutionStatus() == LOCAL_EXECUTION_DISABLED) { /* transaction already connected to localhost */ diff --git a/src/backend/distributed/test/local_plan_cache.c b/src/backend/distributed/test/local_plan_cache.c new file mode 100644 index 00000000000..b85589f5711 --- /dev/null +++ b/src/backend/distributed/test/local_plan_cache.c @@ -0,0 +1,73 @@ +/*------------------------------------------------------------------------- + * + * local_plan_cache.c + * + * This file contains functions to test local plan caching. + * + * Copyright (c) Citus Data, Inc. + * + *------------------------------------------------------------------------- + */ + +#include "postgres.h" + +#include "fmgr.h" + +#include "commands/prepare.h" +#include "utils/builtins.h" + +#include "distributed/citus_custom_scan.h" +#include "distributed/distributed_planner.h" +#include "distributed/listutils.h" +#include "distributed/multi_physical_planner.h" + + +PG_FUNCTION_INFO_V1(local_plan_cache_entry_count); + + +/* + * local_plan_cache_entry_count returns the number of local plans cached by the + * generic plan of the named prepared statement. + */ +Datum +local_plan_cache_entry_count(PG_FUNCTION_ARGS) +{ + char *statementName = text_to_cstring(PG_GETARG_TEXT_PP(0)); + PreparedStatement *preparedStatement = FetchPreparedStatement(statementName, true); + CachedPlan *genericPlan = preparedStatement->plansource->gplan; + + if (genericPlan == NULL || !genericPlan->is_valid) + { + ereport(ERROR, (errmsg("prepared statement \"%s\" has no valid generic plan", + statementName))); + } + + int entryCount = 0; + bool foundDistributedPlan = false; + PlannedStmt *plannedStatement = NULL; + foreach_declared_ptr(plannedStatement, genericPlan->stmt_list) + { + CustomScan *customScan = + FetchCitusCustomScanIfExists(plannedStatement->planTree); + if (customScan == NULL) + { + continue; + } + + DistributedPlan *distributedPlan = GetDistributedPlan(customScan); + if (distributedPlan->workerJob != NULL) + { + foundDistributedPlan = true; + entryCount += list_length( + distributedPlan->workerJob->localPlannedStatements); + } + } + + if (!foundDistributedPlan) + { + ereport(ERROR, (errmsg("prepared statement \"%s\" has no distributed plan", + statementName))); + } + + PG_RETURN_INT32(entryCount); +} diff --git a/src/include/distributed/local_plan_cache.h b/src/include/distributed/local_plan_cache.h index 510e7b706e2..4e60c149a67 100644 --- a/src/include/distributed/local_plan_cache.h +++ b/src/include/distributed/local_plan_cache.h @@ -1,11 +1,12 @@ #ifndef LOCAL_PLAN_CACHE #define LOCAL_PLAN_CACHE -extern bool IsLocalPlanCachingSupported(Job *currentJob, - DistributedPlan *originalDistributedPlan); -extern PlannedStmt * GetCachedLocalPlan(Task *task, DistributedPlan *distributedPlan); -extern void CacheLocalPlanForShardQuery(Task *task, - DistributedPlan *originalDistributedPlan, - ParamListInfo paramListInfo); +extern LocalPlannedStatement * GetCachedLocalPlan(Task *task, + DistributedPlan *distributedPlan); +extern LocalPlannedStatement * CacheLocalPlanForShardQuery(Job *currentJob, + DistributedPlan * + originalDistributedPlan, + ParamListInfo paramListInfo, + bool *planAddedToCache); #endif /* LOCAL_PLAN_CACHE */ diff --git a/src/test/regress/expected/citus_stat_tenants.out b/src/test/regress/expected/citus_stat_tenants.out index 721e9eaf217..dd70bb4abc1 100644 --- a/src/test/regress/expected/citus_stat_tenants.out +++ b/src/test/regress/expected/citus_stat_tenants.out @@ -1145,6 +1145,13 @@ ALTER TABLE referencing ADD CONSTRAINT fkey FOREIGN KEY (shard_key, other_key) R INSERT INTO referenced VALUES (0, 1), (0, 2), (1, 2); \c - - - :worker_2_port SET search_path TO citus_stat_tenants; +SET citus.enable_metadata_sync TO OFF; +CREATE FUNCTION local_plan_cache_entry_count(text) +RETURNS integer +AS 'citus', $$local_plan_cache_entry_count$$ +LANGUAGE C STRICT; +RESET citus.enable_metadata_sync; +SET plan_cache_mode TO force_generic_plan; PREPARE prep_stmt (bigint, int, bigint, int) AS INSERT INTO referencing (shard_key, other_key) VALUES ($1, $2), ($3, $4); EXECUTE prep_stmt(0, 1, 0, 2); EXECUTE prep_stmt(0, 1, 0, 2); @@ -1153,7 +1160,21 @@ EXECUTE prep_stmt(0, 1, 0, 2); EXECUTE prep_stmt(0, 1, 0, 2); EXECUTE prep_stmt(0, 1, 0, 2); EXECUTE prep_stmt(0, 1, 0, 2); +-- Repeated executions for the same shard should reuse one local plan. +SELECT local_plan_cache_entry_count('prep_stmt'); + local_plan_cache_entry_count +--------------------------------------------------------------------- + 1 +(1 row) + EXECUTE prep_stmt(0, 1, 1, 2); -- multi-shard query shouldn't use local cache and fail +-- A multi-shard execution should not add or reuse a single-shard local plan. +SELECT local_plan_cache_entry_count('prep_stmt'); + local_plan_cache_entry_count +--------------------------------------------------------------------- + 1 +(1 row) + SELECT shard_key, other_key, count(1) FROM referencing GROUP BY shard_key, other_key @@ -1165,6 +1186,11 @@ ORDER BY shard_key, other_key; 1 | 2 | 1 (3 rows) +DEALLOCATE prep_stmt; +RESET plan_cache_mode; +SET citus.enable_metadata_sync TO OFF; +DROP FUNCTION local_plan_cache_entry_count(text); +RESET citus.enable_metadata_sync; \c - - - :master_port SET client_min_messages TO ERROR; DROP SCHEMA citus_stat_tenants CASCADE; diff --git a/src/test/regress/sql/citus_stat_tenants.sql b/src/test/regress/sql/citus_stat_tenants.sql index 7ff68c181da..d47aa626b06 100644 --- a/src/test/regress/sql/citus_stat_tenants.sql +++ b/src/test/regress/sql/citus_stat_tenants.sql @@ -436,6 +436,15 @@ INSERT INTO referenced VALUES (0, 1), (0, 2), (1, 2); \c - - - :worker_2_port SET search_path TO citus_stat_tenants; +SET citus.enable_metadata_sync TO OFF; +CREATE FUNCTION local_plan_cache_entry_count(text) +RETURNS integer +AS 'citus', $$local_plan_cache_entry_count$$ +LANGUAGE C STRICT; +RESET citus.enable_metadata_sync; + +SET plan_cache_mode TO force_generic_plan; + PREPARE prep_stmt (bigint, int, bigint, int) AS INSERT INTO referencing (shard_key, other_key) VALUES ($1, $2), ($3, $4); EXECUTE prep_stmt(0, 1, 0, 2); @@ -445,13 +454,26 @@ EXECUTE prep_stmt(0, 1, 0, 2); EXECUTE prep_stmt(0, 1, 0, 2); EXECUTE prep_stmt(0, 1, 0, 2); EXECUTE prep_stmt(0, 1, 0, 2); + +-- Repeated executions for the same shard should reuse one local plan. +SELECT local_plan_cache_entry_count('prep_stmt'); + EXECUTE prep_stmt(0, 1, 1, 2); -- multi-shard query shouldn't use local cache and fail +-- A multi-shard execution should not add or reuse a single-shard local plan. +SELECT local_plan_cache_entry_count('prep_stmt'); + SELECT shard_key, other_key, count(1) FROM referencing GROUP BY shard_key, other_key ORDER BY shard_key, other_key; +DEALLOCATE prep_stmt; +RESET plan_cache_mode; +SET citus.enable_metadata_sync TO OFF; +DROP FUNCTION local_plan_cache_entry_count(text); +RESET citus.enable_metadata_sync; + \c - - - :master_port SET client_min_messages TO ERROR; DROP SCHEMA citus_stat_tenants CASCADE;