From ea3c5ca84aa1921a453bae4b6d26cccd84c265ed Mon Sep 17 00:00:00 2001 From: Colm McHugh Date: Fri, 27 Mar 2026 10:40:32 +0000 Subject: [PATCH 1/2] Cahce plans for fast path prepared statements on workers. DESCRIPTION: Adds citus.enable_prepared_statement_caching to cache prepared statement plans on worker connections for fast-path queries Repeated executions of a prepared statement re-parse and re-plan the shard query on the worker every time, because Citus sends the query via PQsendQueryParams, which uses libpq's unnamed prepared statement. The plan is built and discarded on each execution. The coordinator also deparses the query tree to SQL for every execution. This adds citus.enable_prepared_statement_caching (PGC_USERSET, default off). When enabled, single-shard fast-path queries are sent using named prepared statements: PQprepare on first use, then PQsendQueryPrepared for subsequent executions. The worker keeps the CachedPlanSource/CachedPlan for the lifetime of the connection, so later executions skip parse and plan entirely and go straight to bind/execute. This complements Citus's local plan caching, which already caches prepared statement plans for local execution ((see `local_plan_cache.c`, `local_executor.c`), by caching plans for remote execution so achieving plan caching for all executions of a prepared statement for a fast path query. Cache ----- Each MultiConnection gains a hash table keyed by (planId, shardId) mapping to the statement name prepared on that connection. The cache holds only names and parameter types (~264 bytes/entry); the plans themselves live in the worker backend. It is capped at `MAX_CACHED_STMTS_PER_CONNECTION` entries; beyond that, executions fall back to plain parameterized SQL. Entries are dropped with the connection, and a lost connection simply re-prepares on reconnect. Design ------ Caching logic is encapsulated in prepared_statement_cache.c, which owns both the hash table and the decisions about when caching applies. Call sites invoke the API unconditionally and carry no policy of their own; each entry point returns early when the GUC is off: `PreparedStatementCacheTryFastPath()` build a Task straight from ParamListInfo, skipping replan `PreparedStatementCacheSaveTemplate()` save the parameterized job query on the original plan `PreparedStatementCacheAttachToTasks()` stamp planId and template onto each task `PreparedStatementCacheSendQuery()` lookup, prepare and dispatch on a worker connection `SendQuery()` reports its outcome as a PreparedStatementSendStatus value so the executor can distinguish a dispatched query from a lost connection or a cache-full fallback, for which the deparsed SQL is returned to the caller to send down the existing path. Executor -------- `SendNextQuery()` consults the cache before the existing parameterized-SQL path. Because parameters must survive to be bound remotely, `MarkUnreferencedExternParams()` is skipped when caching is enabled. Fast path --------- `CitusBeginReadOnlyScan()` and `CitusBeginModifyScan()` try the cache-hit fast path on second and subsequent executions of a generic plan. Instead of deep-copying the `DistributedPlan`, evaluating coordinator expressions and regenerating the task, the distribution key value is read directly from `ParamListInfo`, the shard is found with `FindShardInterval()`, and a minimal Task is built. The parameterized job query is saved once on the original plan (Job->savedJobQueryForCaching) and reused, avoiding a per-execution `copyObject()`. DML additionally acquires metadata locks and assigns the first replica. Applies to single-shard router queries with a parameterized distribution key: SELECT, UPDATE, DELETE and single-row INSERT. Multi-row INSERT is excluded: each shard's task carries only its own subset of VALUES rows, so a statement deparsed once from the whole job query would send every row to every shard. Behaviour is unchanged when the GUC is off. --- .../connection/connection_management.c | 6 + .../distributed/connection/remote_commands.c | 61 ++ .../distributed/executor/adaptive_executor.c | 33 +- .../distributed/executor/citus_custom_scan.c | 96 ++- .../distributed/executor/local_executor.c | 34 +- .../executor/prepared_statement_cache.c | 535 ++++++++++++++++ src/backend/distributed/metadata/XXnahQZC | 0 .../distributed/planner/deparse_shard_query.c | 29 + .../distributed/planner/distributed_planner.c | 19 +- .../planner/fast_path_router_planner.c | 2 + .../planner/multi_router_planner.c | 28 + src/backend/distributed/shared_library_init.c | 13 + src/backend/distributed/utils/citus_clauses.c | 29 +- .../distributed/utils/citus_copyfuncs.c | 7 + .../distributed/utils/citus_outfuncs.c | 4 + src/include/distributed/citus_clauses.h | 4 + src/include/distributed/citus_custom_scan.h | 1 + .../distributed/connection_management.h | 3 + src/include/distributed/distributed_planner.h | 8 + .../distributed/multi_physical_planner.h | 16 + .../distributed/prepared_statement_cache.h | 102 +++ src/include/distributed/remote_commands.h | 5 + src/test/regress/citus_tests/run_test.py | 3 + .../expected/prepared_statement_caching.out | 606 ++++++++++++++++++ src/test/regress/multi_schedule | 1 + .../sql/prepared_statement_caching.sql | 306 +++++++++ 26 files changed, 1935 insertions(+), 16 deletions(-) create mode 100644 src/backend/distributed/executor/prepared_statement_cache.c create mode 100644 src/backend/distributed/metadata/XXnahQZC create mode 100644 src/include/distributed/prepared_statement_cache.h create mode 100644 src/test/regress/expected/prepared_statement_caching.out create mode 100644 src/test/regress/sql/prepared_statement_caching.sql diff --git a/src/backend/distributed/connection/connection_management.c b/src/backend/distributed/connection/connection_management.c index 0cb2b7c30c1..1491110d513 100644 --- a/src/backend/distributed/connection/connection_management.c +++ b/src/backend/distributed/connection/connection_management.c @@ -36,6 +36,7 @@ #include "distributed/memutils.h" #include "distributed/metadata_cache.h" #include "distributed/placement_connection.h" +#include "distributed/prepared_statement_cache.h" #include "distributed/remote_commands.h" #include "distributed/run_from_same_connection.h" #include "distributed/shared_connection_stats.h" @@ -794,6 +795,9 @@ ShutdownConnection(MultiConnection *connection) { SendCancelationRequest(connection); } + + PreparedStatementCacheDestroy(&connection->preparedStatementCache); + CitusPQFinish(connection); } @@ -1232,6 +1236,8 @@ CloseNotReadyMultiConnectionStates(List *connectionStates) static void CitusPQFinish(MultiConnection *connection) { + PreparedStatementCacheDestroy(&connection->preparedStatementCache); + if (connection->pgConn != NULL) { PQfinish(connection->pgConn); diff --git a/src/backend/distributed/connection/remote_commands.c b/src/backend/distributed/connection/remote_commands.c index cd0c1951999..ba3c71dead8 100644 --- a/src/backend/distributed/connection/remote_commands.c +++ b/src/backend/distributed/connection/remote_commands.c @@ -576,6 +576,67 @@ SendRemoteCommand(MultiConnection *connection, const char *command) } +/* + * SendRemotePrepare wraps PQprepare() to prepare a statement on a worker + * connection. This is a synchronous call that blocks until the prepare + * completes. Returns 1 on success, 0 on failure. + */ +int +SendRemotePrepare(MultiConnection *connection, const char *stmtName, + const char *query, int nParams, const Oid *paramTypes) +{ + PGconn *pgConn = connection->pgConn; + + LogRemoteCommand(connection, query); + + if (!pgConn || PQstatus(pgConn) != CONNECTION_OK) + { + return 0; + } + + PGresult *result = PQprepare(pgConn, stmtName, query, nParams, paramTypes); + if (PQresultStatus(result) != PGRES_COMMAND_OK) + { + ReportResultError(connection, result, WARNING); + PQclear(result); + return 0; + } + + PQclear(result); + return 1; +} + + +/* + * SendRemotePreparedQuery wraps PQsendQueryPrepared() to asynchronously + * execute a previously prepared statement on a worker connection. Follows + * the same pattern as SendRemoteCommandParams(). Returns the result of + * PQsendQueryPrepared (1 on success, 0 on failure). + */ +int +SendRemotePreparedQuery(MultiConnection *connection, const char *stmtName, + int nParams, const char *const *paramValues, + bool binaryResults) +{ + PGconn *pgConn = connection->pgConn; + + /* log the statement name for debugging */ + LogRemoteCommand(connection, stmtName); + + if (!pgConn || PQstatus(pgConn) != CONNECTION_OK) + { + return 0; + } + + Assert(PQisnonblocking(pgConn)); + + int rc = PQsendQueryPrepared(pgConn, stmtName, nParams, paramValues, + NULL, NULL, binaryResults ? 1 : 0); + + return rc; +} + + /* * ExecuteRemoteCommandAndCheckResult executes the given command in the remote node and * checks if the result is equal to the expected result. If the result is equal to the diff --git a/src/backend/distributed/executor/adaptive_executor.c b/src/backend/distributed/executor/adaptive_executor.c index 443e9ab2ac5..3f97049f4cd 100644 --- a/src/backend/distributed/executor/adaptive_executor.c +++ b/src/backend/distributed/executor/adaptive_executor.c @@ -149,6 +149,7 @@ #include "distributed/backend_data.h" #include "distributed/cancel_utils.h" #include "distributed/citus_custom_scan.h" +#include "distributed/citus_ruleutils.h" #include "distributed/citus_safe_lib.h" #include "distributed/commands/multi_copy.h" #include "distributed/connection_management.h" @@ -166,6 +167,7 @@ #include "distributed/param_utils.h" #include "distributed/placement_access.h" #include "distributed/placement_connection.h" +#include "distributed/prepared_statement_cache.h" #include "distributed/relation_access_tracking.h" #include "distributed/remote_commands.h" #include "distributed/repartition_join_execution.h" @@ -924,8 +926,13 @@ AdaptiveExecutorStart(CitusScanState *scanState) * and never used in the query, mark such parameters' type as Invalid(0), * which will be used later in ExtractParametersFromParamList() to map them * to a generic datatype. Skip for dynamic parameters. + * + * When prepared statement caching is enabled, skip this step entirely: + * the params appear "unreferenced" in job->jobQuery because they were + * resolved to constants there, but they are still needed with their + * original types for the parameterized query sent via PQprepare. */ - if (paramListInfo && !paramListInfo->paramFetch) + if (paramListInfo && !paramListInfo->paramFetch && !EnablePreparedStatementCaching) { paramListInfo = copyParamList(paramListInfo); MarkUnreferencedExternParams((Node *) job->jobQuery, paramListInfo); @@ -4256,7 +4263,29 @@ SendNextQuery(TaskPlacementExecution *placementExecution, uint32 queryIndex = placementExecution->queryIndex; Assert(queryIndex < task->queryCount); - char *queryString = TaskQueryStringAtIndex(task, queryIndex); + char *queryString = NULL; + + /* + * Must be attempted before the paramListInfo path below: a fast-path task + * reports its parameters as resolved, so that path would send plain SQL. + */ + PreparedStatementSendStatus cacheStatus = + PreparedStatementCacheSendQuery(connection, task, paramListInfo, + binaryResults, &queryString); + if (cacheStatus == PREPARED_STMT_SENT) + { + return true; + } + else if (cacheStatus == PREPARED_STMT_FAILED) + { + return false; + } + + /* resolve queryString if not already set (e.g. from cache-full fallback) */ + if (queryString == NULL) + { + queryString = TaskQueryStringAtIndex(task, queryIndex); + } if (paramListInfo != NULL && !task->parametersInQueryStringResolved) { diff --git a/src/backend/distributed/executor/citus_custom_scan.c b/src/backend/distributed/executor/citus_custom_scan.c index 87fca7422b1..c58c79cc481 100644 --- a/src/backend/distributed/executor/citus_custom_scan.c +++ b/src/backend/distributed/executor/citus_custom_scan.c @@ -43,10 +43,13 @@ #include "distributed/local_plan_cache.h" #include "distributed/merge_executor.h" #include "distributed/merge_planner.h" +#include "distributed/metadata_cache.h" #include "distributed/multi_executor.h" #include "distributed/multi_router_planner.h" #include "distributed/multi_server_executor.h" +#include "distributed/prepared_statement_cache.h" #include "distributed/shard_utils.h" +#include "distributed/shardinterval_utils.h" #include "distributed/sorted_merge.h" #include "distributed/stats/query_stats.h" #include "distributed/stats/stat_counters.h" @@ -82,7 +85,7 @@ static void SortedMergeReScan(CustomScanState *node); static void CitusEndScanCommon(CitusScanState *scanState); static void CitusReScanCommon(CustomScanState *node); static void EnsureForceDelegationDistributionKey(Job *job); -static void EnsureAnchorShardsInJobExist(Job *job); + static bool AnchorShardsInTaskListExist(List *taskList); static void TryToRerouteFastPathModifyQuery(Job *job); static void CheckQueryDeparseSafety(Query *query); @@ -355,6 +358,28 @@ CitusBeginReadOnlyScan(CustomScanState *node, EState *estate, int eflags) return; } + Job *workerJob = originalDistributedPlan->workerJob; + + /* + * Always clear stale task pointers from the previous fast-path execution. + * For deferred-pruning plans, workerJob->taskList is NIL until populated + * per-execution by RegenerateTaskForFasthPathQuery (or the fast path below). + * Without this reset, if a previous execution took the fast path and then the + * GUC is disabled, CopyDistributedPlanWithoutCache would deep-copy a stale + * taskList pointing into freed memory. + */ + workerJob->taskList = NIL; + workerJob->parametersInJobQueryResolved = false; + + /* + * A cached plan can build its task straight from the bound parameters, + * skipping the plan copy, coordinator evaluation and task regeneration. + */ + if (PreparedStatementCacheTryFastPath(scanState, estate, false)) + { + return; + } + /* * Create a copy of the generic plan for the current execution, but make a shallow * copy of the plan cache. That means we'll be able to access the plan cache via @@ -365,8 +390,8 @@ CitusBeginReadOnlyScan(CustomScanState *node, EState *estate, int eflags) CopyDistributedPlanWithoutCache(originalDistributedPlan); scanState->distributedPlan = currentPlan; - Job *workerJob = currentPlan->workerJob; - Query *jobQuery = workerJob->jobQuery; + Job *currentJob = currentPlan->workerJob; + Query *jobQuery = currentJob->jobQuery; PlanState *planState = &(scanState->customScanState.ss.ps); /* @@ -375,6 +400,8 @@ CitusBeginReadOnlyScan(CustomScanState *node, EState *estate, int eflags) */ Assert(currentPlan->fastPathRouterPlan || !EnableFastPathRouterPlanner); + Query *savedJobQuery = PreparedStatementCacheSaveTemplate(originalDistributedPlan); + /* * Evaluate parameters, because the parameters are only available on the * coordinator and are required for pruning. @@ -388,14 +415,16 @@ CitusBeginReadOnlyScan(CustomScanState *node, EState *estate, int eflags) ExecuteCoordinatorEvaluableExpressions(jobQuery, planState); /* job query no longer has parameters, so we should not send any */ - workerJob->parametersInJobQueryResolved = true; + currentJob->parametersInJobQueryResolved = true; /* parameters are filled in, so we can generate a task for this execution */ - RegenerateTaskForFasthPathQuery(workerJob); + RegenerateTaskForFasthPathQuery(currentJob); - if (IsLocalPlanCachingSupported(workerJob, originalDistributedPlan)) + PreparedStatementCacheAttachToTasks(currentPlan, currentJob, savedJobQuery); + + if (IsLocalPlanCachingSupported(currentJob, originalDistributedPlan)) { - Task *task = linitial(workerJob->taskList); + Task *task = linitial(currentJob->taskList); /* * We are going to execute this task locally. If it's not already in @@ -426,6 +455,28 @@ CitusBeginModifyScan(CustomScanState *node, EState *estate, int eflags) CitusScanState *scanState = (CitusScanState *) node; PlanState *planState = &(scanState->customScanState.ss.ps); DistributedPlan *originalDistributedPlan = scanState->distributedPlan; + Job *origWorkerJob = originalDistributedPlan->workerJob; + + /* + * Safety: clear stale task pointers from a previous fast-path execution. + * For deferred-pruning plans, workerJob->taskList is NIL until populated + * per-execution. Without this reset, CopyDistributedPlanWithoutCache + * would deep-copy a stale taskList pointing into freed memory. + */ + if (origWorkerJob->deferredPruning) + { + origWorkerJob->taskList = NIL; + origWorkerJob->parametersInJobQueryResolved = false; + } + + /* + * A cached plan can build its task straight from the bound parameters, + * skipping the plan copy, coordinator evaluation and task regeneration. + */ + if (PreparedStatementCacheTryFastPath(scanState, estate, true)) + { + return; + } MemoryContext localContext = AllocSetContextCreate(CurrentMemoryContext, "CitusBeginModifyScan", @@ -440,6 +491,8 @@ CitusBeginModifyScan(CustomScanState *node, EState *estate, int eflags) Query *jobQuery = workerJob->jobQuery; + Query *savedJobQuery = PreparedStatementCacheSaveTemplate(originalDistributedPlan); + if (ModifyJobNeedsEvaluation(workerJob)) { ExecuteCoordinatorEvaluableExpressions(jobQuery, planState); @@ -472,6 +525,8 @@ CitusBeginModifyScan(CustomScanState *node, EState *estate, int eflags) { RegenerateTaskForFasthPathQuery(workerJob); } + + PreparedStatementCacheAttachToTasks(currentPlan, workerJob, savedJobQuery); } else if (workerJob->requiresCoordinatorEvaluation) { @@ -574,7 +629,7 @@ TryToRerouteFastPathModifyQuery(Job *job) * EnsureAnchorShardsInJobExist ensures all shards are valid in job. * If it finds a non-existent shard in given job, it throws an error. */ -static void +void EnsureAnchorShardsInJobExist(Job *job) { if (!AnchorShardsInTaskListExist(job->taskList)) @@ -648,13 +703,17 @@ CopyDistributedPlanWithoutCache(DistributedPlan *originalDistributedPlan) { List *localPlannedStatements = originalDistributedPlan->workerJob->localPlannedStatements; + Query *savedJobQueryForCaching = + originalDistributedPlan->workerJob->savedJobQueryForCaching; originalDistributedPlan->workerJob->localPlannedStatements = NIL; + originalDistributedPlan->workerJob->savedJobQueryForCaching = NULL; DistributedPlan *distributedPlan = copyObject(originalDistributedPlan); - /* set back the immutable field */ + /* set back the immutable/cached fields */ originalDistributedPlan->workerJob->localPlannedStatements = localPlannedStatements; distributedPlan->workerJob->localPlannedStatements = localPlannedStatements; + originalDistributedPlan->workerJob->savedJobQueryForCaching = savedJobQueryForCaching; return distributedPlan; } @@ -915,6 +974,25 @@ CitusEndScanCommon(CitusScanState *scanState) /* queries without partition key are also recorded */ CitusQueryStatsExecutorsEntry(queryId, executorType, partitionKeyString); } + + /* + * Clear mutable per-execution state so the cached plan is clean for + * the next execution. The cache-hit fast paths in CitusBeginReadOnlyScan() + * and CitusBeginModifyScan() store a Task list directly on the original + * plan's workerJob; those Tasks live in the per-execution memory context + * and become dangling after EndScan. In assert-checking builds the next + * execution's GetDistributedPlan() → copyObject() would traverse freed + * memory without this reset. + * + * Only deferred-pruning plans need this: their taskList is rebuilt + * per-execution. Non-deferred plans carry their real taskList from + * planning and must not be touched. + */ + if (workerJob != NULL && workerJob->deferredPruning) + { + workerJob->taskList = NIL; + workerJob->parametersInJobQueryResolved = false; + } } diff --git a/src/backend/distributed/executor/local_executor.c b/src/backend/distributed/executor/local_executor.c index 5480a1d142d..bb3d9d02f19 100644 --- a/src/backend/distributed/executor/local_executor.c +++ b/src/backend/distributed/executor/local_executor.c @@ -392,7 +392,39 @@ ExecuteLocalTaskListExtended(List *taskList, continue; } - if (taskType != TASK_QUERY_LOCAL_PLAN) + if (taskType == TASK_QUERY_NULL && task->jobQueryForPrepare != NULL) + { + /* + * Fast-path task from prepared statement caching: the task + * was built without a query string to avoid expensive + * deparsing. Deparse from the saved job query template now + * so we can plan locally. + */ + Query *queryForDeparse = copyObject(task->jobQueryForPrepare); + StringInfoData buf; + initStringInfo(&buf); + + if (queryForDeparse->commandType == CMD_INSERT) + { + deparse_shard_query(queryForDeparse, + task->anchorDistributedTableId, + task->anchorShardId, &buf); + } + else + { + UpdateRelationToShardNames((Node *) queryForDeparse, + task->relationShardList); + pg_get_query_def(queryForDeparse, &buf); + } + + Query *shardQuery = ParseQueryString(buf.data, + taskParameterTypes, + taskNumParams); + localPlan = planner(shardQuery, NULL, CURSOR_OPT_PARALLEL_OK, + paramListInfo); + pfree(buf.data); + } + else if (taskType != TASK_QUERY_LOCAL_PLAN) { Query *shardQuery = ParseQueryString(TaskQueryString(task), taskParameterTypes, diff --git a/src/backend/distributed/executor/prepared_statement_cache.c b/src/backend/distributed/executor/prepared_statement_cache.c new file mode 100644 index 00000000000..d24230752f8 --- /dev/null +++ b/src/backend/distributed/executor/prepared_statement_cache.c @@ -0,0 +1,535 @@ +/*------------------------------------------------------------------------- + * + * prepared_statement_cache.c + * Per-connection cache for prepared statements on worker connections. + * + * When citus.enable_prepared_statement_caching is ON, the coordinator + * uses PQprepare/PQsendQueryPrepared on worker connections instead + * of PQsendQuery for fast-path prepared statement executions (generic + * plan, execution 6+). This module manages the per-connection hash + * table that tracks which statements have already been prepared on + * each connection. + * + * It also owns the two integration points that make the cache usable: + * the planner-side fast path that builds a Task straight from + * ParamListInfo (skipping replanning), and the executor-side dispatch + * that prepares or reuses the named statement on a connection. + * + * Copyright (c) Citus Data, Inc. + * + *------------------------------------------------------------------------- + */ + +#include "postgres.h" + +#include "access/htup_details.h" +#include "nodes/makefuncs.h" +#include "utils/hsearch.h" +#include "utils/lsyscache.h" +#include "utils/memutils.h" + +#include "distributed/citus_custom_scan.h" +#include "distributed/citus_ruleutils.h" +#include "distributed/citus_safe_lib.h" +#include "distributed/deparse_shard_query.h" +#include "distributed/distributed_execution_locks.h" +#include "distributed/executor_util.h" +#include "distributed/listutils.h" +#include "distributed/local_plan_cache.h" +#include "distributed/metadata_cache.h" +#include "distributed/multi_executor.h" +#include "distributed/multi_router_planner.h" +#include "distributed/prepared_statement_cache.h" +#include "distributed/remote_commands.h" +#include "distributed/shard_cleaner.h" +#include "distributed/shardinterval_utils.h" + + +/* GUC: citus.enable_prepared_statement_caching */ +bool EnablePreparedStatementCaching = false; + + +/* + * PreparedStatementCacheCreate allocates a new hash table for caching + * prepared statement entries on a single worker connection. The hash + * table is allocated in TopMemoryContext so it survives across + * transactions (matching the connection lifetime). + */ +HTAB * +PreparedStatementCacheCreate(void) +{ + HASHCTL info; + + memset(&info, 0, sizeof(info)); + info.keysize = sizeof(PreparedStatementCacheKey); + info.entrysize = sizeof(PreparedStatementCacheEntry); + info.hcxt = TopMemoryContext; + + HTAB *cache = hash_create("Prepared Statement Cache", + 32, /* initial size */ + &info, + HASH_ELEM | HASH_BLOBS | HASH_CONTEXT); + + return cache; +} + + +/* + * PreparedStatementCacheLookup looks up a cache entry by (planId, shardId). + * Returns the entry if found, NULL otherwise. + */ +PreparedStatementCacheEntry * +PreparedStatementCacheLookup(HTAB *cache, uint64 planId, uint64 shardId) +{ + PreparedStatementCacheKey key; + + memset(&key, 0, sizeof(key)); + key.planId = planId; + key.shardId = shardId; + + PreparedStatementCacheEntry *entry = + (PreparedStatementCacheEntry *) hash_search(cache, &key, + HASH_FIND, NULL); + + return entry; +} + + +/* + * PreparedStatementCacheInsert inserts a new entry for (planId, shardId). + * Returns the new entry on success, or NULL if the cache has reached + * MAX_CACHED_STMTS_PER_CONNECTION (caller should fall back to plain SQL). + * + * The caller is responsible for filling in the returned entry's fields + * (stmtName, paramTypes, paramCount, parameterizedQueryString). + */ +PreparedStatementCacheEntry * +PreparedStatementCacheInsert(HTAB *cache, uint64 planId, uint64 shardId) +{ + if (hash_get_num_entries(cache) >= MAX_CACHED_STMTS_PER_CONNECTION) + { + return NULL; + } + + PreparedStatementCacheKey key; + + memset(&key, 0, sizeof(key)); + key.planId = planId; + key.shardId = shardId; + + bool found = false; + PreparedStatementCacheEntry *entry = + (PreparedStatementCacheEntry *) hash_search(cache, &key, + HASH_ENTER, &found); + + if (found) + { + /* already exists — return existing entry */ + return entry; + } + + /* initialize the new entry with auto-generated statement name */ + SafeSnprintf(entry->stmtName, MAX_STMT_NAME_LENGTH, + "__citus_stmt_%ld", (long) hash_get_num_entries(cache)); + entry->paramTypes = NULL; + entry->paramCount = 0; + entry->parameterizedQueryString = NULL; + + return entry; +} + + +/* + * PreparedStatementCacheDestroy frees all memory used by the cache. + * Safe to call if the pointed to cache is NULL (NOP) + * Otherwise, frees all memory used by the cache and + * sets it to NULL. + */ +void +PreparedStatementCacheDestroy(HTAB **cache_ptr) +{ + Assert(cache_ptr != NULL); + HTAB *cache = *cache_ptr; + + if (cache == NULL) + { + return; + } + + /* + * Free dynamically allocated fields in each entry before destroying + * the hash table itself. + */ + HASH_SEQ_STATUS status; + PreparedStatementCacheEntry *entry; + + hash_seq_init(&status, cache); + while ((entry = hash_seq_search(&status)) != NULL) + { + if (entry->paramTypes != NULL) + { + pfree(entry->paramTypes); + } + if (entry->parameterizedQueryString != NULL) + { + pfree(entry->parameterizedQueryString); + } + } + + hash_destroy(cache); + *cache_ptr = NULL; +} + + +/* + * DeparseTaskTemplate deparses the task's parameterized job query into SQL + * targeting the task's shard, leaving Param nodes ($1, ...) intact. + * + * The template is shared across executions, so a working copy is deparsed. + */ +static char * +DeparseTaskTemplate(Task *task) +{ + Query *queryCopy = copyObject(task->jobQueryForPrepare); + StringInfoData buf; + initStringInfo(&buf); + + if (queryCopy->commandType == CMD_INSERT) + { + /* upserts reference the target by name, which becomes the shard name */ + AddInsertAliasIfNeeded(queryCopy); + deparse_shard_query(queryCopy, task->anchorDistributedTableId, + task->anchorShardId, &buf); + } + else + { + UpdateRelationToShardNames((Node *) queryCopy, task->relationShardList); + pg_get_query_def(queryCopy, &buf); + } + + return buf.data; +} + + +/* + * PreparedStatementCacheSendQuery dispatches a task on a worker connection + * using a named prepared statement, preparing it first if this connection has + * not seen it before. + * + * Returns PREPARED_STMT_NOT_APPLICABLE when the task carries no template, in + * which case the caller must use its normal send path. When the connection's + * cache is full it returns PREPARED_STMT_FALLBACK and sets fallbackQueryString + * to parameterized SQL for the caller to send instead. + */ +PreparedStatementSendStatus +PreparedStatementCacheSendQuery(MultiConnection *connection, Task *task, + ParamListInfo paramListInfo, bool binaryResults, + char **fallbackQueryString) +{ + if (!EnablePreparedStatementCaching || task->jobQueryForPrepare == NULL || + paramListInfo == NULL) + { + return PREPARED_STMT_NOT_APPLICABLE; + } + + Oid *parameterTypes = NULL; + const char **parameterValues = NULL; + + /* force evaluation of bound params */ + paramListInfo = copyParamList(paramListInfo); + int parameterCount = paramListInfo->numParams; + + ExtractParametersForRemoteExecution(paramListInfo, ¶meterTypes, + ¶meterValues); + + if (connection->preparedStatementCache == NULL) + { + connection->preparedStatementCache = PreparedStatementCacheCreate(); + } + + PreparedStatementCacheEntry *cacheEntry = + PreparedStatementCacheLookup(connection->preparedStatementCache, + task->preparedStatementPlanId, + task->anchorShardId); + if (cacheEntry == NULL) + { + elog(DEBUG2, "prepared statement cache miss: plan " UINT64_FORMAT + " shard " UINT64_FORMAT, + task->preparedStatementPlanId, task->anchorShardId); + + cacheEntry = PreparedStatementCacheInsert(connection->preparedStatementCache, + task->preparedStatementPlanId, + task->anchorShardId); + if (cacheEntry == NULL) + { + /* + * Cache full. The fast-path task has no query string of its own, so + * hand the caller freshly deparsed SQL. Clearing the resolved flag + * routes it through the parameterized send path. + */ + *fallbackQueryString = DeparseTaskTemplate(task); + task->parametersInQueryStringResolved = false; + return PREPARED_STMT_FALLBACK; + } + + char *queryString = DeparseTaskTemplate(task); + + if (SendRemotePrepare(connection, cacheEntry->stmtName, queryString, + parameterCount, parameterTypes) == 0) + { + connection->connectionState = MULTI_CONNECTION_LOST; + return PREPARED_STMT_FAILED; + } + + Size paramTypesSize = parameterCount * sizeof(Oid); + cacheEntry->paramTypes = MemoryContextAlloc(TopMemoryContext, paramTypesSize); + memcpy_s(cacheEntry->paramTypes, paramTypesSize, parameterTypes, + paramTypesSize); + cacheEntry->paramCount = parameterCount; + cacheEntry->parameterizedQueryString = + MemoryContextStrdup(TopMemoryContext, queryString); + + pfree(queryString); + } + else + { + elog(DEBUG2, "prepared statement cache hit: plan " UINT64_FORMAT + " shard " UINT64_FORMAT " stmt %s", + task->preparedStatementPlanId, task->anchorShardId, + cacheEntry->stmtName); + } + + if (SendRemotePreparedQuery(connection, cacheEntry->stmtName, parameterCount, + parameterValues, binaryResults) == 0) + { + connection->connectionState = MULTI_CONNECTION_LOST; + return PREPARED_STMT_FAILED; + } + + if (PQsetSingleRowMode(connection->pgConn) == 0) + { + connection->connectionState = MULTI_CONNECTION_LOST; + return PREPARED_STMT_FAILED; + } + + return PREPARED_STMT_SENT; +} + + +/* + * PreparedStatementCacheSaveTemplate returns the job query template used to + * build parameterized statements, saving it on the original plan on first use. + * + * The template must be captured before coordinator evaluation resolves Param + * nodes, and is kept on the original (cached) plan so later executions reuse + * it instead of paying for a copyObject each time. + * + * Returns NULL when the plan is not eligible for caching. + */ +Query * +PreparedStatementCacheSaveTemplate(DistributedPlan *originalPlan) +{ + Job *originalJob = originalPlan->workerJob; + Query *jobQuery = originalJob->jobQuery; + + if (!EnablePreparedStatementCaching) + { + return NULL; + } + + if (jobQuery->commandType != CMD_SELECT) + { + /* + * Multi-row INSERT can't be cached: each shard's task carries only its + * own subset of VALUES rows, but the statement is deparsed once from + * the whole job query, so every shard would receive every row. + */ + if (!originalJob->deferredPruning || + ExtractDistributedInsertValuesRTE(jobQuery) != NULL) + { + return NULL; + } + } + + if (originalJob->savedJobQueryForCaching == NULL) + { + MemoryContext oldContext = + MemoryContextSwitchTo(GetMemoryChunkContext(originalPlan)); + originalJob->savedJobQueryForCaching = copyObject(jobQuery); + MemoryContextSwitchTo(oldContext); + } + + return originalJob->savedJobQueryForCaching; +} + + +/* + * PreparedStatementCacheAttachToTasks stamps the cache key and query template + * onto each task so SendNextQuery can look them up. Does nothing when the plan + * is not eligible for caching. + */ +void +PreparedStatementCacheAttachToTasks(DistributedPlan *currentPlan, Job *workerJob, + Query *savedJobQuery) +{ + if (!EnablePreparedStatementCaching || savedJobQuery == NULL) + { + return; + } + + bool isInsert = workerJob->jobQuery->commandType == CMD_INSERT; + + Task *task = NULL; + foreach_declared_ptr(task, workerJob->taskList) + { + task->preparedStatementPlanId = currentPlan->planId; + task->jobQueryForPrepare = savedJobQuery; + + /* deparse_shard_query needs the target relation for INSERT */ + if (isInsert) + { + task->anchorDistributedTableId = linitial_oid(currentPlan->relationIdList); + } + } +} + + +/* + * FastPathShardInterval returns the shard the distribution key parameter routes + * to, or NULL if the fast path cannot be used for this execution. + */ +static ShardInterval * +FastPathShardInterval(DistributedPlan *plan, Job *workerJob, EState *estate) +{ + int paramId = workerJob->distributionKeyParamId; + ParamListInfo paramListInfo = estate->es_param_list_info; + + if (paramId < 1 || paramListInfo == NULL || paramId > paramListInfo->numParams) + { + return NULL; + } + + ParamExternData *param = ¶mListInfo->params[paramId - 1]; + if (!OidIsValid(param->ptype) || param->isnull) + { + return NULL; + } + + Oid relationId = linitial_oid(plan->relationIdList); + CitusTableCacheEntry *tableEntry = GetCitusTableCacheEntry(relationId); + + return FindShardInterval(param->value, tableEntry); +} + + +/* + * BuildFastPathTask builds the minimal Task for a single-shard execution, + * bypassing plan copying, coordinator evaluation and task regeneration. + */ +static Task * +BuildFastPathTask(DistributedPlan *plan, Job *workerJob, EState *estate, + ShardInterval *shardInterval, bool isModify) +{ + List *shardIntervalListList = list_make1(list_make1(shardInterval)); + bool shardsPresent = false; + List *relationShardList = + RelationShardListForShardIntervalList(shardIntervalListList, &shardsPresent); + List *placementList = + CreateTaskPlacementListForShardIntervals(shardIntervalListList, shardsPresent, + true, false); + + Task *task = CitusMakeNode(Task); + task->taskType = isModify ? MODIFY_TASK : READ_TASK; + task->anchorShardId = shardInterval->shardId; + task->anchorDistributedTableId = linitial_oid(plan->relationIdList); + task->taskPlacementList = placementList; + task->queryCount = 1; + task->parametersInQueryStringResolved = true; + task->preparedStatementPlanId = plan->planId; + task->jobQueryForPrepare = workerJob->savedJobQueryForCaching; + task->relationShardList = relationShardList; + task->colocationId = workerJob->colocationId; + + ParamExternData *param = + &estate->es_param_list_info->params[workerJob->distributionKeyParamId - 1]; + int16 typeLength; + bool typeByValue; + get_typlenbyval(param->ptype, &typeLength, &typeByValue); + task->partitionKeyValue = makeConst(param->ptype, -1, InvalidOid, + (int) typeLength, param->value, false, + typeByValue); + + return task; +} + + +/* + * PreparedStatementCacheTryFastPath builds the task for this execution directly + * from the bound parameters, reusing the plan and query template saved on the + * first execution. + * + * Returns false when the fast path does not apply, in which case the caller + * must fall through to normal planning. + */ +bool +PreparedStatementCacheTryFastPath(struct CitusScanState *scanStateArg, EState *estate, + bool isModify) +{ + CitusScanState *scanState = (CitusScanState *) scanStateArg; + DistributedPlan *originalPlan = scanState->distributedPlan; + Job *workerJob = originalPlan->workerJob; + + /* the first execution populates the template the fast path depends on */ + if (!EnablePreparedStatementCaching || + originalPlan->numberOfTimesExecuted == 0 || + workerJob->savedJobQueryForCaching == NULL) + { + return false; + } + + /* volatile functions must be re-evaluated per execution on the coordinator */ + if (isModify && + (!workerJob->deferredPruning || workerJob->requiresCoordinatorEvaluation)) + { + return false; + } + + ShardInterval *shardInterval = FastPathShardInterval(originalPlan, workerJob, + estate); + if (shardInterval == NULL || (isModify && !ShardExists(shardInterval->shardId))) + { + return false; + } + + Task *task = BuildFastPathTask(originalPlan, workerJob, estate, shardInterval, + isModify); + + workerJob->taskList = list_make1(task); + workerJob->parametersInJobQueryResolved = true; + + elog(DEBUG2, "prepared statement cache-hit fast path%s: plan " UINT64_FORMAT + " shard " UINT64_FORMAT, + isModify ? " (DML)" : "", originalPlan->planId, shardInterval->shardId); + + /* the executor reads the plan back from the scan state */ + scanState->distributedPlan = originalPlan; + + if (isModify) + { + AcquireMetadataLocks(workerJob->taskList); + EnsureAnchorShardsInJobExist(workerJob); + workerJob->taskList = FirstReplicaAssignTaskList(workerJob->taskList); + } + + /* + * A fast-path task has no query string, so local execution needs a cached + * local plan rather than deparsed SQL. + */ + if (IsLocalPlanCachingSupported(workerJob, originalPlan)) + { + CacheLocalPlanForShardQuery(linitial(workerJob->taskList), originalPlan, + estate->es_param_list_info); + } + + return true; +} diff --git a/src/backend/distributed/metadata/XXnahQZC b/src/backend/distributed/metadata/XXnahQZC new file mode 100644 index 00000000000..e69de29bb2d diff --git a/src/backend/distributed/planner/deparse_shard_query.c b/src/backend/distributed/planner/deparse_shard_query.c index d38491baef2..18dffc90363 100644 --- a/src/backend/distributed/planner/deparse_shard_query.c +++ b/src/backend/distributed/planner/deparse_shard_query.c @@ -790,6 +790,35 @@ TaskQueryString(Task *task) int taskQueryType = GetTaskQueryType(task); if (taskQueryType == TASK_QUERY_NULL) { + if (task->jobQueryForPrepare != NULL) + { + /* + * Fast-path task from prepared statement caching: deparse the + * query from the saved template and cache the result on the task + * so subsequent calls don't re-deparse. + */ + Query *queryForDeparse = copyObject(task->jobQueryForPrepare); + StringInfoData buf; + initStringInfo(&buf); + + if (queryForDeparse->commandType == CMD_INSERT) + { + deparse_shard_query(queryForDeparse, + task->anchorDistributedTableId, + task->anchorShardId, &buf); + } + else + { + UpdateRelationToShardNames((Node *) queryForDeparse, + task->relationShardList); + pg_get_query_def(queryForDeparse, &buf); + } + + SetTaskQueryString(task, buf.data); + pfree(buf.data); + return task->taskQuery.data.queryStringLazy; + } + /* if task query type is TASK_QUERY_NULL then the data will be NULL, * this is unexpected state */ ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), diff --git a/src/backend/distributed/planner/distributed_planner.c b/src/backend/distributed/planner/distributed_planner.c index 949bfbd7519..a59eb810617 100644 --- a/src/backend/distributed/planner/distributed_planner.c +++ b/src/backend/distributed/planner/distributed_planner.c @@ -1448,11 +1448,22 @@ GetDistributedPlan(CustomScan *customScan) Node *node = (Node *) linitial(customScan->custom_private); Assert(CitusIsA(node, DistributedPlan)); - CheckNodeCopyAndSerialization(node); + /* + * Clear stale task pointers before CheckNodeCopyAndSerialization. + * When prepared statement caching fast-path is active, the previous + * execution may have set workerJob->taskList on the original plan. + * That task's memory has been freed (portal context destroyed), so + * serializing the plan now would dereference dangling pointers. + */ + DistributedPlan *plan = (DistributedPlan *) node; + if (plan->workerJob != NULL && plan->workerJob->deferredPruning) + { + plan->workerJob->taskList = NIL; + } - DistributedPlan *distributedPlan = (DistributedPlan *) node; + CheckNodeCopyAndSerialization(node); - return distributedPlan; + return plan; } @@ -2677,6 +2688,8 @@ CreateAndPushPlannerRestrictionContext(DistributedPlanningContext *planContext, fastPathRestrictionContext->distributionKeyValue; plannersFastPathCtx->distributionKeyHasParam = fastPathRestrictionContext->distributionKeyHasParam; + plannersFastPathCtx->distributionKeyParamId = + fastPathRestrictionContext->distributionKeyParamId; plannersFastPathCtx->delayFastPathPlanning = fastPathRestrictionContext->delayFastPathPlanning; } diff --git a/src/backend/distributed/planner/fast_path_router_planner.c b/src/backend/distributed/planner/fast_path_router_planner.c index 63c68f03af5..01493398092 100644 --- a/src/backend/distributed/planner/fast_path_router_planner.c +++ b/src/backend/distributed/planner/fast_path_router_planner.c @@ -180,7 +180,9 @@ InitializeFastPathContext(FastPathRestrictionContext *fastPathContext, } else if (IsA(distributionKeyValue, Param)) { + Param *distributionKeyParam = (Param *) distributionKeyValue; fastPathContext->distributionKeyHasParam = true; + fastPathContext->distributionKeyParamId = distributionKeyParam->paramid; } /* diff --git a/src/backend/distributed/planner/multi_router_planner.c b/src/backend/distributed/planner/multi_router_planner.c index 4bc038b3d4d..69b7b896c73 100644 --- a/src/backend/distributed/planner/multi_router_planner.c +++ b/src/backend/distributed/planner/multi_router_planner.c @@ -1768,6 +1768,33 @@ RouterInsertJob(Query *originalQuery) job->deferredPruning = true; job->partitionKeyValue = ExtractInsertPartitionKeyValue(originalQuery); + /* + * For single-row INSERTs with a parameterized distribution key, + * capture the Param index so the executor fast path can extract + * the value from ParamListInfo at execution time. + */ + if (!isMultiRowInsert && job->partitionKeyValue == NULL && + !job->requiresCoordinatorEvaluation) + { + Oid distributedTableId = ExtractFirstCitusTableId(originalQuery); + if (HasDistributionKey(distributedTableId)) + { + Var *partitionColumn = PartitionColumn(distributedTableId, 1); + TargetEntry *targetEntry = get_tle_by_resno(originalQuery->targetList, + partitionColumn->varattno); + if (targetEntry != NULL) + { + Node *targetExpression = strip_implicit_coercions( + (Node *) targetEntry->expr); + if (IsA(targetExpression, Param)) + { + job->distributionKeyParamId = + ((Param *) targetExpression)->paramid; + } + } + } + } + return job; } @@ -1955,6 +1982,7 @@ RouterJob(Query *originalQuery, PlannerRestrictionContext *plannerRestrictionCon { Job *job = CreateJob(originalQuery); job->deferredPruning = true; + job->distributionKeyParamId = fastPathRestrictionContext->distributionKeyParamId; ereport(DEBUG2, (errmsg("Deferred pruning for a fast-path router " "query"))); diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c index 603f20f9384..bb7051d5b44 100644 --- a/src/backend/distributed/shared_library_init.c +++ b/src/backend/distributed/shared_library_init.c @@ -91,6 +91,7 @@ #include "distributed/multi_server_executor.h" #include "distributed/pg_dist_partition.h" #include "distributed/placement_connection.h" +#include "distributed/prepared_statement_cache.h" #include "distributed/priority.h" #include "distributed/procedure_body_analysis.h" #include "distributed/query_pushdown_planning.h" @@ -1602,6 +1603,18 @@ RegisterCitusConfigVariables(void) GUC_NO_SHOW_ALL | GUC_NOT_IN_SAMPLE, WarnIfDeprecatedPseudoconstantQualsGucIsSet, NULL, NULL); + DefineCustomBoolVariable( + "citus.enable_prepared_statement_caching", + gettext_noop("Enables caching prepared statement plans on worker connections."), + gettext_noop("When enabled, Citus uses PQprepare/PQsendQueryPrepared on " + "worker connections for fast-path prepared statement executions, " + "eliminating redundant parse/plan cycles on workers."), + &EnablePreparedStatementCaching, + false, + PGC_USERSET, + GUC_STANDARD, + NULL, NULL, NULL); + DefineCustomBoolVariable( "citus.enable_procedure_transaction_skip", gettext_noop("Skip coordinated transactions for single-statement, " diff --git a/src/backend/distributed/utils/citus_clauses.c b/src/backend/distributed/utils/citus_clauses.c index 7086cec4eab..386536b8831 100644 --- a/src/backend/distributed/utils/citus_clauses.c +++ b/src/backend/distributed/utils/citus_clauses.c @@ -84,6 +84,25 @@ ExecuteCoordinatorEvaluableExpressions(Query *query, PlanState *planState) } +/* + * ExecuteCoordinatorEvaluableFunctions evaluates function calls on the + * coordinator (e.g. now(), nextval()) but leaves Param nodes untouched. + * This is used for prepared statement caching where we need to evaluate + * volatile/stable functions but preserve parameter placeholders for the + * prepared statement template. + */ +void +ExecuteCoordinatorEvaluableFunctions(Query *query, PlanState *planState) +{ + CoordinatorEvaluationContext coordinatorEvaluationContext; + + coordinatorEvaluationContext.planState = planState; + coordinatorEvaluationContext.evaluationMode = EVALUATE_FUNCTIONS; + + PartiallyEvaluateExpression((Node *) query, &coordinatorEvaluationContext); +} + + /* * PartiallyEvaluateExpression descends into an expression tree to evaluate * expressions that can be resolved to a constant on the master. Expressions @@ -112,6 +131,13 @@ PartiallyEvaluateExpression(Node *expression, { Assert(((Param *) expression)->paramkind != PARAM_MULTIEXPR && ((Param *) expression)->paramkind != PARAM_SUBLINK); + + /* in EVALUATE_FUNCTIONS mode, leave Param nodes untouched */ + if (coordinatorEvaluationContext->evaluationMode == EVALUATE_FUNCTIONS) + { + return expression; + } + return (Node *) citus_evaluate_expr((Expr *) expression, exprType(expression), exprTypmod(expression), @@ -224,7 +250,8 @@ ShouldEvaluateFunctions(CoordinatorEvaluationContext *evaluationContext) return true; } - return evaluationContext->evaluationMode == EVALUATE_FUNCTIONS_PARAMS; + return evaluationContext->evaluationMode == EVALUATE_FUNCTIONS_PARAMS || + evaluationContext->evaluationMode == EVALUATE_FUNCTIONS; } diff --git a/src/backend/distributed/utils/citus_copyfuncs.c b/src/backend/distributed/utils/citus_copyfuncs.c index 58ef1697564..06ecc9bb05a 100644 --- a/src/backend/distributed/utils/citus_copyfuncs.c +++ b/src/backend/distributed/utils/citus_copyfuncs.c @@ -99,9 +99,12 @@ copyJobInfo(Job *newnode, Job *from) COPY_SCALAR_FIELD(subqueryPushdown); COPY_SCALAR_FIELD(requiresCoordinatorEvaluation); COPY_SCALAR_FIELD(deferredPruning); + COPY_SCALAR_FIELD(distributionKeyParamId); COPY_NODE_FIELD(partitionKeyValue); COPY_NODE_FIELD(localPlannedStatements); COPY_SCALAR_FIELD(parametersInJobQueryResolved); + COPY_SCALAR_FIELD(colocationId); + COPY_NODE_FIELD(savedJobQueryForCaching); } @@ -359,6 +362,8 @@ CopyNodeTask(COPYFUNC_ARGS) COPY_NODE_FIELD(rowValuesLists); COPY_SCALAR_FIELD(partiallyLocalOrRemote); COPY_SCALAR_FIELD(parametersInQueryStringResolved); + COPY_SCALAR_FIELD(preparedStatementPlanId); + COPY_NODE_FIELD(jobQueryForPrepare); COPY_SCALAR_FIELD(tupleDest); COPY_SCALAR_FIELD(queryCount); COPY_SCALAR_FIELD(totalReceivedTupleData); @@ -367,6 +372,8 @@ CopyNodeTask(COPYFUNC_ARGS) COPY_SCALAR_FIELD(fetchedExplainAnalyzeExecutionDuration); COPY_SCALAR_FIELD(isLocalTableModification); COPY_SCALAR_FIELD(cannotBeExecutedInTransaction); + COPY_NODE_FIELD(partitionKeyValue); + COPY_SCALAR_FIELD(colocationId); } diff --git a/src/backend/distributed/utils/citus_outfuncs.c b/src/backend/distributed/utils/citus_outfuncs.c index 3ddabbbca2e..d6b6051e4cd 100644 --- a/src/backend/distributed/utils/citus_outfuncs.c +++ b/src/backend/distributed/utils/citus_outfuncs.c @@ -393,9 +393,11 @@ OutJobFields(StringInfo str, const Job *node) WRITE_BOOL_FIELD(subqueryPushdown); WRITE_BOOL_FIELD(requiresCoordinatorEvaluation); WRITE_BOOL_FIELD(deferredPruning); + WRITE_INT_FIELD(distributionKeyParamId); WRITE_NODE_FIELD(partitionKeyValue); WRITE_NODE_FIELD(localPlannedStatements); WRITE_BOOL_FIELD(parametersInJobQueryResolved); + WRITE_UINT_FIELD(colocationId); } @@ -576,6 +578,7 @@ OutTask(OUTFUNC_ARGS) WRITE_NODE_FIELD(rowValuesLists); WRITE_BOOL_FIELD(partiallyLocalOrRemote); WRITE_BOOL_FIELD(parametersInQueryStringResolved); + WRITE_UINT64_FIELD(preparedStatementPlanId); WRITE_INT_FIELD(queryCount); WRITE_UINT64_FIELD(totalReceivedTupleData); WRITE_INT_FIELD(fetchedExplainAnalyzePlacementIndex); @@ -583,6 +586,7 @@ OutTask(OUTFUNC_ARGS) WRITE_FLOAT_FIELD(fetchedExplainAnalyzeExecutionDuration, "%.2f"); WRITE_BOOL_FIELD(isLocalTableModification); WRITE_BOOL_FIELD(cannotBeExecutedInTransaction); + WRITE_INT_FIELD(colocationId); } diff --git a/src/include/distributed/citus_clauses.h b/src/include/distributed/citus_clauses.h index 8cfbe18eeb3..b30cc0e03e7 100644 --- a/src/include/distributed/citus_clauses.h +++ b/src/include/distributed/citus_clauses.h @@ -28,6 +28,9 @@ typedef enum CoordinatorEvaluationMode /* evaluate only external parameters */ EVALUATE_PARAMS, + /* evaluate functions/expressions but leave Param nodes untouched */ + EVALUATE_FUNCTIONS, + /* evaluate both the functions/expressions and the external paramaters */ EVALUATE_FUNCTIONS_PARAMS } CoordinatorEvaluationMode; @@ -48,6 +51,7 @@ extern void ExecuteCoordinatorEvaluableExpressions(Query *query, PlanState *plan extern Node * PartiallyEvaluateExpression(Node *expression, CoordinatorEvaluationContext * coordinatorEvaluationContext); +extern void ExecuteCoordinatorEvaluableFunctions(Query *query, PlanState *planState); extern bool CitusIsVolatileFunction(Node *node); extern bool CitusIsMutableFunction(Node *node); diff --git a/src/include/distributed/citus_custom_scan.h b/src/include/distributed/citus_custom_scan.h index 1c74da08756..5d73733ea3f 100644 --- a/src/include/distributed/citus_custom_scan.h +++ b/src/include/distributed/citus_custom_scan.h @@ -59,5 +59,6 @@ extern bool IsCitusPlan(Plan *plan); extern bool IsCitusCustomScan(Plan *plan); extern void SetJobColocationId(Job *job); +extern void EnsureAnchorShardsInJobExist(Job *job); #endif /* CITUS_CUSTOM_SCAN_H */ diff --git a/src/include/distributed/connection_management.h b/src/include/distributed/connection_management.h index 8032a0ea37b..aa8178f9d26 100644 --- a/src/include/distributed/connection_management.h +++ b/src/include/distributed/connection_management.h @@ -232,6 +232,9 @@ typedef struct MultiConnection bool requiresReplication; MultiConnectionStructInitializationState initializationState; + + /* per-connection cache of prepared statements on this worker connection */ + HTAB *preparedStatementCache; } MultiConnection; diff --git a/src/include/distributed/distributed_planner.h b/src/include/distributed/distributed_planner.h index f07ac78c2e6..92d915b5c29 100644 --- a/src/include/distributed/distributed_planner.h +++ b/src/include/distributed/distributed_planner.h @@ -112,6 +112,14 @@ typedef struct FastPathRestrictionContext */ bool distributionKeyHasParam; + /* + * When distributionKeyHasParam is true, stores the Param's paramid + * so the executor can extract the distribution key value directly + * from ParamListInfo without walking the query tree. 0 when not set + * (paramid is 1-based in PostgreSQL). + */ + int distributionKeyParamId; + /* * Indicates to hold off calling the fast path planner until its * known if the shard is local or not. diff --git a/src/include/distributed/multi_physical_planner.h b/src/include/distributed/multi_physical_planner.h index 977b51825cd..cd55fe270a0 100644 --- a/src/include/distributed/multi_physical_planner.h +++ b/src/include/distributed/multi_physical_planner.h @@ -141,6 +141,7 @@ typedef struct Job bool subqueryPushdown; bool requiresCoordinatorEvaluation; /* only applies to modify jobs */ bool deferredPruning; + int distributionKeyParamId; Const *partitionKeyValue; /* for local shard queries, we may save the local plan here */ @@ -153,6 +154,15 @@ typedef struct Job */ bool parametersInJobQueryResolved; uint32 colocationId; /* common colocation group ID of the relations */ + + /* + * Cached copy of jobQuery with Param nodes ($1, $2, ...) still intact, + * used for prepared statement caching. Populated lazily on first + * execution and reused across subsequent executions to avoid a + * per-execution copyObject(jobQuery). Lives on originalDistributedPlan + * only; the per-execution copy (currentPlan) gets a NULL here. + */ + Query *savedJobQueryForCaching; } Job; @@ -311,6 +321,12 @@ typedef struct Task */ bool parametersInQueryStringResolved; + /* prepared statement cache key: plan ID from DistributedPlan */ + uint64 preparedStatementPlanId; + + /* pre-evaluation job query with Param nodes intact, for PQprepare on cache miss */ + Query *jobQueryForPrepare; + /* * Destination of tuples generated as a result of executing this task. Can be * NULL, in which case executor might use a default destination. diff --git a/src/include/distributed/prepared_statement_cache.h b/src/include/distributed/prepared_statement_cache.h new file mode 100644 index 00000000000..ba7373b5aa3 --- /dev/null +++ b/src/include/distributed/prepared_statement_cache.h @@ -0,0 +1,102 @@ +/*------------------------------------------------------------------------- + * prepared_statement_cache.h + * + * Declarations for per-connection prepared statement caching on worker + * connections. + * + * Copyright (c) Citus Data, Inc. + * + *------------------------------------------------------------------------- + */ + +#ifndef PREPARED_STATEMENT_CACHE_H +#define PREPARED_STATEMENT_CACHE_H + +#include "postgres.h" + +#include "nodes/execnodes.h" +#include "nodes/parsenodes.h" +#include "utils/hsearch.h" + +#include "distributed/connection_management.h" +#include "distributed/multi_physical_planner.h" + +struct CitusScanState; + +/* compile-time limit for per-connection cached prepared statements */ +#define MAX_CACHED_STMTS_PER_CONNECTION 1000 + +/* maximum length for generated statement names ("__citus_stmt_NNNN") */ +#define MAX_STMT_NAME_LENGTH 64 + + +/* + * PreparedStatementCacheKey uniquely identifies a prepared statement on a + * given worker connection: planId identifies the cached generic plan, + * shardId identifies the target shard. + */ +typedef struct PreparedStatementCacheKey +{ + uint64 planId; + uint64 shardId; +} PreparedStatementCacheKey; + + +/* + * PreparedStatementCacheEntry stores the prepared statement handle on a + * connection, plus metadata needed to re-execute it. + */ +typedef struct PreparedStatementCacheEntry +{ + PreparedStatementCacheKey key; + + char stmtName[MAX_STMT_NAME_LENGTH]; + Oid *paramTypes; + int paramCount; + char *parameterizedQueryString; +} PreparedStatementCacheEntry; + + +/* + * Outcome of an attempt to dispatch a task through the connection's + * prepared statement cache. + */ +typedef enum PreparedStatementSendStatus +{ + PREPARED_STMT_NOT_APPLICABLE, /* caching off or task carries no template */ + PREPARED_STMT_SENT, /* dispatched on the connection */ + PREPARED_STMT_FAILED, /* connection lost */ + PREPARED_STMT_FALLBACK /* cache full; caller sends the returned SQL */ +} PreparedStatementSendStatus; + + +/* GUC variable */ +extern bool EnablePreparedStatementCaching; + +/* cache lifecycle */ +extern HTAB * PreparedStatementCacheCreate(void); +extern void PreparedStatementCacheDestroy(HTAB **cache_ptr); + +/* cache operations */ +extern PreparedStatementCacheEntry * PreparedStatementCacheLookup(HTAB *cache, uint64 + planId, uint64 shardId); +extern PreparedStatementCacheEntry * PreparedStatementCacheInsert(HTAB *cache, uint64 + planId, uint64 shardId); + +/* planner-side integration (see citus_custom_scan.c) */ +extern bool PreparedStatementCacheTryFastPath(struct CitusScanState *scanState, + EState *estate, bool isModify); +extern Query * PreparedStatementCacheSaveTemplate(DistributedPlan *originalPlan); +extern void PreparedStatementCacheAttachToTasks(DistributedPlan *currentPlan, + Job *workerJob, Query *savedJobQuery); + +/* executor-side integration (see adaptive_executor.c) */ +extern PreparedStatementSendStatus PreparedStatementCacheSendQuery(MultiConnection * + connection, Task *task, + ParamListInfo + paramListInfo, + bool binaryResults, + char ** + fallbackQueryString); + +#endif /* PREPARED_STATEMENT_CACHE_H */ diff --git a/src/include/distributed/remote_commands.h b/src/include/distributed/remote_commands.h index 71cb9dad27f..88662503767 100644 --- a/src/include/distributed/remote_commands.h +++ b/src/include/distributed/remote_commands.h @@ -68,6 +68,11 @@ extern bool PutRemoteCopyEnd(MultiConnection *connection, const char *errormsg); /* waiting for multiple command results */ extern void WaitForAllConnections(List *connectionList, bool raiseInterrupts); +extern int SendRemotePrepare(MultiConnection *connection, const char *stmtName, + const char *query, int nParams, const Oid *paramTypes); +extern int SendRemotePreparedQuery(MultiConnection *connection, const char *stmtName, + int nParams, const char *const *paramValues, + bool binaryResults); extern bool SendCancelationRequest(MultiConnection *connection); extern bool EvaluateSingleQueryResult(MultiConnection *connection, PGresult *queryResult, diff --git a/src/test/regress/citus_tests/run_test.py b/src/test/regress/citus_tests/run_test.py index 77c13287557..bf64c16eb9b 100755 --- a/src/test/regress/citus_tests/run_test.py +++ b/src/test/regress/citus_tests/run_test.py @@ -337,6 +337,9 @@ def extra_tests(self): "multi_subquery_window_functions": TestDeps( "minimal_schedule", ["multi_behavioral_analytics_create_table"] ), + "subquery_prepared_statements": TestDeps( + "minimal_schedule", ["multi_behavioral_analytics_create_table"] + ), } diff --git a/src/test/regress/expected/prepared_statement_caching.out b/src/test/regress/expected/prepared_statement_caching.out new file mode 100644 index 00000000000..4261cc9155c --- /dev/null +++ b/src/test/regress/expected/prepared_statement_caching.out @@ -0,0 +1,606 @@ +-- +-- PREPARED_STATEMENT_CACHING +-- +-- Tests for citus.enable_prepared_statement_caching, which enables +-- worker-side prepared statement plan caching for fast-path queries. +-- +CREATE SCHEMA prepared_stmt_caching; +SET search_path TO prepared_stmt_caching; +-- Create test tables +CREATE TABLE dist_table ( + key int PRIMARY KEY, + value int, + label text +); +SELECT create_distributed_table('dist_table', 'key'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +-- Insert base data +INSERT INTO dist_table SELECT i, i * 10, 'label-' || i FROM generate_series(1, 20) i; +CREATE TABLE dist_table_ts ( + key int PRIMARY KEY, + value int, + created_at timestamptz DEFAULT now() +); +SELECT create_distributed_table('dist_table_ts', 'key'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +-- ============================================================ +-- Test 1: GUC toggle — verify default is OFF, SET to ON succeeds +-- ============================================================ +SHOW citus.enable_prepared_statement_caching; + citus.enable_prepared_statement_caching +--------------------------------------------------------------------- + off +(1 row) + +SET citus.enable_prepared_statement_caching = on; +SHOW citus.enable_prepared_statement_caching; + citus.enable_prepared_statement_caching +--------------------------------------------------------------------- + on +(1 row) + +SET citus.enable_prepared_statement_caching = off; +SHOW citus.enable_prepared_statement_caching; + citus.enable_prepared_statement_caching +--------------------------------------------------------------------- + off +(1 row) + +-- ============================================================ +-- Test 2: Basic caching — PREPARE a single-shard SELECT, EXECUTE +-- 10 times with GUC ON, verify correct results +-- ============================================================ +SET citus.enable_prepared_statement_caching = on; +set citus.max_cached_connection_lifetime to '60min'; +SET search_path TO prepared_stmt_caching; +PREPARE cached_select(int) AS + SELECT key, value FROM dist_table WHERE key = $1; +-- Execute 10 times to ensure generic plan path and cache hit path +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +DEALLOCATE cached_select; +-- ============================================================ +-- Test 3: Multi-shard-value — EXECUTE with different partition +-- key values routing to different shards +-- ============================================================ +PREPARE cached_multi_shard(int) AS + SELECT key, value FROM dist_table WHERE key = $1; +-- Different keys likely route to different shards +EXECUTE cached_multi_shard(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_multi_shard(2); + key | value +--------------------------------------------------------------------- + 2 | 20 +(1 row) + +EXECUTE cached_multi_shard(3); + key | value +--------------------------------------------------------------------- + 3 | 30 +(1 row) + +EXECUTE cached_multi_shard(4); + key | value +--------------------------------------------------------------------- + 4 | 40 +(1 row) + +EXECUTE cached_multi_shard(5); + key | value +--------------------------------------------------------------------- + 5 | 50 +(1 row) + +EXECUTE cached_multi_shard(6); + key | value +--------------------------------------------------------------------- + 6 | 60 +(1 row) + +EXECUTE cached_multi_shard(7); + key | value +--------------------------------------------------------------------- + 7 | 70 +(1 row) + +EXECUTE cached_multi_shard(8); + key | value +--------------------------------------------------------------------- + 8 | 80 +(1 row) + +EXECUTE cached_multi_shard(9); + key | value +--------------------------------------------------------------------- + 9 | 90 +(1 row) + +EXECUTE cached_multi_shard(10); + key | value +--------------------------------------------------------------------- + 10 | 100 +(1 row) + +DEALLOCATE cached_multi_shard; +-- ============================================================ +-- Test 4: INSERT/UPDATE/DELETE with caching ON +-- Include now() to verify coordinator-side function +-- evaluation still works. +-- Single-row INSERT, UPDATE, and DELETE all use +-- the cached prepared statement path. +-- ============================================================ +-- INSERT (cached via deparse_shard_query path) +PREPARE cached_insert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2); +EXECUTE cached_insert(100, 1000); +EXECUTE cached_insert(101, 1010); +EXECUTE cached_insert(102, 1020); +EXECUTE cached_insert(103, 1030); +EXECUTE cached_insert(104, 1040); +EXECUTE cached_insert(105, 1050); +EXECUTE cached_insert(106, 1060); +EXECUTE cached_insert(107, 1070); +EXECUTE cached_insert(108, 1080); +EXECUTE cached_insert(109, 1090); +EXECUTE cached_insert(110, 1100); +-- Verify inserts +SELECT key, value FROM dist_table_ts WHERE key >= 100 ORDER BY key; + key | value +--------------------------------------------------------------------- + 100 | 1000 + 101 | 1010 + 102 | 1020 + 103 | 1030 + 104 | 1040 + 105 | 1050 + 106 | 1060 + 107 | 1070 + 108 | 1080 + 109 | 1090 + 110 | 1100 +(11 rows) + +DEALLOCATE cached_insert; +-- UPDATE +PREPARE cached_update(int, int) AS + UPDATE dist_table SET value = $2 WHERE key = $1; +EXECUTE cached_update(1, 100); +EXECUTE cached_update(2, 200); +EXECUTE cached_update(3, 300); +EXECUTE cached_update(4, 400); +EXECUTE cached_update(5, 500); +EXECUTE cached_update(6, 600); +EXECUTE cached_update(7, 700); +SELECT key, value FROM dist_table WHERE key <= 7 ORDER BY key; + key | value +--------------------------------------------------------------------- + 1 | 100 + 2 | 200 + 3 | 300 + 4 | 400 + 5 | 500 + 6 | 600 + 7 | 700 +(7 rows) + +DEALLOCATE cached_update; +-- DELETE +PREPARE cached_delete(int) AS + DELETE FROM dist_table WHERE key = $1; +EXECUTE cached_delete(18); +EXECUTE cached_delete(19); +EXECUTE cached_delete(20); +EXECUTE cached_delete(18); +EXECUTE cached_delete(19); +EXECUTE cached_delete(20); +EXECUTE cached_delete(18); +-- Verify deletes +SELECT count(*) FROM dist_table WHERE key >= 18; + count +--------------------------------------------------------------------- + 0 +(1 row) + +DEALLOCATE cached_delete; +-- INSERT with now() function evaluation +PREPARE cached_insert_ts(int) AS + INSERT INTO dist_table_ts (key, value, created_at) VALUES ($1, $1 * 10, now()); +EXECUTE cached_insert_ts(200); +EXECUTE cached_insert_ts(201); +EXECUTE cached_insert_ts(202); +EXECUTE cached_insert_ts(203); +EXECUTE cached_insert_ts(204); +EXECUTE cached_insert_ts(205); +EXECUTE cached_insert_ts(206); +-- Verify that each row has created_at populated (functions were evaluated) +SELECT key, value, created_at IS NOT NULL AS has_ts FROM dist_table_ts + WHERE key >= 200 ORDER BY key; + key | value | has_ts +--------------------------------------------------------------------- + 200 | 2000 | t + 201 | 2010 | t + 202 | 2020 | t + 203 | 2030 | t + 204 | 2040 | t + 205 | 2050 | t + 206 | 2060 | t +(7 rows) + +DEALLOCATE cached_insert_ts; +-- INSERT ... ON CONFLICT DO UPDATE (upsert). The qualified reference to the +-- target table in DO UPDATE must resolve to the shard alias on the worker. +PREPARE cached_upsert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2) + ON CONFLICT (key) DO UPDATE SET value = dist_table_ts.value + EXCLUDED.value; +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(501, 5); +EXECUTE cached_upsert(501, 5); +-- key 500 accumulated 6 increments, key 501 accumulated 2 +SELECT key, value FROM dist_table_ts WHERE key >= 500 AND key < 600 ORDER BY key; + key | value +--------------------------------------------------------------------- + 500 | 6 + 501 | 10 +(2 rows) + +DEALLOCATE cached_upsert; +-- Multi-row INSERT is not cacheable: each shard's task carries only its own +-- subset of VALUES rows, so the rows must not be duplicated across shards. +PREPARE cached_multirow(int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, 1), ($1 + 1, 2); +EXECUTE cached_multirow(600); +EXECUTE cached_multirow(602); +EXECUTE cached_multirow(604); +EXECUTE cached_multirow(606); +EXECUTE cached_multirow(608); +EXECUTE cached_multirow(610); +-- exactly 12 rows, one per key, no duplicates +SELECT count(*) AS row_count, count(DISTINCT key) AS distinct_keys + FROM dist_table_ts WHERE key >= 600; + row_count | distinct_keys +--------------------------------------------------------------------- + 12 | 12 +(1 row) + +DEALLOCATE cached_multirow; +-- ============================================================ +-- Test 5: GUC OFF baseline — same queries produce identical results +-- ============================================================ +SET citus.enable_prepared_statement_caching = off; +PREPARE uncached_select(int) AS + SELECT key, value FROM dist_table WHERE key = $1; +EXECUTE uncached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE uncached_select(2); + key | value +--------------------------------------------------------------------- + 2 | 200 +(1 row) + +EXECUTE uncached_select(3); + key | value +--------------------------------------------------------------------- + 3 | 300 +(1 row) + +EXECUTE uncached_select(4); + key | value +--------------------------------------------------------------------- + 4 | 400 +(1 row) + +EXECUTE uncached_select(5); + key | value +--------------------------------------------------------------------- + 5 | 500 +(1 row) + +EXECUTE uncached_select(6); + key | value +--------------------------------------------------------------------- + 6 | 600 +(1 row) + +EXECUTE uncached_select(7); + key | value +--------------------------------------------------------------------- + 7 | 700 +(1 row) + +EXECUTE uncached_select(8); + key | value +--------------------------------------------------------------------- + 8 | 80 +(1 row) + +EXECUTE uncached_select(9); + key | value +--------------------------------------------------------------------- + 9 | 90 +(1 row) + +EXECUTE uncached_select(10); + key | value +--------------------------------------------------------------------- + 10 | 100 +(1 row) + +DEALLOCATE uncached_select; +-- ============================================================ +-- Test 6: Multiple prepared statements in same session +-- ============================================================ +SET citus.enable_prepared_statement_caching = on; +PREPARE stmt_a(int) AS SELECT key, value FROM dist_table WHERE key = $1; +PREPARE stmt_b(int) AS SELECT key, label FROM dist_table WHERE key = $1; +PREPARE stmt_c(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2); +-- Interleave executions to verify independent caching +EXECUTE stmt_a(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE stmt_b(1); + key | label +--------------------------------------------------------------------- + 1 | label-1 +(1 row) + +EXECUTE stmt_a(2); + key | value +--------------------------------------------------------------------- + 2 | 200 +(1 row) + +EXECUTE stmt_b(2); + key | label +--------------------------------------------------------------------- + 2 | label-2 +(1 row) + +EXECUTE stmt_a(3); + key | value +--------------------------------------------------------------------- + 3 | 300 +(1 row) + +EXECUTE stmt_b(3); + key | label +--------------------------------------------------------------------- + 3 | label-3 +(1 row) + +EXECUTE stmt_a(4); + key | value +--------------------------------------------------------------------- + 4 | 400 +(1 row) + +EXECUTE stmt_b(4); + key | label +--------------------------------------------------------------------- + 4 | label-4 +(1 row) + +EXECUTE stmt_a(5); + key | value +--------------------------------------------------------------------- + 5 | 500 +(1 row) + +EXECUTE stmt_b(5); + key | label +--------------------------------------------------------------------- + 5 | label-5 +(1 row) + +EXECUTE stmt_a(6); + key | value +--------------------------------------------------------------------- + 6 | 600 +(1 row) + +EXECUTE stmt_b(6); + key | label +--------------------------------------------------------------------- + 6 | label-6 +(1 row) + +EXECUTE stmt_a(7); + key | value +--------------------------------------------------------------------- + 7 | 700 +(1 row) + +EXECUTE stmt_b(7); + key | label +--------------------------------------------------------------------- + 7 | label-7 +(1 row) + +EXECUTE stmt_c(300, 3000); +EXECUTE stmt_c(301, 3010); +EXECUTE stmt_c(302, 3020); +EXECUTE stmt_c(303, 3030); +EXECUTE stmt_c(304, 3040); +EXECUTE stmt_c(305, 3050); +EXECUTE stmt_c(306, 3060); +-- Verify inserts from stmt_c +SELECT key, value FROM dist_table_ts WHERE key >= 300 AND key < 400 ORDER BY key; + key | value +--------------------------------------------------------------------- + 300 | 3000 + 301 | 3010 + 302 | 3020 + 303 | 3030 + 304 | 3040 + 305 | 3050 + 306 | 3060 +(7 rows) + +DEALLOCATE stmt_a; +DEALLOCATE stmt_b; +DEALLOCATE stmt_c; +-- ============================================================ +-- Test 7: Connection loss re-prepare — force worker connection +-- close, verify subsequent EXECUTE still works +-- ============================================================ +SET citus.enable_prepared_statement_caching = on; +PREPARE reconnect_test(int) AS + SELECT key, value FROM dist_table WHERE key = $1; +-- Execute enough times to get into generic plan + cache hit +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +-- Force all cached connections to be dropped by setting lifetime to 0 +SET citus.max_cached_connection_lifetime TO '0s'; +-- The next execution should get a new connection, re-prepare, and succeed +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE reconnect_test(2); + key | value +--------------------------------------------------------------------- + 2 | 200 +(1 row) + +EXECUTE reconnect_test(3); + key | value +--------------------------------------------------------------------- + 3 | 300 +(1 row) + +-- Restore default +RESET citus.max_cached_connection_lifetime; +DEALLOCATE reconnect_test; +-- ============================================================ +-- Cleanup +-- ============================================================ +SET citus.enable_prepared_statement_caching = off; +DROP SCHEMA prepared_stmt_caching CASCADE; +NOTICE: drop cascades to 2 other objects +DETAIL: drop cascades to table dist_table +drop cascades to table dist_table_ts diff --git a/src/test/regress/multi_schedule b/src/test/regress/multi_schedule index 0960831df83..24fbaa169db 100644 --- a/src/test/regress/multi_schedule +++ b/src/test/regress/multi_schedule @@ -61,6 +61,7 @@ test: set_operation_and_local_tables test: subqueries_deep subquery_view subquery_partitioning subqueries_not_supported test: subquery_in_targetlist subquery_in_where subquery_complex_target_list subquery_append test: subquery_prepared_statements +test: prepared_statement_caching test: non_colocated_leaf_subquery_joins non_colocated_subquery_joins test: cte_inline recursive_view_local_table values sequences_with_different_types multi_level_recursive_queries test: pg13 pg12 diff --git a/src/test/regress/sql/prepared_statement_caching.sql b/src/test/regress/sql/prepared_statement_caching.sql new file mode 100644 index 00000000000..8880b165508 --- /dev/null +++ b/src/test/regress/sql/prepared_statement_caching.sql @@ -0,0 +1,306 @@ +-- +-- PREPARED_STATEMENT_CACHING +-- +-- Tests for citus.enable_prepared_statement_caching, which enables +-- worker-side prepared statement plan caching for fast-path queries. +-- + +CREATE SCHEMA prepared_stmt_caching; +SET search_path TO prepared_stmt_caching; + +-- Create test tables +CREATE TABLE dist_table ( + key int PRIMARY KEY, + value int, + label text +); +SELECT create_distributed_table('dist_table', 'key'); + +-- Insert base data +INSERT INTO dist_table SELECT i, i * 10, 'label-' || i FROM generate_series(1, 20) i; + +CREATE TABLE dist_table_ts ( + key int PRIMARY KEY, + value int, + created_at timestamptz DEFAULT now() +); +SELECT create_distributed_table('dist_table_ts', 'key'); + +-- ============================================================ +-- Test 1: GUC toggle — verify default is OFF, SET to ON succeeds +-- ============================================================ + +SHOW citus.enable_prepared_statement_caching; +SET citus.enable_prepared_statement_caching = on; +SHOW citus.enable_prepared_statement_caching; +SET citus.enable_prepared_statement_caching = off; +SHOW citus.enable_prepared_statement_caching; + +-- ============================================================ +-- Test 2: Basic caching — PREPARE a single-shard SELECT, EXECUTE +-- 10 times with GUC ON, verify correct results +-- ============================================================ + +SET citus.enable_prepared_statement_caching = on; +set citus.max_cached_connection_lifetime to '60min'; +SET search_path TO prepared_stmt_caching; + +PREPARE cached_select(int) AS + SELECT key, value FROM dist_table WHERE key = $1; + +-- Execute 10 times to ensure generic plan path and cache hit path +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); + +DEALLOCATE cached_select; + +-- ============================================================ +-- Test 3: Multi-shard-value — EXECUTE with different partition +-- key values routing to different shards +-- ============================================================ + +PREPARE cached_multi_shard(int) AS + SELECT key, value FROM dist_table WHERE key = $1; + +-- Different keys likely route to different shards +EXECUTE cached_multi_shard(1); +EXECUTE cached_multi_shard(2); +EXECUTE cached_multi_shard(3); +EXECUTE cached_multi_shard(4); +EXECUTE cached_multi_shard(5); +EXECUTE cached_multi_shard(6); +EXECUTE cached_multi_shard(7); +EXECUTE cached_multi_shard(8); +EXECUTE cached_multi_shard(9); +EXECUTE cached_multi_shard(10); + +DEALLOCATE cached_multi_shard; + +-- ============================================================ +-- Test 4: INSERT/UPDATE/DELETE with caching ON +-- Include now() to verify coordinator-side function +-- evaluation still works. +-- Single-row INSERT, UPDATE, and DELETE all use +-- the cached prepared statement path. +-- ============================================================ + +-- INSERT (cached via deparse_shard_query path) +PREPARE cached_insert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2); + +EXECUTE cached_insert(100, 1000); +EXECUTE cached_insert(101, 1010); +EXECUTE cached_insert(102, 1020); +EXECUTE cached_insert(103, 1030); +EXECUTE cached_insert(104, 1040); +EXECUTE cached_insert(105, 1050); +EXECUTE cached_insert(106, 1060); + +EXECUTE cached_insert(107, 1070); +EXECUTE cached_insert(108, 1080); +EXECUTE cached_insert(109, 1090); +EXECUTE cached_insert(110, 1100); + +-- Verify inserts +SELECT key, value FROM dist_table_ts WHERE key >= 100 ORDER BY key; + +DEALLOCATE cached_insert; + +-- UPDATE +PREPARE cached_update(int, int) AS + UPDATE dist_table SET value = $2 WHERE key = $1; + +EXECUTE cached_update(1, 100); +EXECUTE cached_update(2, 200); +EXECUTE cached_update(3, 300); +EXECUTE cached_update(4, 400); +EXECUTE cached_update(5, 500); +EXECUTE cached_update(6, 600); +EXECUTE cached_update(7, 700); + +SELECT key, value FROM dist_table WHERE key <= 7 ORDER BY key; + +DEALLOCATE cached_update; + +-- DELETE +PREPARE cached_delete(int) AS + DELETE FROM dist_table WHERE key = $1; + +EXECUTE cached_delete(18); +EXECUTE cached_delete(19); +EXECUTE cached_delete(20); +EXECUTE cached_delete(18); +EXECUTE cached_delete(19); +EXECUTE cached_delete(20); +EXECUTE cached_delete(18); + +-- Verify deletes +SELECT count(*) FROM dist_table WHERE key >= 18; + +DEALLOCATE cached_delete; + +-- INSERT with now() function evaluation +PREPARE cached_insert_ts(int) AS + INSERT INTO dist_table_ts (key, value, created_at) VALUES ($1, $1 * 10, now()); + +EXECUTE cached_insert_ts(200); +EXECUTE cached_insert_ts(201); +EXECUTE cached_insert_ts(202); +EXECUTE cached_insert_ts(203); +EXECUTE cached_insert_ts(204); +EXECUTE cached_insert_ts(205); +EXECUTE cached_insert_ts(206); + +-- Verify that each row has created_at populated (functions were evaluated) +SELECT key, value, created_at IS NOT NULL AS has_ts FROM dist_table_ts + WHERE key >= 200 ORDER BY key; + +DEALLOCATE cached_insert_ts; + +-- INSERT ... ON CONFLICT DO UPDATE (upsert). The qualified reference to the +-- target table in DO UPDATE must resolve to the shard alias on the worker. +PREPARE cached_upsert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2) + ON CONFLICT (key) DO UPDATE SET value = dist_table_ts.value + EXCLUDED.value; + +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(501, 5); +EXECUTE cached_upsert(501, 5); + +-- key 500 accumulated 6 increments, key 501 accumulated 2 +SELECT key, value FROM dist_table_ts WHERE key >= 500 AND key < 600 ORDER BY key; + +DEALLOCATE cached_upsert; + +-- Multi-row INSERT is not cacheable: each shard's task carries only its own +-- subset of VALUES rows, so the rows must not be duplicated across shards. +PREPARE cached_multirow(int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, 1), ($1 + 1, 2); + +EXECUTE cached_multirow(600); +EXECUTE cached_multirow(602); +EXECUTE cached_multirow(604); +EXECUTE cached_multirow(606); +EXECUTE cached_multirow(608); +EXECUTE cached_multirow(610); + +-- exactly 12 rows, one per key, no duplicates +SELECT count(*) AS row_count, count(DISTINCT key) AS distinct_keys + FROM dist_table_ts WHERE key >= 600; + +DEALLOCATE cached_multirow; + +-- ============================================================ +-- Test 5: GUC OFF baseline — same queries produce identical results +-- ============================================================ + +SET citus.enable_prepared_statement_caching = off; + +PREPARE uncached_select(int) AS + SELECT key, value FROM dist_table WHERE key = $1; + +EXECUTE uncached_select(1); +EXECUTE uncached_select(2); +EXECUTE uncached_select(3); +EXECUTE uncached_select(4); +EXECUTE uncached_select(5); +EXECUTE uncached_select(6); +EXECUTE uncached_select(7); +EXECUTE uncached_select(8); +EXECUTE uncached_select(9); +EXECUTE uncached_select(10); + +DEALLOCATE uncached_select; + +-- ============================================================ +-- Test 6: Multiple prepared statements in same session +-- ============================================================ + +SET citus.enable_prepared_statement_caching = on; + +PREPARE stmt_a(int) AS SELECT key, value FROM dist_table WHERE key = $1; +PREPARE stmt_b(int) AS SELECT key, label FROM dist_table WHERE key = $1; +PREPARE stmt_c(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2); + +-- Interleave executions to verify independent caching +EXECUTE stmt_a(1); +EXECUTE stmt_b(1); +EXECUTE stmt_a(2); +EXECUTE stmt_b(2); +EXECUTE stmt_a(3); +EXECUTE stmt_b(3); +EXECUTE stmt_a(4); +EXECUTE stmt_b(4); +EXECUTE stmt_a(5); +EXECUTE stmt_b(5); +EXECUTE stmt_a(6); +EXECUTE stmt_b(6); +EXECUTE stmt_a(7); +EXECUTE stmt_b(7); +EXECUTE stmt_c(300, 3000); +EXECUTE stmt_c(301, 3010); +EXECUTE stmt_c(302, 3020); +EXECUTE stmt_c(303, 3030); +EXECUTE stmt_c(304, 3040); +EXECUTE stmt_c(305, 3050); +EXECUTE stmt_c(306, 3060); + +-- Verify inserts from stmt_c +SELECT key, value FROM dist_table_ts WHERE key >= 300 AND key < 400 ORDER BY key; + +DEALLOCATE stmt_a; +DEALLOCATE stmt_b; +DEALLOCATE stmt_c; + +-- ============================================================ +-- Test 7: Connection loss re-prepare — force worker connection +-- close, verify subsequent EXECUTE still works +-- ============================================================ + +SET citus.enable_prepared_statement_caching = on; + +PREPARE reconnect_test(int) AS + SELECT key, value FROM dist_table WHERE key = $1; + +-- Execute enough times to get into generic plan + cache hit +EXECUTE reconnect_test(1); +EXECUTE reconnect_test(1); +EXECUTE reconnect_test(1); +EXECUTE reconnect_test(1); +EXECUTE reconnect_test(1); +EXECUTE reconnect_test(1); +EXECUTE reconnect_test(1); + +-- Force all cached connections to be dropped by setting lifetime to 0 +SET citus.max_cached_connection_lifetime TO '0s'; + +-- The next execution should get a new connection, re-prepare, and succeed +EXECUTE reconnect_test(1); +EXECUTE reconnect_test(2); +EXECUTE reconnect_test(3); + +-- Restore default +RESET citus.max_cached_connection_lifetime; + +DEALLOCATE reconnect_test; + +-- ============================================================ +-- Cleanup +-- ============================================================ + +SET citus.enable_prepared_statement_caching = off; +DROP SCHEMA prepared_stmt_caching CASCADE; From d49e8af873b49959f37da6d3c1c087540b7ac583 Mon Sep 17 00:00:00 2001 From: Colm McHugh Date: Wed, 9 Sep 2026 16:34:18 +0000 Subject: [PATCH 2/2] Address co-pilot CR comments Review comments: - Drop the write-only paramTypes, paramCount and parameterizedQueryString from PreparedStatementCacheEntry; free queryString on the PREPARED_STMT_FAILED path; add an insert alias when deparsing shard-INSERTs. - Carry `relationRowLockList` on the fast-path task, so `SELECT ... FOR UPDATE` still opens a remote transaction and holds the row lock to the end of the enclosing one. - Publish the cache entry only after the worker accepts the statement; a rejected `PQprepare` was leaving an entry naming a statement that does not exist. - Recheck `ShardExists()` after `AcquireMetadataLocks()` and decline, so a concurrent split reroutes through normal planning instead of erroring. - Don't attach templates to custom plans: their planId is new per execution, so every execution prepared again and consumed a cache slot it could never hit. Gate on plan reuse, as `IsLocalPlanCachingSupported()` does. - Set `workerJob->partitionKeyValue` per execution; local execution and query statistics read it from the job. Save the planner's value first and put it back at end of scan rather than nulling the field: the fast path's Const dies with the portal, but clearing it made `ModifyJobNeedsEvaluation()` treat the distribution column value as unknown and forced every deferred-pruning INSERT through coordinator evaluation, with or without the cache. - `PQprepare()` blocks uninterruptibly, so a cancel against a slow worker was lost. Use `PQsendPrepare()` with `GetRemoteCommandResult()`. The wait is still synchronous for sibling placements; full executor integration is follow-up. - Share `SetRemoteRowMode()` with `SendNextQuery()`; caching was hard-coded to single-row mode and ignored `citus.executor_chunk_size`. - Decline caching while `citus.stat_tenants_track` is on. The tenant id travels in the query text, and a reused statement would freeze the first tenant rather than attribute per execution. Also address issues discovered with running multi, multi-1, and multi-mx test suites with `citus.enable_prepared_statement_caching` enabled: - Crash on shard key mismatch: `strip_implicit_coercions()` records the raw Param, so `PREPARE p(int)` on a numeric column hashed an int4 Datum as numeric. Instead, coerce via `TransformPartitionRestrictionValue()` and decline if not coercible. - Don't bypass coordinator evaluation and incorrectly ship nextval() to a worker. Correctly propagate `requiresCoordinatorEvaluation` on the deferred-pruning path. - Assert fail on unused params of a user-defined type; now mark against `jobQueryForPrepare` when that is present, instead of depending on the GUC value. - All three surfaced as the same assert: a rejected `PQprepare` was reported as MULTI_CONNECTION_LOST and sent the executor down the reconnect path with placements still attached. - `BuildFastPathTask()` should not ignore round-robin task assignment. - `parametersInQueryStringResolved` was set unconditionally, suppressing parameter types for a template that still carries $n. - EXPLAIN ANALYZE inherited `jobQueryForPrepare` through copyObject, and the fast path leaked $1 into the reported worker plan. - Deparse path: dangling pointer from freeing the task's own string, and a per-tuple memory context for a string the task retains. Drop the EVALUATE_FUNCTIONS evaluation mode and ExecuteCoordinatorEvaluableFunctions(). citus_evaluate_expr() bails out for every mode except EVALUATE_FUNCTIONS_PARAMS, so the helper never evaluated anything, and nothing called it; caching declines on requiresCoordinatorEvaluation instead. citus_clauses is untouched by this branch again. Tests asserted only results, which are identical whether or not the cache is engaged. Test 13 now asserts the wire protocol, Test 14 the remote transaction for FOR UPDATE, Test 15 the INSERT/UPDATE/DELETE dispatch, and Test 7 the re-prepare after a connection drop. DML reaches the cache through its own eligibility and task-building branches: with template saving disabled for non-SELECT, every diff came from Test 15 and no pre-existing test noticed. The suite sets `citus.stat_tenants_track = 'ALL'`, which the cache declines, so these tests set it to 'none'. Adds a CitusPreparedStatementCachingConfig arbitrary-configs class, which inherits the 'none' default and is the main CI coverage. Known limitation: caching is disabled while `citus.stat_tenants_track` is enabled. No regression test for: the SendRemotePrepare changes, the concurrent-split recheck, and the stale cache entry; all verified by hand. --- .../distributed/connection/remote_commands.c | 33 +- .../distributed/executor/adaptive_executor.c | 38 +- .../distributed/executor/citus_custom_scan.c | 20 +- .../distributed/executor/local_executor.c | 3 + .../executor/prepared_statement_cache.c | 263 ++++--- src/backend/distributed/metadata/XXnahQZC | 0 .../distributed/planner/deparse_shard_query.c | 13 +- .../distributed/planner/distributed_planner.c | 14 +- .../distributed/planner/multi_explain.c | 3 + .../planner/multi_router_planner.c | 16 + src/backend/distributed/utils/citus_clauses.c | 29 +- .../distributed/utils/citus_copyfuncs.c | 1 + src/include/distributed/citus_clauses.h | 4 - .../distributed/multi_physical_planner.h | 8 + .../distributed/multi_router_planner.h | 1 + .../distributed/prepared_statement_cache.h | 10 +- src/test/regress/bin/normalize.sed | 3 + src/test/regress/citus_tests/config.py | 8 + .../expected/prepared_statement_caching.out | 693 +++++++++++++++++- src/test/regress/expected/single_node.out | 80 +- .../sql/prepared_statement_caching.sql | 407 +++++++++- src/test/regress/sql/single_node.sql | 55 ++ 22 files changed, 1538 insertions(+), 164 deletions(-) delete mode 100644 src/backend/distributed/metadata/XXnahQZC diff --git a/src/backend/distributed/connection/remote_commands.c b/src/backend/distributed/connection/remote_commands.c index ba3c71dead8..82f8f3777b3 100644 --- a/src/backend/distributed/connection/remote_commands.c +++ b/src/backend/distributed/connection/remote_commands.c @@ -577,9 +577,9 @@ SendRemoteCommand(MultiConnection *connection, const char *command) /* - * SendRemotePrepare wraps PQprepare() to prepare a statement on a worker - * connection. This is a synchronous call that blocks until the prepare - * completes. Returns 1 on success, 0 on failure. + * SendRemotePrepare prepares a statement on a worker connection and waits for + * the worker to confirm it. The wait is interruptible, so a cancel is honoured + * even when the worker is slow to answer. Returns 1 on success, 0 on failure. */ int SendRemotePrepare(MultiConnection *connection, const char *stmtName, @@ -594,15 +594,34 @@ SendRemotePrepare(MultiConnection *connection, const char *stmtName, return 0; } - PGresult *result = PQprepare(pgConn, stmtName, query, nParams, paramTypes); - if (PQresultStatus(result) != PGRES_COMMAND_OK) + Assert(PQisnonblocking(pgConn)); + + if (PQsendPrepare(pgConn, stmtName, query, nParams, paramTypes) == 0) + { + return 0; + } + + bool raiseInterrupts = true; + PGresult *result = GetRemoteCommandResult(connection, raiseInterrupts); + if (result == NULL) { - ReportResultError(connection, result, WARNING); - PQclear(result); return 0; } + if (PQresultStatus(result) != PGRES_COMMAND_OK) + { + /* + * Raise for a lost connection as well, which GetRemoteCommandResult() + * reports as a synthetic PGRES_FATAL_ERROR. Returning failure instead + * marks the connection lost while placements are still attached, and + * RestartConnection() asserts that they are not. + */ + ReportResultError(connection, result, ERROR); + } + PQclear(result); + ForgetResults(connection); + return 1; } diff --git a/src/backend/distributed/executor/adaptive_executor.c b/src/backend/distributed/executor/adaptive_executor.c index 3f97049f4cd..7845bad38a6 100644 --- a/src/backend/distributed/executor/adaptive_executor.c +++ b/src/backend/distributed/executor/adaptive_executor.c @@ -741,6 +741,7 @@ static bool StartPlacementExecutionOnSession(TaskPlacementExecution *placementEx WorkerSession *session); static bool SendNextQuery(TaskPlacementExecution *placementExecution, WorkerSession *session); +static bool SetRemoteRowMode(MultiConnection *connection); static void ConnectionStateMachine(WorkerSession *session); static bool HasUnfinishedTaskForSession(WorkerSession *session); static void HandleMultiConnectionSuccess(WorkerSession *session, bool newConnection); @@ -926,16 +927,26 @@ AdaptiveExecutorStart(CitusScanState *scanState) * and never used in the query, mark such parameters' type as Invalid(0), * which will be used later in ExtractParametersFromParamList() to map them * to a generic datatype. Skip for dynamic parameters. - * - * When prepared statement caching is enabled, skip this step entirely: - * the params appear "unreferenced" in job->jobQuery because they were - * resolved to constants there, but they are still needed with their - * original types for the parameterized query sent via PQprepare. */ - if (paramListInfo && !paramListInfo->paramFetch && !EnablePreparedStatementCaching) + if (paramListInfo && !paramListInfo->paramFetch) { + /* + * A cached statement is deparsed from the saved template, which still + * carries the Params. job->jobQuery has them resolved to constants, so + * marking against it would strip types the worker still needs. + */ + Node *queryForParams = (Node *) job->jobQuery; + if (taskList != NIL) + { + Task *firstTask = (Task *) linitial(taskList); + if (firstTask->jobQueryForPrepare != NULL) + { + queryForParams = (Node *) firstTask->jobQueryForPrepare; + } + } + paramListInfo = copyParamList(paramListInfo); - MarkUnreferencedExternParams((Node *) job->jobQuery, paramListInfo); + MarkUnreferencedExternParams(queryForParams, paramListInfo); } DistributedExecution *execution = CreateDistributedExecution( @@ -4274,7 +4285,7 @@ SendNextQuery(TaskPlacementExecution *placementExecution, binaryResults, &queryString); if (cacheStatus == PREPARED_STMT_SENT) { - return true; + return SetRemoteRowMode(connection); } else if (cacheStatus == PREPARED_STMT_FAILED) { @@ -4333,6 +4344,17 @@ SendNextQuery(TaskPlacementExecution *placementExecution, return false; } + return SetRemoteRowMode(connection); +} + + +/* + * SetRemoteRowMode puts the connection into the incremental result mode the + * executor expects, chunked where libpq supports it. + */ +static bool +SetRemoteRowMode(MultiConnection *connection) +{ #ifdef LIBPQ_HAS_CHUNK_MODE int rowMode = PQsetChunkedRowsMode(connection->pgConn, ExecutorChunkSize); #else diff --git a/src/backend/distributed/executor/citus_custom_scan.c b/src/backend/distributed/executor/citus_custom_scan.c index c58c79cc481..ebac9263fcc 100644 --- a/src/backend/distributed/executor/citus_custom_scan.c +++ b/src/backend/distributed/executor/citus_custom_scan.c @@ -978,20 +978,22 @@ CitusEndScanCommon(CitusScanState *scanState) /* * Clear mutable per-execution state so the cached plan is clean for * the next execution. The cache-hit fast paths in CitusBeginReadOnlyScan() - * and CitusBeginModifyScan() store a Task list directly on the original - * plan's workerJob; those Tasks live in the per-execution memory context - * and become dangling after EndScan. In assert-checking builds the next - * execution's GetDistributedPlan() → copyObject() would traverse freed - * memory without this reset. + * and CitusBeginModifyScan() store a Task list and a partition key Const + * directly on the original plan's workerJob; both live in the per-execution + * memory context and become dangling after EndScan. In assert-checking + * builds the next execution's GetDistributedPlan() → copyObject() would + * traverse freed memory without this reset. * - * Only deferred-pruning plans need this: their taskList is rebuilt - * per-execution. Non-deferred plans carry their real taskList from - * planning and must not be touched. + * partitionKeyValue goes back to what the planner set rather than to NULL: + * ModifyJobNeedsEvaluation() reads it to decide whether the distribution + * column value is already known, and nulling it forces every later + * execution through needless coordinator evaluation. */ - if (workerJob != NULL && workerJob->deferredPruning) + if (workerJob != NULL && workerJob->savedJobQueryForCaching != NULL) { workerJob->taskList = NIL; workerJob->parametersInJobQueryResolved = false; + workerJob->partitionKeyValue = workerJob->plannerPartitionKeyValue; } } diff --git a/src/backend/distributed/executor/local_executor.c b/src/backend/distributed/executor/local_executor.c index bb3d9d02f19..d19c48977bf 100644 --- a/src/backend/distributed/executor/local_executor.c +++ b/src/backend/distributed/executor/local_executor.c @@ -406,6 +406,9 @@ ExecuteLocalTaskListExtended(List *taskList, if (queryForDeparse->commandType == CMD_INSERT) { + /* upserts reference the target by name, which becomes the shard name */ + AddInsertAliasIfNeeded(queryForDeparse); + deparse_shard_query(queryForDeparse, task->anchorDistributedTableId, task->anchorShardId, &buf); diff --git a/src/backend/distributed/executor/prepared_statement_cache.c b/src/backend/distributed/executor/prepared_statement_cache.c index d24230752f8..2e23806447a 100644 --- a/src/backend/distributed/executor/prepared_statement_cache.c +++ b/src/backend/distributed/executor/prepared_statement_cache.c @@ -37,18 +37,40 @@ #include "distributed/listutils.h" #include "distributed/local_plan_cache.h" #include "distributed/metadata_cache.h" +#include "distributed/metadata_utility.h" #include "distributed/multi_executor.h" +#include "distributed/multi_explain.h" +#include "distributed/multi_physical_planner.h" #include "distributed/multi_router_planner.h" #include "distributed/prepared_statement_cache.h" #include "distributed/remote_commands.h" #include "distributed/shard_cleaner.h" +#include "distributed/shard_pruning.h" #include "distributed/shardinterval_utils.h" +#include "distributed/stats/stat_tenants.h" /* GUC: citus.enable_prepared_statement_caching */ bool EnablePreparedStatementCaching = false; +/* + * PreparedStatementCachingUsable returns whether the cache may be used at all + * for this execution. + * + * Tenant statistics reach the worker as a comment on the query text. A reused + * prepared statement cannot carry that per execution -- it would freeze the + * tenant of whichever execution prepared it -- so decline while tracking is on + * rather than mis-attribute. + */ +static bool +PreparedStatementCachingUsable(void) +{ + return EnablePreparedStatementCaching && + StatTenantsTrack == STAT_TENANTS_TRACK_NONE; +} + + /* * PreparedStatementCacheCreate allocates a new hash table for caching * prepared statement entries on a single worker connection. The hash @@ -96,21 +118,17 @@ PreparedStatementCacheLookup(HTAB *cache, uint64 planId, uint64 shardId) /* - * PreparedStatementCacheInsert inserts a new entry for (planId, shardId). - * Returns the new entry on success, or NULL if the cache has reached - * MAX_CACHED_STMTS_PER_CONNECTION (caller should fall back to plain SQL). + * PreparedStatementCacheInsert records that stmtName has been prepared on this + * connection for (planId, shardId). * - * The caller is responsible for filling in the returned entry's fields - * (stmtName, paramTypes, paramCount, parameterizedQueryString). + * Only call this once the worker has accepted the statement: an entry for a + * statement that does not exist makes every later execution on this connection + * fail with "prepared statement does not exist". */ PreparedStatementCacheEntry * -PreparedStatementCacheInsert(HTAB *cache, uint64 planId, uint64 shardId) +PreparedStatementCacheInsert(HTAB *cache, uint64 planId, uint64 shardId, + const char *stmtName) { - if (hash_get_num_entries(cache) >= MAX_CACHED_STMTS_PER_CONNECTION) - { - return NULL; - } - PreparedStatementCacheKey key; memset(&key, 0, sizeof(key)); @@ -122,19 +140,11 @@ PreparedStatementCacheInsert(HTAB *cache, uint64 planId, uint64 shardId) (PreparedStatementCacheEntry *) hash_search(cache, &key, HASH_ENTER, &found); - if (found) + if (!found) { - /* already exists — return existing entry */ - return entry; + strlcpy(entry->stmtName, stmtName, MAX_STMT_NAME_LENGTH); } - /* initialize the new entry with auto-generated statement name */ - SafeSnprintf(entry->stmtName, MAX_STMT_NAME_LENGTH, - "__citus_stmt_%ld", (long) hash_get_num_entries(cache)); - entry->paramTypes = NULL; - entry->paramCount = 0; - entry->parameterizedQueryString = NULL; - return entry; } @@ -156,26 +166,6 @@ PreparedStatementCacheDestroy(HTAB **cache_ptr) return; } - /* - * Free dynamically allocated fields in each entry before destroying - * the hash table itself. - */ - HASH_SEQ_STATUS status; - PreparedStatementCacheEntry *entry; - - hash_seq_init(&status, cache); - while ((entry = hash_seq_search(&status)) != NULL) - { - if (entry->paramTypes != NULL) - { - pfree(entry->paramTypes); - } - if (entry->parameterizedQueryString != NULL) - { - pfree(entry->parameterizedQueryString); - } - } - hash_destroy(cache); *cache_ptr = NULL; } @@ -226,7 +216,7 @@ PreparedStatementCacheSendQuery(MultiConnection *connection, Task *task, ParamListInfo paramListInfo, bool binaryResults, char **fallbackQueryString) { - if (!EnablePreparedStatementCaching || task->jobQueryForPrepare == NULL || + if (!PreparedStatementCachingUsable() || task->jobQueryForPrepare == NULL || paramListInfo == NULL) { return PREPARED_STMT_NOT_APPLICABLE; @@ -257,10 +247,9 @@ PreparedStatementCacheSendQuery(MultiConnection *connection, Task *task, " shard " UINT64_FORMAT, task->preparedStatementPlanId, task->anchorShardId); - cacheEntry = PreparedStatementCacheInsert(connection->preparedStatementCache, - task->preparedStatementPlanId, - task->anchorShardId); - if (cacheEntry == NULL) + HTAB *cache = connection->preparedStatementCache; + + if (hash_get_num_entries(cache) >= MAX_CACHED_STMTS_PER_CONNECTION) { /* * Cache full. The fast-path task has no query string of its own, so @@ -272,24 +261,30 @@ PreparedStatementCacheSendQuery(MultiConnection *connection, Task *task, return PREPARED_STMT_FALLBACK; } + char stmtName[MAX_STMT_NAME_LENGTH]; + SafeSnprintf(stmtName, MAX_STMT_NAME_LENGTH, "__citus_stmt_%ld", + (long) hash_get_num_entries(cache) + 1); + char *queryString = DeparseTaskTemplate(task); - if (SendRemotePrepare(connection, cacheEntry->stmtName, queryString, + if (SendRemotePrepare(connection, stmtName, queryString, parameterCount, parameterTypes) == 0) { + pfree(queryString); connection->connectionState = MULTI_CONNECTION_LOST; return PREPARED_STMT_FAILED; } - Size paramTypesSize = parameterCount * sizeof(Oid); - cacheEntry->paramTypes = MemoryContextAlloc(TopMemoryContext, paramTypesSize); - memcpy_s(cacheEntry->paramTypes, paramTypesSize, parameterTypes, - paramTypesSize); - cacheEntry->paramCount = parameterCount; - cacheEntry->parameterizedQueryString = - MemoryContextStrdup(TopMemoryContext, queryString); - pfree(queryString); + + /* + * Publish only now: SendRemotePrepare raises on a rejected statement, so + * an earlier insert would leave an entry naming a statement the worker + * never created. + */ + cacheEntry = PreparedStatementCacheInsert(cache, + task->preparedStatementPlanId, + task->anchorShardId, stmtName); } else { @@ -306,12 +301,6 @@ PreparedStatementCacheSendQuery(MultiConnection *connection, Task *task, return PREPARED_STMT_FAILED; } - if (PQsetSingleRowMode(connection->pgConn) == 0) - { - connection->connectionState = MULTI_CONNECTION_LOST; - return PREPARED_STMT_FAILED; - } - return PREPARED_STMT_SENT; } @@ -332,13 +321,23 @@ PreparedStatementCacheSaveTemplate(DistributedPlan *originalPlan) Job *originalJob = originalPlan->workerJob; Query *jobQuery = originalJob->jobQuery; - if (!EnablePreparedStatementCaching) + if (!PreparedStatementCachingUsable()) { return NULL; } if (jobQuery->commandType != CMD_SELECT) { + /* + * nextval(), now() and other coordinator-evaluated expressions are + * resolved into the per-execution copy of the query, never into this + * template, so a cached statement would evaluate them on the worker. + */ + if (originalJob->requiresCoordinatorEvaluation) + { + return NULL; + } + /* * Multi-row INSERT can't be cached: each shard's task carries only its * own subset of VALUES rows, but the statement is deparsed once from @@ -357,6 +356,9 @@ PreparedStatementCacheSaveTemplate(DistributedPlan *originalPlan) MemoryContextSwitchTo(GetMemoryChunkContext(originalPlan)); originalJob->savedJobQueryForCaching = copyObject(jobQuery); MemoryContextSwitchTo(oldContext); + + /* the fast path has not run yet, so this is still the planner's value */ + originalJob->plannerPartitionKeyValue = originalJob->partitionKeyValue; } return originalJob->savedJobQueryForCaching; @@ -372,7 +374,17 @@ void PreparedStatementCacheAttachToTasks(DistributedPlan *currentPlan, Job *workerJob, Query *savedJobQuery) { - if (!EnablePreparedStatementCaching || savedJobQuery == NULL) + if (!PreparedStatementCachingUsable() || savedJobQuery == NULL) + { + return; + } + + /* + * Only attach to a plan that is being reused. A custom plan is rebuilt with a + * fresh planId per execution, so its entries could never be hit and would fill + * the connection's cache. Same reuse test as IsLocalPlanCachingSupported(). + */ + if (currentPlan->numberOfTimesExecuted < 1) { return; } @@ -397,9 +409,14 @@ PreparedStatementCacheAttachToTasks(DistributedPlan *currentPlan, Job *workerJob /* * FastPathShardInterval returns the shard the distribution key parameter routes * to, or NULL if the fast path cannot be used for this execution. + * + * The value is returned via partitionKeyValue, coerced to the distribution + * column's type. Single-row INSERT records the Param with implicit coercions + * stripped, so the parameter type can differ from the column type. */ static ShardInterval * -FastPathShardInterval(DistributedPlan *plan, Job *workerJob, EState *estate) +FastPathShardInterval(DistributedPlan *plan, Job *workerJob, EState *estate, + Const **partitionKeyValue) { int paramId = workerJob->distributionKeyParamId; ParamListInfo paramListInfo = estate->es_param_list_info; @@ -417,18 +434,47 @@ FastPathShardInterval(DistributedPlan *plan, Job *workerJob, EState *estate) Oid relationId = linitial_oid(plan->relationIdList); CitusTableCacheEntry *tableEntry = GetCitusTableCacheEntry(relationId); + Var *partitionColumn = tableEntry->partitionColumn; + + if (partitionColumn == NULL) + { + return NULL; + } - return FindShardInterval(param->value, tableEntry); + int16 typeLength; + bool typeByValue; + get_typlenbyval(param->ptype, &typeLength, &typeByValue); + Const *valueConst = makeConst(param->ptype, -1, InvalidOid, (int) typeLength, + param->value, false, typeByValue); + + if (param->ptype != partitionColumn->vartype) + { + bool missingOk = true; + valueConst = TransformPartitionRestrictionValue(partitionColumn, valueConst, + missingOk); + if (valueConst == NULL || valueConst->constisnull) + { + return NULL; + } + } + + *partitionKeyValue = valueConst; + + return FindShardInterval(valueConst->constvalue, tableEntry); } /* * BuildFastPathTask builds the minimal Task for a single-shard execution, * bypassing plan copying, coordinator evaluation and task regeneration. + * + * Returns NULL when the shard has no placement to run on, so the caller can + * fall back to normal planning and raise its "found no worker with all shard + * placements" error. */ static Task * -BuildFastPathTask(DistributedPlan *plan, Job *workerJob, EState *estate, - ShardInterval *shardInterval, bool isModify) +BuildFastPathTask(DistributedPlan *plan, Job *workerJob, ShardInterval *shardInterval, + Const *partitionKeyValue, bool isModify) { List *shardIntervalListList = list_make1(list_make1(shardInterval)); bool shardsPresent = false; @@ -438,26 +484,34 @@ BuildFastPathTask(DistributedPlan *plan, Job *workerJob, EState *estate, CreateTaskPlacementListForShardIntervals(shardIntervalListList, shardsPresent, true, false); + if (placementList == NIL) + { + return NULL; + } + + /* + * Modifications run on every placement and are assigned first-replica by the + * caller, so only reads honour citus.task_assignment_policy. + */ + if (!isModify && TaskAssignmentPolicy == TASK_ASSIGNMENT_ROUND_ROBIN) + { + placementList = RemoveCoordinatorPlacementIfNotSingleNode(placementList); + placementList = RoundRobinReorder(placementList); + } + Task *task = CitusMakeNode(Task); task->taskType = isModify ? MODIFY_TASK : READ_TASK; task->anchorShardId = shardInterval->shardId; task->anchorDistributedTableId = linitial_oid(plan->relationIdList); task->taskPlacementList = placementList; task->queryCount = 1; - task->parametersInQueryStringResolved = true; task->preparedStatementPlanId = plan->planId; task->jobQueryForPrepare = workerJob->savedJobQueryForCaching; task->relationShardList = relationShardList; + task->relationRowLockList = + RelationRowLockListForQuery(workerJob->savedJobQueryForCaching); task->colocationId = workerJob->colocationId; - - ParamExternData *param = - &estate->es_param_list_info->params[workerJob->distributionKeyParamId - 1]; - int16 typeLength; - bool typeByValue; - get_typlenbyval(param->ptype, &typeLength, &typeByValue); - task->partitionKeyValue = makeConst(param->ptype, -1, InvalidOid, - (int) typeLength, param->value, false, - typeByValue); + task->partitionKeyValue = partitionKeyValue; return task; } @@ -480,13 +534,22 @@ PreparedStatementCacheTryFastPath(struct CitusScanState *scanStateArg, EState *e Job *workerJob = originalPlan->workerJob; /* the first execution populates the template the fast path depends on */ - if (!EnablePreparedStatementCaching || + if (!PreparedStatementCachingUsable() || originalPlan->numberOfTimesExecuted == 0 || workerJob->savedJobQueryForCaching == NULL) { return false; } + /* + * EXPLAIN ANALYZE wraps the task's query string, which would then carry the + * template's Params into the reported worker plan. + */ + if (RequestedForExplainAnalyze(scanState)) + { + return false; + } + /* volatile functions must be re-evaluated per execution on the coordinator */ if (isModify && (!workerJob->deferredPruning || workerJob->requiresCoordinatorEvaluation)) @@ -494,33 +557,55 @@ PreparedStatementCacheTryFastPath(struct CitusScanState *scanStateArg, EState *e return false; } + Const *partitionKeyValue = NULL; ShardInterval *shardInterval = FastPathShardInterval(originalPlan, workerJob, - estate); - if (shardInterval == NULL || (isModify && !ShardExists(shardInterval->shardId))) + estate, &partitionKeyValue); + if (shardInterval == NULL) { return false; } - Task *task = BuildFastPathTask(originalPlan, workerJob, estate, shardInterval, - isModify); + Task *task = BuildFastPathTask(originalPlan, workerJob, shardInterval, + partitionKeyValue, isModify); + if (task == NULL) + { + return false; + } workerJob->taskList = list_make1(task); workerJob->parametersInJobQueryResolved = true; - elog(DEBUG2, "prepared statement cache-hit fast path%s: plan " UINT64_FORMAT - " shard " UINT64_FORMAT, - isModify ? " (DML)" : "", originalPlan->planId, shardInterval->shardId); - - /* the executor reads the plan back from the scan state */ - scanState->distributedPlan = originalPlan; + /* local execution and query stats read the key from the job, not the task */ + workerJob->partitionKeyValue = partitionKeyValue; if (isModify) { AcquireMetadataLocks(workerJob->taskList); - EnsureAnchorShardsInJobExist(workerJob); + + /* + * A concurrent split may have dropped the shard between pruning and + * locking. Normal planning reroutes here, but it does so by rewriting + * jobQuery in place, which would pin the cached plan to this execution's + * shards. Decline instead, so normal planning reroutes on its own copy. + */ + if (!ShardExists(shardInterval->shardId)) + { + workerJob->taskList = NIL; + workerJob->parametersInJobQueryResolved = false; + workerJob->partitionKeyValue = workerJob->plannerPartitionKeyValue; + return false; + } + workerJob->taskList = FirstReplicaAssignTaskList(workerJob->taskList); } + elog(DEBUG2, "prepared statement cache-hit fast path%s: plan " UINT64_FORMAT + " shard " UINT64_FORMAT, + isModify ? " (DML)" : "", originalPlan->planId, shardInterval->shardId); + + /* the executor reads the plan back from the scan state */ + scanState->distributedPlan = originalPlan; + /* * A fast-path task has no query string, so local execution needs a cached * local plan rather than deparsed SQL. diff --git a/src/backend/distributed/metadata/XXnahQZC b/src/backend/distributed/metadata/XXnahQZC deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/src/backend/distributed/planner/deparse_shard_query.c b/src/backend/distributed/planner/deparse_shard_query.c index 18dffc90363..9724d692bc0 100644 --- a/src/backend/distributed/planner/deparse_shard_query.c +++ b/src/backend/distributed/planner/deparse_shard_query.c @@ -798,11 +798,23 @@ TaskQueryString(Task *task) * so subsequent calls don't re-deparse. */ Query *queryForDeparse = copyObject(task->jobQueryForPrepare); + + /* + * The task retains the string, so allocate it in the task's own + * context rather than CurrentMemoryContext, which may be a + * short-lived per-tuple context. + */ + MemoryContext previousContext = + MemoryContextSwitchTo(GetMemoryChunkContext(task)); StringInfoData buf; initStringInfo(&buf); + MemoryContextSwitchTo(previousContext); if (queryForDeparse->commandType == CMD_INSERT) { + /* upserts reference the target by name, which becomes the shard name */ + AddInsertAliasIfNeeded(queryForDeparse); + deparse_shard_query(queryForDeparse, task->anchorDistributedTableId, task->anchorShardId, &buf); @@ -815,7 +827,6 @@ TaskQueryString(Task *task) } SetTaskQueryString(task, buf.data); - pfree(buf.data); return task->taskQuery.data.queryStringLazy; } diff --git a/src/backend/distributed/planner/distributed_planner.c b/src/backend/distributed/planner/distributed_planner.c index a59eb810617..8633d1a5dab 100644 --- a/src/backend/distributed/planner/distributed_planner.c +++ b/src/backend/distributed/planner/distributed_planner.c @@ -1449,16 +1449,18 @@ GetDistributedPlan(CustomScan *customScan) Assert(CitusIsA(node, DistributedPlan)); /* - * Clear stale task pointers before CheckNodeCopyAndSerialization. - * When prepared statement caching fast-path is active, the previous - * execution may have set workerJob->taskList on the original plan. - * That task's memory has been freed (portal context destroyed), so - * serializing the plan now would dereference dangling pointers. + * Undo what the previous execution's fast path left on the original plan. + * Its Task and partition key Const were allocated in a portal context that + * is now gone, so serializing the plan would follow dangling pointers. + * CitusEndScanCommon() does this too, but an execution that errors out + * never reaches it. */ DistributedPlan *plan = (DistributedPlan *) node; - if (plan->workerJob != NULL && plan->workerJob->deferredPruning) + if (plan->workerJob != NULL && plan->workerJob->savedJobQueryForCaching != NULL) { plan->workerJob->taskList = NIL; + plan->workerJob->parametersInJobQueryResolved = false; + plan->workerJob->partitionKeyValue = plan->workerJob->plannerPartitionKeyValue; } CheckNodeCopyAndSerialization(node); diff --git a/src/backend/distributed/planner/multi_explain.c b/src/backend/distributed/planner/multi_explain.c index 9c00f34c2fa..84a868e0cb1 100644 --- a/src/backend/distributed/planner/multi_explain.c +++ b/src/backend/distributed/planner/multi_explain.c @@ -1938,6 +1938,9 @@ ExplainAnalyzeTaskList(List *originalTaskList, SetTaskQueryStringList(explainAnalyzeTask, list_make2(wrappedQuery, fetchQuery)); + /* the wrapper replaces the query, so the saved template no longer applies */ + explainAnalyzeTask->jobQueryForPrepare = NULL; + TupleDestination *originalTaskDest = originalTask->tupleDest ? originalTask->tupleDest : defaultTupleDest; diff --git a/src/backend/distributed/planner/multi_router_planner.c b/src/backend/distributed/planner/multi_router_planner.c index 69b7b896c73..4c95e290088 100644 --- a/src/backend/distributed/planner/multi_router_planner.c +++ b/src/backend/distributed/planner/multi_router_planner.c @@ -1983,6 +1983,7 @@ RouterJob(Query *originalQuery, PlannerRestrictionContext *plannerRestrictionCon Job *job = CreateJob(originalQuery); job->deferredPruning = true; job->distributionKeyParamId = fastPathRestrictionContext->distributionKeyParamId; + job->requiresCoordinatorEvaluation = requiresCoordinatorEvaluation; ereport(DEBUG2, (errmsg("Deferred pruning for a fast-path router " "query"))); @@ -2749,6 +2750,21 @@ RowLocksOnRelations(Node *node, List **relationRowLockList) } +/* + * RelationRowLockListForQuery returns the FOR UPDATE/SHARE row locks the query + * takes on Citus tables. + */ +List * +RelationRowLockListForQuery(Query *query) +{ + List *relationRowLockList = NIL; + + RowLocksOnRelations((Node *) query, &relationRowLockList); + + return relationRowLockList; +} + + /* * SelectsFromDistributedTable checks if there is a select on a distributed * table by looking into range table entries. diff --git a/src/backend/distributed/utils/citus_clauses.c b/src/backend/distributed/utils/citus_clauses.c index 386536b8831..7086cec4eab 100644 --- a/src/backend/distributed/utils/citus_clauses.c +++ b/src/backend/distributed/utils/citus_clauses.c @@ -84,25 +84,6 @@ ExecuteCoordinatorEvaluableExpressions(Query *query, PlanState *planState) } -/* - * ExecuteCoordinatorEvaluableFunctions evaluates function calls on the - * coordinator (e.g. now(), nextval()) but leaves Param nodes untouched. - * This is used for prepared statement caching where we need to evaluate - * volatile/stable functions but preserve parameter placeholders for the - * prepared statement template. - */ -void -ExecuteCoordinatorEvaluableFunctions(Query *query, PlanState *planState) -{ - CoordinatorEvaluationContext coordinatorEvaluationContext; - - coordinatorEvaluationContext.planState = planState; - coordinatorEvaluationContext.evaluationMode = EVALUATE_FUNCTIONS; - - PartiallyEvaluateExpression((Node *) query, &coordinatorEvaluationContext); -} - - /* * PartiallyEvaluateExpression descends into an expression tree to evaluate * expressions that can be resolved to a constant on the master. Expressions @@ -131,13 +112,6 @@ PartiallyEvaluateExpression(Node *expression, { Assert(((Param *) expression)->paramkind != PARAM_MULTIEXPR && ((Param *) expression)->paramkind != PARAM_SUBLINK); - - /* in EVALUATE_FUNCTIONS mode, leave Param nodes untouched */ - if (coordinatorEvaluationContext->evaluationMode == EVALUATE_FUNCTIONS) - { - return expression; - } - return (Node *) citus_evaluate_expr((Expr *) expression, exprType(expression), exprTypmod(expression), @@ -250,8 +224,7 @@ ShouldEvaluateFunctions(CoordinatorEvaluationContext *evaluationContext) return true; } - return evaluationContext->evaluationMode == EVALUATE_FUNCTIONS_PARAMS || - evaluationContext->evaluationMode == EVALUATE_FUNCTIONS; + return evaluationContext->evaluationMode == EVALUATE_FUNCTIONS_PARAMS; } diff --git a/src/backend/distributed/utils/citus_copyfuncs.c b/src/backend/distributed/utils/citus_copyfuncs.c index 06ecc9bb05a..c8cb29ff3d2 100644 --- a/src/backend/distributed/utils/citus_copyfuncs.c +++ b/src/backend/distributed/utils/citus_copyfuncs.c @@ -105,6 +105,7 @@ copyJobInfo(Job *newnode, Job *from) COPY_SCALAR_FIELD(parametersInJobQueryResolved); COPY_SCALAR_FIELD(colocationId); COPY_NODE_FIELD(savedJobQueryForCaching); + COPY_NODE_FIELD(plannerPartitionKeyValue); } diff --git a/src/include/distributed/citus_clauses.h b/src/include/distributed/citus_clauses.h index b30cc0e03e7..8cfbe18eeb3 100644 --- a/src/include/distributed/citus_clauses.h +++ b/src/include/distributed/citus_clauses.h @@ -28,9 +28,6 @@ typedef enum CoordinatorEvaluationMode /* evaluate only external parameters */ EVALUATE_PARAMS, - /* evaluate functions/expressions but leave Param nodes untouched */ - EVALUATE_FUNCTIONS, - /* evaluate both the functions/expressions and the external paramaters */ EVALUATE_FUNCTIONS_PARAMS } CoordinatorEvaluationMode; @@ -51,7 +48,6 @@ extern void ExecuteCoordinatorEvaluableExpressions(Query *query, PlanState *plan extern Node * PartiallyEvaluateExpression(Node *expression, CoordinatorEvaluationContext * coordinatorEvaluationContext); -extern void ExecuteCoordinatorEvaluableFunctions(Query *query, PlanState *planState); extern bool CitusIsVolatileFunction(Node *node); extern bool CitusIsMutableFunction(Node *node); diff --git a/src/include/distributed/multi_physical_planner.h b/src/include/distributed/multi_physical_planner.h index cd55fe270a0..f899e0b16bd 100644 --- a/src/include/distributed/multi_physical_planner.h +++ b/src/include/distributed/multi_physical_planner.h @@ -163,6 +163,14 @@ typedef struct Job * only; the per-execution copy (currentPlan) gets a NULL here. */ Query *savedJobQueryForCaching; + + /* + * partitionKeyValue as the planner left it, captured alongside + * savedJobQueryForCaching. The cache fast path overwrites + * partitionKeyValue with a per-execution Const, and this is what it gets + * put back to once that Const is about to become dangling. + */ + Const *plannerPartitionKeyValue; } Job; diff --git a/src/include/distributed/multi_router_planner.h b/src/include/distributed/multi_router_planner.h index 434fd966fe1..9441e0e61a0 100644 --- a/src/include/distributed/multi_router_planner.h +++ b/src/include/distributed/multi_router_planner.h @@ -56,6 +56,7 @@ extern List * CreateTaskPlacementListForShardIntervals(List *shardIntervalList, bool hasLocalRelation); extern List * RouterInsertTaskList(Query *query, bool parametersInQueryResolved, DeferredErrorMessage **planningError); +extern List * RelationRowLockListForQuery(Query *query); extern Const * ExtractInsertPartitionKeyValue(Query *query); extern List * TargetShardIntervalsForRestrictInfo(RelationRestrictionContext * restrictionContext, diff --git a/src/include/distributed/prepared_statement_cache.h b/src/include/distributed/prepared_statement_cache.h index ba7373b5aa3..789ecc97bb7 100644 --- a/src/include/distributed/prepared_statement_cache.h +++ b/src/include/distributed/prepared_statement_cache.h @@ -43,17 +43,14 @@ typedef struct PreparedStatementCacheKey /* - * PreparedStatementCacheEntry stores the prepared statement handle on a - * connection, plus metadata needed to re-execute it. + * PreparedStatementCacheEntry names the statement prepared on a connection + * for a given (planId, shardId). */ typedef struct PreparedStatementCacheEntry { PreparedStatementCacheKey key; char stmtName[MAX_STMT_NAME_LENGTH]; - Oid *paramTypes; - int paramCount; - char *parameterizedQueryString; } PreparedStatementCacheEntry; @@ -81,7 +78,8 @@ extern void PreparedStatementCacheDestroy(HTAB **cache_ptr); extern PreparedStatementCacheEntry * PreparedStatementCacheLookup(HTAB *cache, uint64 planId, uint64 shardId); extern PreparedStatementCacheEntry * PreparedStatementCacheInsert(HTAB *cache, uint64 - planId, uint64 shardId); + planId, uint64 shardId, + const char *stmtName); /* planner-side integration (see citus_custom_scan.c) */ extern bool PreparedStatementCacheTryFastPath(struct CitusScanState *scanState, diff --git a/src/test/regress/bin/normalize.sed b/src/test/regress/bin/normalize.sed index 34762fe1538..90b19035320 100644 --- a/src/test/regress/bin/normalize.sed +++ b/src/test/regress/bin/normalize.sed @@ -88,6 +88,9 @@ s/(NOTICE: [a-z]+ cascades to table ".*)_[0-9]{5,}"/\1_xxxxx"/g # connection id s/connectionId: [0-9]+/connectionId: xxxxxxx/g +# cached prepared statement names are numbered per connection +s/issuing __citus_stmt_[0-9]+/issuing __citus_stmt_xxx/g + # Remove trailing whitespace s/ *$//g diff --git a/src/test/regress/citus_tests/config.py b/src/test/regress/citus_tests/config.py index 1d3d122be4f..dd1684f0382 100644 --- a/src/test/regress/citus_tests/config.py +++ b/src/test/regress/citus_tests/config.py @@ -296,6 +296,14 @@ def __init__(self, arguments): } +class CitusPreparedStatementCachingConfig(CitusDefaultClusterConfig): + def __init__(self, arguments): + super().__init__(arguments) + self.new_settings = { + "citus.enable_prepared_statement_caching": True, + } + + class CitusUnusualExecutorConfig(CitusDefaultClusterConfig): def __init__(self, arguments): super().__init__(arguments) diff --git a/src/test/regress/expected/prepared_statement_caching.out b/src/test/regress/expected/prepared_statement_caching.out index 4261cc9155c..65d21e51cb2 100644 --- a/src/test/regress/expected/prepared_statement_caching.out +++ b/src/test/regress/expected/prepared_statement_caching.out @@ -6,6 +6,13 @@ -- CREATE SCHEMA prepared_stmt_caching; SET search_path TO prepared_stmt_caching; +-- Test 8 prints shard names in EXPLAIN output, so pin the shard ids +SET citus.next_shard_id TO 105000; +SET citus.shard_count TO 4; +SET citus.shard_replication_factor TO 1; +-- the regression suite runs with citus.stat_tenants_track = 'ALL', which the +-- cache declines in order to keep per-execution tenant attribution +SET citus.stat_tenants_track TO 'none'; -- Create test tables CREATE TABLE dist_table ( key int PRIMARY KEY, @@ -524,7 +531,10 @@ DEALLOCATE stmt_b; DEALLOCATE stmt_c; -- ============================================================ -- Test 7: Connection loss re-prepare — force worker connection --- close, verify subsequent EXECUTE still works +-- close, verify the statement is re-prepared on the new +-- connection. Results alone cannot show this (see Test 13), +-- so assert the wire protocol: the new connection must be +-- sent the parameterized SQL before a statement name. -- ============================================================ SET citus.enable_prepared_statement_caching = on; PREPARE reconnect_test(int) AS @@ -574,13 +584,30 @@ EXECUTE reconnect_test(1); -- Force all cached connections to be dropped by setting lifetime to 0 SET citus.max_cached_connection_lifetime TO '0s'; --- The next execution should get a new connection, re-prepare, and succeed +-- The lifetime is only applied when the connection is released at end of +-- transaction, so the first execution below still reuses the prepared +-- connection and the second lands on a fresh one. Both use the same key, +-- so a re-prepare can only be caused by the new connection. +SET citus.log_remote_commands TO on; EXECUTE reconnect_test(1); +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx key | value --------------------------------------------------------------------- 1 | 100 (1 row) +EXECUTE reconnect_test(1); +NOTICE: issuing SELECT key, value FROM prepared_stmt_caching.dist_table_105000 dist_table WHERE (key OPERATOR(pg_catalog.=) $1) +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +SET citus.log_remote_commands TO off; EXECUTE reconnect_test(2); key | value --------------------------------------------------------------------- @@ -597,10 +624,670 @@ EXECUTE reconnect_test(3); RESET citus.max_cached_connection_lifetime; DEALLOCATE reconnect_test; -- ============================================================ +-- Test 8: EXPLAIN a cached fast-path statement. EXPLAIN builds the +-- shard query from the saved template rather than from the +-- executor's cache path, and the result is retained on the +-- task, so it must outlive the deparse. +-- ============================================================ +SET citus.enable_prepared_statement_caching = on; +PREPARE explain_select(int) AS + SELECT key, value FROM dist_table WHERE key = $1; +-- reach the generic plan and the cache-hit fast path first +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXPLAIN (COSTS OFF) EXECUTE explain_select(1); + QUERY PLAN +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Index Scan using dist_table_pkey_105000 on dist_table_105000 dist_table + Index Cond: (key = 1) +(7 rows) + +-- the statement must still execute correctly afterwards +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +DEALLOCATE explain_select; +-- INSERT takes the other deparse branch +PREPARE explain_insert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2); +EXECUTE explain_insert(700, 7000); +EXECUTE explain_insert(701, 7010); +EXECUTE explain_insert(702, 7020); +EXECUTE explain_insert(703, 7030); +EXECUTE explain_insert(704, 7040); +EXECUTE explain_insert(705, 7050); +EXECUTE explain_insert(706, 7060); +EXPLAIN (COSTS OFF) EXECUTE explain_insert(707, 7070); + QUERY PLAN +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on dist_table_ts_105005 + -> Result +(7 rows) + +-- EXPLAIN without ANALYZE must not have inserted key 707 +SELECT count(*) AS inserted_rows FROM dist_table_ts WHERE key >= 700 AND key < 800; + inserted_rows +--------------------------------------------------------------------- + 7 +(1 row) + +DEALLOCATE explain_insert; +-- An upsert is what actually needs the insert alias: the DO UPDATE reference to +-- the target table has to resolve against the shard alias, not the shard name. +PREPARE explain_upsert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2) + ON CONFLICT (key) DO UPDATE SET value = dist_table_ts.value + EXCLUDED.value; +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXPLAIN (COSTS OFF) EXECUTE explain_upsert(800, 1); + QUERY PLAN +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on dist_table_ts_105005 citus_table_alias + Conflict Resolution: UPDATE + Conflict Arbiter Indexes: dist_table_ts_pkey_105005 + -> Result +(9 rows) + +-- seven increments, and EXPLAIN must not have added an eighth +SELECT key, value FROM dist_table_ts WHERE key = 800; + key | value +--------------------------------------------------------------------- + 800 | 7 +(1 row) + +DEALLOCATE explain_upsert; +-- ============================================================ +-- Test 9: citus.task_assignment_policy must still apply on the +-- cache-hit fast path, which builds its task directly +-- instead of going through GenerateSingleShardRouterTaskList(). +-- ============================================================ +-- returns 'shardId@port' for the placement the task was assigned to +CREATE OR REPLACE FUNCTION parse_explain_output(in qry text, in table_name text, out r text) +RETURNS SETOF TEXT AS $$ +DECLARE + portOfTheTask text; + shardOfTheTask text; +begin + for r in execute qry loop + IF r LIKE '%port%' THEN + portOfTheTask = substring(r, '([0-9]{1,10})'); + END IF; + + IF r LIKE '%' || table_name || '%' THEN + shardOfTheTask = substring(r, '([0-9]{5,10})'); + END IF; + + end loop; + return QUERY SELECT shardOfTheTask || '@' || portOfTheTask; +end; $$ language plpgsql; +-- round-robin only has something to choose between when shards are replicated +SET citus.shard_replication_factor TO 2; +CREATE TABLE replicated_table (key int PRIMARY KEY, value int); +SELECT create_distributed_table('replicated_table', 'key'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +SET citus.shard_replication_factor TO 1; +INSERT INTO replicated_table SELECT i, i * 10 FROM generate_series(1, 20) i; +SET citus.task_assignment_policy TO 'round-robin'; +SET citus.explain_distributed_queries TO on; +PREPARE round_robin_select(int) AS + SELECT value FROM replicated_table WHERE key = $1; +-- reach the generic plan so that later executions take the fast path +EXECUTE round_robin_select(1); + value +--------------------------------------------------------------------- + 10 +(1 row) + +EXECUTE round_robin_select(1); + value +--------------------------------------------------------------------- + 10 +(1 row) + +EXECUTE round_robin_select(1); + value +--------------------------------------------------------------------- + 10 +(1 row) + +EXECUTE round_robin_select(1); + value +--------------------------------------------------------------------- + 10 +(1 row) + +EXECUTE round_robin_select(1); + value +--------------------------------------------------------------------- + 10 +(1 row) + +EXECUTE round_robin_select(1); + value +--------------------------------------------------------------------- + 10 +(1 row) + +EXECUTE round_robin_select(1); + value +--------------------------------------------------------------------- + 10 +(1 row) + +CREATE TEMPORARY TABLE explain_outputs (value text); +INSERT INTO explain_outputs + SELECT parse_explain_output('EXPLAIN EXECUTE round_robin_select(1)', 'replicated_table'); +INSERT INTO explain_outputs + SELECT parse_explain_output('EXPLAIN EXECUTE round_robin_select(1)', 'replicated_table'); +-- outside a transaction round-robin must alternate placements, so the fast +-- path has to reach both nodes rather than pinning to the first placement +SELECT count(DISTINCT value) FROM explain_outputs; + count +--------------------------------------------------------------------- + 2 +(1 row) + +DROP TABLE explain_outputs; +DEALLOCATE round_robin_select; +RESET citus.task_assignment_policy; +RESET citus.explain_distributed_queries; +-- ============================================================ +-- Test 10: the distribution key parameter need not have the same type +-- as the distribution column. A single-row INSERT records the +-- Param with implicit coercions stripped, so $1 here is int4 +-- while the column is numeric. +-- ============================================================ +CREATE TABLE numeric_dist (key numeric PRIMARY KEY, value int); +SELECT create_distributed_table('numeric_dist', 'key'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +PREPARE numeric_insert(int) AS + INSERT INTO numeric_dist (key, value) VALUES ($1, $1 * 10); +EXECUTE numeric_insert(1); +EXECUTE numeric_insert(2); +EXECUTE numeric_insert(3); +EXECUTE numeric_insert(4); +EXECUTE numeric_insert(5); +EXECUTE numeric_insert(6); +EXECUTE numeric_insert(7); +EXECUTE numeric_insert(8); +EXECUTE numeric_insert(9); +EXECUTE numeric_insert(10); +SELECT count(*) AS total FROM numeric_dist; + total +--------------------------------------------------------------------- + 10 +(1 row) + +-- rows written on the fast path must be reachable by a router lookup, which +-- only holds if the value was hashed as numeric rather than as int4 +SELECT key, value FROM numeric_dist WHERE key = 8; + key | value +--------------------------------------------------------------------- + 8 | 80 +(1 row) + +SELECT key, value FROM numeric_dist WHERE key = 9; + key | value +--------------------------------------------------------------------- + 9 | 90 +(1 row) + +SELECT key, value FROM numeric_dist WHERE key = 10; + key | value +--------------------------------------------------------------------- + 10 | 100 +(1 row) + +DEALLOCATE numeric_insert; +-- ============================================================ +-- Test 11: DML whose expressions must be evaluated on the coordinator +-- cannot be cached. The template is copied before evaluation, +-- so a cached statement would evaluate nextval() on the worker, +-- where the sequence does not exist. +-- ============================================================ +CREATE SEQUENCE coord_eval_seq; +CREATE TABLE coord_eval (key int PRIMARY KEY, seq_value bigint); +SELECT create_distributed_table('coord_eval', 'key'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +PREPARE coord_eval_insert(int) AS + INSERT INTO coord_eval (key, seq_value) VALUES ($1, nextval('coord_eval_seq')); +EXECUTE coord_eval_insert(1); +EXECUTE coord_eval_insert(2); +EXECUTE coord_eval_insert(3); +EXECUTE coord_eval_insert(4); +EXECUTE coord_eval_insert(5); +EXECUTE coord_eval_insert(6); +EXECUTE coord_eval_insert(7); +EXECUTE coord_eval_insert(8); +-- values come from the coordinator's sequence, in execution order +SELECT key, seq_value FROM coord_eval ORDER BY key; + key | seq_value +--------------------------------------------------------------------- + 1 | 1 + 2 | 2 + 3 | 3 + 4 | 4 + 5 | 5 + 6 | 6 + 7 | 7 + 8 | 8 +(8 rows) + +DEALLOCATE coord_eval_insert; +-- UPDATE reaches the deferred-pruning router path, which must carry the same +-- coordinator-evaluation flag +ALTER TABLE coord_eval ADD COLUMN ts timestamptz; +PREPARE coord_eval_update(int) AS + UPDATE coord_eval SET ts = now() WHERE key = $1; +BEGIN; +EXECUTE coord_eval_update(1); +EXECUTE coord_eval_update(2); +EXECUTE coord_eval_update(3); +EXECUTE coord_eval_update(4); +EXECUTE coord_eval_update(5); +EXECUTE coord_eval_update(6); +EXECUTE coord_eval_update(7); +EXECUTE coord_eval_update(8); +-- a coordinator-evaluated now() is one value for the whole transaction; +-- evaluated on the workers it would be one value per node +SELECT count(DISTINCT ts) AS distinct_timestamps FROM coord_eval; + distinct_timestamps +--------------------------------------------------------------------- + 1 +(1 row) + +COMMIT; +DEALLOCATE coord_eval_update; +-- ============================================================ +-- Test 12: a prepared statement may carry parameters the query never +-- uses. Their types still have to be normalized, or the worker +-- cannot infer a type for a parameter absent from the SQL. +-- ============================================================ +CREATE TYPE unused_param_type AS (a int); +PREPARE unused_param(int, unused_param_type) AS + SELECT key, value FROM dist_table WHERE key = $1; +EXECUTE unused_param(1, '(1)'); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE unused_param(1, '(1)'); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE unused_param(1, '(1)'); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE unused_param(1, '(1)'); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE unused_param(1, '(1)'); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE unused_param(1, '(1)'); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE unused_param(1, '(1)'); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +DEALLOCATE unused_param; +-- ============================================================ +-- Test 13: assert the cache is actually engaged. Every other test +-- here checks only query results, which are identical +-- whether or not the feature is active, so observe the +-- wire protocol instead: PQprepare logs the parameterized +-- SQL, PQsendQueryPrepared logs the statement name. +-- ============================================================ +CREATE TABLE observe_cache (key int PRIMARY KEY, value int); +SELECT create_distributed_table('observe_cache', 'key'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +INSERT INTO observe_cache SELECT i, i FROM generate_series(1, 20) i; +PREPARE observe_select(int) AS + SELECT value FROM observe_cache WHERE key = $1; +-- reach the generic plan before enabling logging; the cache is populated on +-- the second use of the generic plan, so this needs one more than the +-- five executions it takes to get there +EXECUTE observe_select(1); + value +--------------------------------------------------------------------- + 1 +(1 row) + +EXECUTE observe_select(1); + value +--------------------------------------------------------------------- + 1 +(1 row) + +EXECUTE observe_select(1); + value +--------------------------------------------------------------------- + 1 +(1 row) + +EXECUTE observe_select(1); + value +--------------------------------------------------------------------- + 1 +(1 row) + +EXECUTE observe_select(1); + value +--------------------------------------------------------------------- + 1 +(1 row) + +EXECUTE observe_select(1); + value +--------------------------------------------------------------------- + 1 +(1 row) + +EXECUTE observe_select(1); + value +--------------------------------------------------------------------- + 1 +(1 row) + +SET citus.log_remote_commands TO on; +-- key 1's shard is already prepared on this connection, so this reuses +-- the statement by name rather than sending SQL +EXECUTE observe_select(1); +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx + value +--------------------------------------------------------------------- + 1 +(1 row) + +-- an unprepared shard prepares first, then executes by name +EXECUTE observe_select(3); +NOTICE: issuing SELECT value FROM prepared_stmt_caching.observe_cache_105021 observe_cache WHERE (key OPERATOR(pg_catalog.=) $1) +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx + value +--------------------------------------------------------------------- + 3 +(1 row) + +EXECUTE observe_select(3); +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx + value +--------------------------------------------------------------------- + 3 +(1 row) + +SET citus.log_remote_commands TO off; +DEALLOCATE observe_select; +-- ============================================================ +-- Test 14: SELECT ... FOR UPDATE is fast-path eligible, so the +-- cached task must carry its row locks. Without them the +-- executor treats it as an ordinary read and, with +-- select_opens_transaction_block off, opens no remote +-- transaction — releasing the row lock at statement end +-- rather than at the end of the enclosing transaction. +-- ============================================================ +SET citus.select_opens_transaction_block TO off; +PREPARE lock_row(int) AS + SELECT value FROM observe_cache WHERE key = $1 FOR UPDATE; +-- reach the generic plan before enabling logging +EXECUTE lock_row(5); + value +--------------------------------------------------------------------- + 5 +(1 row) + +EXECUTE lock_row(5); + value +--------------------------------------------------------------------- + 5 +(1 row) + +EXECUTE lock_row(5); + value +--------------------------------------------------------------------- + 5 +(1 row) + +EXECUTE lock_row(5); + value +--------------------------------------------------------------------- + 5 +(1 row) + +EXECUTE lock_row(5); + value +--------------------------------------------------------------------- + 5 +(1 row) + +EXECUTE lock_row(5); + value +--------------------------------------------------------------------- + 5 +(1 row) + +EXECUTE lock_row(5); + value +--------------------------------------------------------------------- + 5 +(1 row) + +-- the worker must still be given a transaction to hold the row lock in +SET citus.log_remote_commands TO on; +BEGIN; +EXECUTE lock_row(5); +NOTICE: issuing BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;SELECT assign_distributed_transaction_id(xx, xx, 'xxxxxxx'); +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx + value +--------------------------------------------------------------------- + 5 +(1 row) + +COMMIT; +NOTICE: issuing COMMIT +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +SET citus.log_remote_commands TO off; +RESET citus.select_opens_transaction_block; +DEALLOCATE lock_row; +-- ============================================================ +-- Test 15: the same wire-protocol assertion for the remote DML +-- paths. INSERT and UPDATE/DELETE reach the cache through +-- their own eligibility and task-building branches, and a +-- silent fallback in either is invisible in query results. +-- ============================================================ +-- no primary key, so the same shard can be hit repeatedly by INSERT +CREATE TABLE observe_dml (key int, value int); +SELECT create_distributed_table('observe_dml', 'key'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +PREPARE observe_insert(int, int) AS + INSERT INTO observe_dml (key, value) VALUES ($1, $2); +-- reach the generic plan and populate the cache before enabling logging +EXECUTE observe_insert(1, 1); +EXECUTE observe_insert(1, 2); +EXECUTE observe_insert(1, 3); +EXECUTE observe_insert(1, 4); +EXECUTE observe_insert(1, 5); +EXECUTE observe_insert(1, 6); +EXECUTE observe_insert(1, 7); +SET citus.log_remote_commands TO on; +-- key 1's shard is already prepared on this connection, so this reuses +-- the statement by name rather than sending SQL +EXECUTE observe_insert(1, 8); +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +-- an unprepared shard prepares first, then executes by name +EXECUTE observe_insert(3, 9); +NOTICE: issuing INSERT INTO prepared_stmt_caching.observe_dml_105025 (key, value) VALUES ($1, $2) +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +EXECUTE observe_insert(3, 10); +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +SET citus.log_remote_commands TO off; +PREPARE observe_update(int, int) AS + UPDATE observe_dml SET value = $2 WHERE key = $1; +EXECUTE observe_update(1, 11); +EXECUTE observe_update(1, 12); +EXECUTE observe_update(1, 13); +EXECUTE observe_update(1, 14); +EXECUTE observe_update(1, 15); +EXECUTE observe_update(1, 16); +EXECUTE observe_update(1, 17); +SET citus.log_remote_commands TO on; +EXECUTE observe_update(1, 18); +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +SET citus.log_remote_commands TO off; +PREPARE observe_delete(int) AS + DELETE FROM observe_dml WHERE key = $1; +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +SET citus.log_remote_commands TO on; +EXECUTE observe_delete(3); +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +SET citus.log_remote_commands TO off; +SELECT key, value FROM observe_dml ORDER BY key, value; + key | value +--------------------------------------------------------------------- + 1 | 18 + 1 | 18 + 1 | 18 + 1 | 18 + 1 | 18 + 1 | 18 + 1 | 18 + 1 | 18 +(8 rows) + +DEALLOCATE observe_insert; +DEALLOCATE observe_update; +DEALLOCATE observe_delete; +-- ============================================================ -- Cleanup -- ============================================================ SET citus.enable_prepared_statement_caching = off; DROP SCHEMA prepared_stmt_caching CASCADE; -NOTICE: drop cascades to 2 other objects +NOTICE: drop cascades to 10 other objects DETAIL: drop cascades to table dist_table drop cascades to table dist_table_ts +drop cascades to function parse_explain_output(text,text) +drop cascades to table replicated_table +drop cascades to table numeric_dist +drop cascades to sequence coord_eval_seq +drop cascades to table coord_eval +drop cascades to type unused_param_type +drop cascades to table observe_cache +drop cascades to table observe_dml diff --git a/src/test/regress/expected/single_node.out b/src/test/regress/expected/single_node.out index bdd89459875..e0038152510 100644 --- a/src/test/regress/expected/single_node.out +++ b/src/test/regress/expected/single_node.out @@ -1103,7 +1103,7 @@ ROLLBACK; -- explain analyze should work on a single node EXPLAIN (COSTS FALSE, ANALYZE TRUE, TIMING FALSE, SUMMARY FALSE, BUFFERS OFF) SELECT * FROM test; - QUERY PLAN + QUERY PLAN --------------------------------------------------------------------- Custom Scan (Citus Adaptive) (actual rows=5 loops=1) Task Count: 4 @@ -2629,6 +2629,84 @@ NOTICE: executing the command locally: UPDATE single_node.another_schema_table_ (1 row) ROLLBACK; +-- A cached fast-path upsert executed locally is deparsed for local planning, +-- and needs the insert alias so the DO UPDATE reference to the target table +-- resolves against the shard alias. +SET citus.shard_replication_factor TO 1; +SET citus.shard_count TO 4; +SET citus.enable_prepared_statement_caching TO on; +-- the suite runs with citus.stat_tenants_track = 'ALL', which the cache declines +SET citus.stat_tenants_track TO 'none'; +CREATE TABLE single_node.upsert_cache_test (key int PRIMARY KEY, value int); +SELECT create_distributed_table('single_node.upsert_cache_test', 'key'); +NOTICE: executing the command locally: SELECT worker_apply_shard_ddl_command (90830508, 'single_node', 'CREATE TABLE single_node.upsert_cache_test (key integer NOT NULL, value integer) USING heap');SELECT worker_apply_shard_ddl_command (90830508, 'single_node', 'ALTER TABLE single_node.upsert_cache_test OWNER TO postgres');SELECT worker_apply_shard_ddl_command (90830508, 'single_node', 'ALTER TABLE single_node.upsert_cache_test ADD CONSTRAINT upsert_cache_test_pkey PRIMARY KEY (key)') +NOTICE: executing the command locally: SELECT worker_apply_shard_ddl_command (90830509, 'single_node', 'CREATE TABLE single_node.upsert_cache_test (key integer NOT NULL, value integer) USING heap');SELECT worker_apply_shard_ddl_command (90830509, 'single_node', 'ALTER TABLE single_node.upsert_cache_test OWNER TO postgres');SELECT worker_apply_shard_ddl_command (90830509, 'single_node', 'ALTER TABLE single_node.upsert_cache_test ADD CONSTRAINT upsert_cache_test_pkey PRIMARY KEY (key)') +NOTICE: executing the command locally: SELECT worker_apply_shard_ddl_command (90830510, 'single_node', 'CREATE TABLE single_node.upsert_cache_test (key integer NOT NULL, value integer) USING heap');SELECT worker_apply_shard_ddl_command (90830510, 'single_node', 'ALTER TABLE single_node.upsert_cache_test OWNER TO postgres');SELECT worker_apply_shard_ddl_command (90830510, 'single_node', 'ALTER TABLE single_node.upsert_cache_test ADD CONSTRAINT upsert_cache_test_pkey PRIMARY KEY (key)') +NOTICE: executing the command locally: SELECT worker_apply_shard_ddl_command (90830511, 'single_node', 'CREATE TABLE single_node.upsert_cache_test (key integer NOT NULL, value integer) USING heap');SELECT worker_apply_shard_ddl_command (90830511, 'single_node', 'ALTER TABLE single_node.upsert_cache_test OWNER TO postgres');SELECT worker_apply_shard_ddl_command (90830511, 'single_node', 'ALTER TABLE single_node.upsert_cache_test ADD CONSTRAINT upsert_cache_test_pkey PRIMARY KEY (key)') + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +PREPARE local_cached_upsert(int, int) AS + INSERT INTO single_node.upsert_cache_test (key, value) VALUES ($1, $2) + ON CONFLICT (key) DO UPDATE + SET value = upsert_cache_test.value + EXCLUDED.value; +-- with local command logging on, LogLocalCommand() deparses the task first, so +-- the notice below shows the shard SQL that the alias has to appear in +EXECUTE local_cached_upsert(1, 10); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830508 AS citus_table_alias (key, value) VALUES (1, 10) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +EXECUTE local_cached_upsert(1, 10); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830508 AS citus_table_alias (key, value) VALUES (1, 10) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +EXECUTE local_cached_upsert(1, 10); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830508 AS citus_table_alias (key, value) VALUES (1, 10) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +EXECUTE local_cached_upsert(1, 10); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830508 AS citus_table_alias (key, value) VALUES (1, 10) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +EXECUTE local_cached_upsert(1, 10); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830508 AS citus_table_alias (key, value) VALUES (1, 10) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +EXECUTE local_cached_upsert(1, 10); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830508 AS citus_table_alias (key, value) VALUES (1, 10) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +EXECUTE local_cached_upsert(1, 10); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830508 AS citus_table_alias (key, value) VALUES ($1, $2) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +EXECUTE local_cached_upsert(2, 5); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830511 AS citus_table_alias (key, value) VALUES ($1, $2) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +-- key 1 accumulated seven increments, key 2 was inserted once +SELECT key, value FROM single_node.upsert_cache_test ORDER BY key; +NOTICE: executing the command locally: SELECT key, value FROM single_node.upsert_cache_test_90830508 upsert_cache_test WHERE true ORDER BY key +NOTICE: executing the command locally: SELECT key, value FROM single_node.upsert_cache_test_90830509 upsert_cache_test WHERE true ORDER BY key +NOTICE: executing the command locally: SELECT key, value FROM single_node.upsert_cache_test_90830510 upsert_cache_test WHERE true ORDER BY key +NOTICE: executing the command locally: SELECT key, value FROM single_node.upsert_cache_test_90830511 upsert_cache_test WHERE true ORDER BY key + key | value +--------------------------------------------------------------------- + 1 | 70 + 2 | 5 +(2 rows) + +-- with logging off nothing deparses the task up front, so the local executor +-- reaches its own fast-path deparse instead +SET citus.log_local_commands TO off; +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(4, 7); +SELECT key, value FROM single_node.upsert_cache_test WHERE key >= 3 ORDER BY key; + key | value +--------------------------------------------------------------------- + 3 | 700 + 4 | 7 +(2 rows) + +RESET citus.log_local_commands; +DEALLOCATE local_cached_upsert; +DROP TABLE single_node.upsert_cache_test; +RESET citus.enable_prepared_statement_caching; +RESET citus.stat_tenants_track; +RESET citus.shard_count; +RESET citus.shard_replication_factor; -- if the local execution is disabled, we cannot failover to -- local execution and the queries would fail SET citus.enable_local_execution TO false; diff --git a/src/test/regress/sql/prepared_statement_caching.sql b/src/test/regress/sql/prepared_statement_caching.sql index 8880b165508..85bc1cd14d7 100644 --- a/src/test/regress/sql/prepared_statement_caching.sql +++ b/src/test/regress/sql/prepared_statement_caching.sql @@ -8,6 +8,15 @@ CREATE SCHEMA prepared_stmt_caching; SET search_path TO prepared_stmt_caching; +-- Test 8 prints shard names in EXPLAIN output, so pin the shard ids +SET citus.next_shard_id TO 105000; +SET citus.shard_count TO 4; +SET citus.shard_replication_factor TO 1; + +-- the regression suite runs with citus.stat_tenants_track = 'ALL', which the +-- cache declines in order to keep per-execution tenant attribution +SET citus.stat_tenants_track TO 'none'; + -- Create test tables CREATE TABLE dist_table ( key int PRIMARY KEY, @@ -268,7 +277,10 @@ DEALLOCATE stmt_c; -- ============================================================ -- Test 7: Connection loss re-prepare — force worker connection --- close, verify subsequent EXECUTE still works +-- close, verify the statement is re-prepared on the new +-- connection. Results alone cannot show this (see Test 13), +-- so assert the wire protocol: the new connection must be +-- sent the parameterized SQL before a statement name. -- ============================================================ SET citus.enable_prepared_statement_caching = on; @@ -288,8 +300,15 @@ EXECUTE reconnect_test(1); -- Force all cached connections to be dropped by setting lifetime to 0 SET citus.max_cached_connection_lifetime TO '0s'; --- The next execution should get a new connection, re-prepare, and succeed +-- The lifetime is only applied when the connection is released at end of +-- transaction, so the first execution below still reuses the prepared +-- connection and the second lands on a fresh one. Both use the same key, +-- so a re-prepare can only be caused by the new connection. +SET citus.log_remote_commands TO on; +EXECUTE reconnect_test(1); EXECUTE reconnect_test(1); +SET citus.log_remote_commands TO off; + EXECUTE reconnect_test(2); EXECUTE reconnect_test(3); @@ -298,6 +317,390 @@ RESET citus.max_cached_connection_lifetime; DEALLOCATE reconnect_test; +-- ============================================================ +-- Test 8: EXPLAIN a cached fast-path statement. EXPLAIN builds the +-- shard query from the saved template rather than from the +-- executor's cache path, and the result is retained on the +-- task, so it must outlive the deparse. +-- ============================================================ + +SET citus.enable_prepared_statement_caching = on; + +PREPARE explain_select(int) AS + SELECT key, value FROM dist_table WHERE key = $1; + +-- reach the generic plan and the cache-hit fast path first +EXECUTE explain_select(1); +EXECUTE explain_select(1); +EXECUTE explain_select(1); +EXECUTE explain_select(1); +EXECUTE explain_select(1); +EXECUTE explain_select(1); +EXECUTE explain_select(1); + +EXPLAIN (COSTS OFF) EXECUTE explain_select(1); + +-- the statement must still execute correctly afterwards +EXECUTE explain_select(1); + +DEALLOCATE explain_select; + +-- INSERT takes the other deparse branch +PREPARE explain_insert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2); + +EXECUTE explain_insert(700, 7000); +EXECUTE explain_insert(701, 7010); +EXECUTE explain_insert(702, 7020); +EXECUTE explain_insert(703, 7030); +EXECUTE explain_insert(704, 7040); +EXECUTE explain_insert(705, 7050); +EXECUTE explain_insert(706, 7060); + +EXPLAIN (COSTS OFF) EXECUTE explain_insert(707, 7070); + +-- EXPLAIN without ANALYZE must not have inserted key 707 +SELECT count(*) AS inserted_rows FROM dist_table_ts WHERE key >= 700 AND key < 800; + +DEALLOCATE explain_insert; + +-- An upsert is what actually needs the insert alias: the DO UPDATE reference to +-- the target table has to resolve against the shard alias, not the shard name. +PREPARE explain_upsert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2) + ON CONFLICT (key) DO UPDATE SET value = dist_table_ts.value + EXCLUDED.value; + +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); + +EXPLAIN (COSTS OFF) EXECUTE explain_upsert(800, 1); + +-- seven increments, and EXPLAIN must not have added an eighth +SELECT key, value FROM dist_table_ts WHERE key = 800; + +DEALLOCATE explain_upsert; + +-- ============================================================ +-- Test 9: citus.task_assignment_policy must still apply on the +-- cache-hit fast path, which builds its task directly +-- instead of going through GenerateSingleShardRouterTaskList(). +-- ============================================================ + +-- returns 'shardId@port' for the placement the task was assigned to +CREATE OR REPLACE FUNCTION parse_explain_output(in qry text, in table_name text, out r text) +RETURNS SETOF TEXT AS $$ +DECLARE + portOfTheTask text; + shardOfTheTask text; +begin + for r in execute qry loop + IF r LIKE '%port%' THEN + portOfTheTask = substring(r, '([0-9]{1,10})'); + END IF; + + IF r LIKE '%' || table_name || '%' THEN + shardOfTheTask = substring(r, '([0-9]{5,10})'); + END IF; + + end loop; + return QUERY SELECT shardOfTheTask || '@' || portOfTheTask; +end; $$ language plpgsql; + +-- round-robin only has something to choose between when shards are replicated +SET citus.shard_replication_factor TO 2; +CREATE TABLE replicated_table (key int PRIMARY KEY, value int); +SELECT create_distributed_table('replicated_table', 'key'); +SET citus.shard_replication_factor TO 1; + +INSERT INTO replicated_table SELECT i, i * 10 FROM generate_series(1, 20) i; + +SET citus.task_assignment_policy TO 'round-robin'; +SET citus.explain_distributed_queries TO on; + +PREPARE round_robin_select(int) AS + SELECT value FROM replicated_table WHERE key = $1; + +-- reach the generic plan so that later executions take the fast path +EXECUTE round_robin_select(1); +EXECUTE round_robin_select(1); +EXECUTE round_robin_select(1); +EXECUTE round_robin_select(1); +EXECUTE round_robin_select(1); +EXECUTE round_robin_select(1); +EXECUTE round_robin_select(1); + +CREATE TEMPORARY TABLE explain_outputs (value text); + +INSERT INTO explain_outputs + SELECT parse_explain_output('EXPLAIN EXECUTE round_robin_select(1)', 'replicated_table'); +INSERT INTO explain_outputs + SELECT parse_explain_output('EXPLAIN EXECUTE round_robin_select(1)', 'replicated_table'); + +-- outside a transaction round-robin must alternate placements, so the fast +-- path has to reach both nodes rather than pinning to the first placement +SELECT count(DISTINCT value) FROM explain_outputs; + +DROP TABLE explain_outputs; +DEALLOCATE round_robin_select; +RESET citus.task_assignment_policy; +RESET citus.explain_distributed_queries; + +-- ============================================================ +-- Test 10: the distribution key parameter need not have the same type +-- as the distribution column. A single-row INSERT records the +-- Param with implicit coercions stripped, so $1 here is int4 +-- while the column is numeric. +-- ============================================================ + +CREATE TABLE numeric_dist (key numeric PRIMARY KEY, value int); +SELECT create_distributed_table('numeric_dist', 'key'); + +PREPARE numeric_insert(int) AS + INSERT INTO numeric_dist (key, value) VALUES ($1, $1 * 10); + +EXECUTE numeric_insert(1); +EXECUTE numeric_insert(2); +EXECUTE numeric_insert(3); +EXECUTE numeric_insert(4); +EXECUTE numeric_insert(5); +EXECUTE numeric_insert(6); +EXECUTE numeric_insert(7); +EXECUTE numeric_insert(8); +EXECUTE numeric_insert(9); +EXECUTE numeric_insert(10); + +SELECT count(*) AS total FROM numeric_dist; + +-- rows written on the fast path must be reachable by a router lookup, which +-- only holds if the value was hashed as numeric rather than as int4 +SELECT key, value FROM numeric_dist WHERE key = 8; +SELECT key, value FROM numeric_dist WHERE key = 9; +SELECT key, value FROM numeric_dist WHERE key = 10; + +DEALLOCATE numeric_insert; + +-- ============================================================ +-- Test 11: DML whose expressions must be evaluated on the coordinator +-- cannot be cached. The template is copied before evaluation, +-- so a cached statement would evaluate nextval() on the worker, +-- where the sequence does not exist. +-- ============================================================ + +CREATE SEQUENCE coord_eval_seq; +CREATE TABLE coord_eval (key int PRIMARY KEY, seq_value bigint); +SELECT create_distributed_table('coord_eval', 'key'); + +PREPARE coord_eval_insert(int) AS + INSERT INTO coord_eval (key, seq_value) VALUES ($1, nextval('coord_eval_seq')); + +EXECUTE coord_eval_insert(1); +EXECUTE coord_eval_insert(2); +EXECUTE coord_eval_insert(3); +EXECUTE coord_eval_insert(4); +EXECUTE coord_eval_insert(5); +EXECUTE coord_eval_insert(6); +EXECUTE coord_eval_insert(7); +EXECUTE coord_eval_insert(8); + +-- values come from the coordinator's sequence, in execution order +SELECT key, seq_value FROM coord_eval ORDER BY key; + +DEALLOCATE coord_eval_insert; + +-- UPDATE reaches the deferred-pruning router path, which must carry the same +-- coordinator-evaluation flag +ALTER TABLE coord_eval ADD COLUMN ts timestamptz; + +PREPARE coord_eval_update(int) AS + UPDATE coord_eval SET ts = now() WHERE key = $1; + +BEGIN; +EXECUTE coord_eval_update(1); +EXECUTE coord_eval_update(2); +EXECUTE coord_eval_update(3); +EXECUTE coord_eval_update(4); +EXECUTE coord_eval_update(5); +EXECUTE coord_eval_update(6); +EXECUTE coord_eval_update(7); +EXECUTE coord_eval_update(8); + +-- a coordinator-evaluated now() is one value for the whole transaction; +-- evaluated on the workers it would be one value per node +SELECT count(DISTINCT ts) AS distinct_timestamps FROM coord_eval; +COMMIT; + +DEALLOCATE coord_eval_update; + +-- ============================================================ +-- Test 12: a prepared statement may carry parameters the query never +-- uses. Their types still have to be normalized, or the worker +-- cannot infer a type for a parameter absent from the SQL. +-- ============================================================ + +CREATE TYPE unused_param_type AS (a int); + +PREPARE unused_param(int, unused_param_type) AS + SELECT key, value FROM dist_table WHERE key = $1; + +EXECUTE unused_param(1, '(1)'); +EXECUTE unused_param(1, '(1)'); +EXECUTE unused_param(1, '(1)'); +EXECUTE unused_param(1, '(1)'); +EXECUTE unused_param(1, '(1)'); +EXECUTE unused_param(1, '(1)'); +EXECUTE unused_param(1, '(1)'); + +DEALLOCATE unused_param; + +-- ============================================================ +-- Test 13: assert the cache is actually engaged. Every other test +-- here checks only query results, which are identical +-- whether or not the feature is active, so observe the +-- wire protocol instead: PQprepare logs the parameterized +-- SQL, PQsendQueryPrepared logs the statement name. +-- ============================================================ + +CREATE TABLE observe_cache (key int PRIMARY KEY, value int); +SELECT create_distributed_table('observe_cache', 'key'); +INSERT INTO observe_cache SELECT i, i FROM generate_series(1, 20) i; + +PREPARE observe_select(int) AS + SELECT value FROM observe_cache WHERE key = $1; + +-- reach the generic plan before enabling logging; the cache is populated on +-- the second use of the generic plan, so this needs one more than the +-- five executions it takes to get there +EXECUTE observe_select(1); +EXECUTE observe_select(1); +EXECUTE observe_select(1); +EXECUTE observe_select(1); +EXECUTE observe_select(1); +EXECUTE observe_select(1); +EXECUTE observe_select(1); + +SET citus.log_remote_commands TO on; + +-- key 1's shard is already prepared on this connection, so this reuses +-- the statement by name rather than sending SQL +EXECUTE observe_select(1); + +-- an unprepared shard prepares first, then executes by name +EXECUTE observe_select(3); +EXECUTE observe_select(3); + +SET citus.log_remote_commands TO off; + +DEALLOCATE observe_select; + +-- ============================================================ +-- Test 14: SELECT ... FOR UPDATE is fast-path eligible, so the +-- cached task must carry its row locks. Without them the +-- executor treats it as an ordinary read and, with +-- select_opens_transaction_block off, opens no remote +-- transaction — releasing the row lock at statement end +-- rather than at the end of the enclosing transaction. +-- ============================================================ + +SET citus.select_opens_transaction_block TO off; + +PREPARE lock_row(int) AS + SELECT value FROM observe_cache WHERE key = $1 FOR UPDATE; + +-- reach the generic plan before enabling logging +EXECUTE lock_row(5); +EXECUTE lock_row(5); +EXECUTE lock_row(5); +EXECUTE lock_row(5); +EXECUTE lock_row(5); +EXECUTE lock_row(5); +EXECUTE lock_row(5); + +-- the worker must still be given a transaction to hold the row lock in +SET citus.log_remote_commands TO on; +BEGIN; +EXECUTE lock_row(5); +COMMIT; +SET citus.log_remote_commands TO off; + +RESET citus.select_opens_transaction_block; +DEALLOCATE lock_row; + +-- ============================================================ +-- Test 15: the same wire-protocol assertion for the remote DML +-- paths. INSERT and UPDATE/DELETE reach the cache through +-- their own eligibility and task-building branches, and a +-- silent fallback in either is invisible in query results. +-- ============================================================ + +-- no primary key, so the same shard can be hit repeatedly by INSERT +CREATE TABLE observe_dml (key int, value int); +SELECT create_distributed_table('observe_dml', 'key'); + +PREPARE observe_insert(int, int) AS + INSERT INTO observe_dml (key, value) VALUES ($1, $2); + +-- reach the generic plan and populate the cache before enabling logging +EXECUTE observe_insert(1, 1); +EXECUTE observe_insert(1, 2); +EXECUTE observe_insert(1, 3); +EXECUTE observe_insert(1, 4); +EXECUTE observe_insert(1, 5); +EXECUTE observe_insert(1, 6); +EXECUTE observe_insert(1, 7); + +SET citus.log_remote_commands TO on; + +-- key 1's shard is already prepared on this connection, so this reuses +-- the statement by name rather than sending SQL +EXECUTE observe_insert(1, 8); + +-- an unprepared shard prepares first, then executes by name +EXECUTE observe_insert(3, 9); +EXECUTE observe_insert(3, 10); + +SET citus.log_remote_commands TO off; + +PREPARE observe_update(int, int) AS + UPDATE observe_dml SET value = $2 WHERE key = $1; + +EXECUTE observe_update(1, 11); +EXECUTE observe_update(1, 12); +EXECUTE observe_update(1, 13); +EXECUTE observe_update(1, 14); +EXECUTE observe_update(1, 15); +EXECUTE observe_update(1, 16); +EXECUTE observe_update(1, 17); + +SET citus.log_remote_commands TO on; +EXECUTE observe_update(1, 18); +SET citus.log_remote_commands TO off; + +PREPARE observe_delete(int) AS + DELETE FROM observe_dml WHERE key = $1; + +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); + +SET citus.log_remote_commands TO on; +EXECUTE observe_delete(3); +SET citus.log_remote_commands TO off; + +SELECT key, value FROM observe_dml ORDER BY key, value; + +DEALLOCATE observe_insert; +DEALLOCATE observe_update; +DEALLOCATE observe_delete; + -- ============================================================ -- Cleanup -- ============================================================ diff --git a/src/test/regress/sql/single_node.sql b/src/test/regress/sql/single_node.sql index 079e59a93d8..93a24857fa3 100644 --- a/src/test/regress/sql/single_node.sql +++ b/src/test/regress/sql/single_node.sql @@ -1333,6 +1333,61 @@ BEGIN; SELECT coordinated_transaction_should_use_2PC(); ROLLBACK; +-- A cached fast-path upsert executed locally is deparsed for local planning, +-- and needs the insert alias so the DO UPDATE reference to the target table +-- resolves against the shard alias. +SET citus.shard_replication_factor TO 1; +SET citus.shard_count TO 4; +SET citus.enable_prepared_statement_caching TO on; + +-- the suite runs with citus.stat_tenants_track = 'ALL', which the cache declines +SET citus.stat_tenants_track TO 'none'; + +CREATE TABLE single_node.upsert_cache_test (key int PRIMARY KEY, value int); +SELECT create_distributed_table('single_node.upsert_cache_test', 'key'); + +PREPARE local_cached_upsert(int, int) AS + INSERT INTO single_node.upsert_cache_test (key, value) VALUES ($1, $2) + ON CONFLICT (key) DO UPDATE + SET value = upsert_cache_test.value + EXCLUDED.value; + +-- with local command logging on, LogLocalCommand() deparses the task first, so +-- the notice below shows the shard SQL that the alias has to appear in +EXECUTE local_cached_upsert(1, 10); +EXECUTE local_cached_upsert(1, 10); +EXECUTE local_cached_upsert(1, 10); +EXECUTE local_cached_upsert(1, 10); +EXECUTE local_cached_upsert(1, 10); +EXECUTE local_cached_upsert(1, 10); +EXECUTE local_cached_upsert(1, 10); +EXECUTE local_cached_upsert(2, 5); + +-- key 1 accumulated seven increments, key 2 was inserted once +SELECT key, value FROM single_node.upsert_cache_test ORDER BY key; + +-- with logging off nothing deparses the task up front, so the local executor +-- reaches its own fast-path deparse instead +SET citus.log_local_commands TO off; + +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(4, 7); + +SELECT key, value FROM single_node.upsert_cache_test WHERE key >= 3 ORDER BY key; + +RESET citus.log_local_commands; +DEALLOCATE local_cached_upsert; +DROP TABLE single_node.upsert_cache_test; +RESET citus.enable_prepared_statement_caching; +RESET citus.stat_tenants_track; +RESET citus.shard_count; +RESET citus.shard_replication_factor; + -- if the local execution is disabled, we cannot failover to -- local execution and the queries would fail SET citus.enable_local_execution TO false;