From e5d5d1ad1e3aa495609e19fa716dcea6a39db1bb Mon Sep 17 00:00:00 2001 From: Michael Paquier Date: Tue, 27 Oct 2026 15:05:59 +0900 Subject: [PATCH] Add custom extension states to QueryDesc Extensions sometimes need to carry private data from one executor hook to another for a given query. ExecutorEnd() can run in a different context than ExecutorStart(), so backend-wide variables cannot be used for that. This commit adds a way for loadable modules to attach private state to a QueryDesc, modeled in the same fashion as what is already done for PlannerGlobal, PlannerInfo, RelOptInfo and ExplainState. GetQueryDescExtensionId() maps an extension name to a backend-local ID. SetQueryDescExtensionState() and GetQueryDescExtensionState() store and fetch an opaque pointer for that ID. The array of pointers is allocated in the memory context of the QueryDesc, and is freed by FreeQueryDesc(). What the pointers refer to is owned by the extensions. There are no users of this facility yet. A follow-up commit will make use of it. Some potential use cases have been identified for pg_stat_statements and auto_explain, where both need to be able to co-exist in peace. Extensions living outside of core may also find a use for it. Author: Michael Paquier Reviewed-by: Sami Imseih Discussion: https://postgr.es/m/CAEcQ1bZrgwKsauC=03otebkLk77Zc_hYNjhTP3OUdpXs++u+BA@mail.gmail.com --- src/backend/tcop/pquery.c | 107 +++++++++++++++++++++++++++++++- src/include/executor/execdesc.h | 9 +++ 2 files changed, 114 insertions(+), 2 deletions(-) diff --git a/src/backend/tcop/pquery.c b/src/backend/tcop/pquery.c index fa2ef20f178..98205aafb05 100644 --- a/src/backend/tcop/pquery.c +++ b/src/backend/tcop/pquery.c @@ -23,18 +23,22 @@ #include "executor/tstoreReceiver.h" #include "miscadmin.h" #include "pg_trace.h" +#include "port/pg_bitutils.h" #include "tcop/pquery.h" #include "tcop/utility.h" #include "utils/memutils.h" #include "utils/snapmgr.h" - /* * ActivePortal is the currently executing Portal (the most closely nested, * if there are several). */ Portal ActivePortal = NULL; +/* Private data for QueryDesc.extension_state. */ +static const char **QueryDescExtensionNameArray = NULL; +static int QueryDescExtensionNamesAssigned = 0; +static int QueryDescExtensionNamesAllocated = 0; static void ProcessQuery(PlannedStmt *plan, const char *sourceText, @@ -97,6 +101,10 @@ CreateQueryDesc(PlannedStmt *plannedstmt, /* not yet executed */ qd->already_executed = false; + /* no extension state yet */ + qd->extension_state = NULL; + qd->extension_state_allocated = 0; + return qd; } @@ -113,10 +121,105 @@ FreeQueryDesc(QueryDesc *qdesc) UnregisterSnapshot(qdesc->snapshot); UnregisterSnapshot(qdesc->crosscheck_snapshot); - /* Only the QueryDesc itself need be freed */ + /* Free extension state pointers */ + if (qdesc->extension_state != NULL) + pfree(qdesc->extension_state); + + /* Free the QueryDesc itself */ pfree(qdesc); } +/* + * Map the name of an extension to an integer ID, to be used with + * SetQueryDescExtensionState() and GetQueryDescExtensionState(). + * + * Within the lifetime of a particular backend, the same name will be mapped + * to the same ID every time. IDs are not stable across backends. + */ +int +GetQueryDescExtensionId(const char *extension_name) +{ + /* Search for an existing extension by this name; if found, return ID. */ + for (int i = 0; i < QueryDescExtensionNamesAssigned; ++i) + if (strcmp(QueryDescExtensionNameArray[i], extension_name) == 0) + return i; + + /* If there is no array yet, create one. */ + if (QueryDescExtensionNameArray == NULL) + { + QueryDescExtensionNamesAllocated = 16; + QueryDescExtensionNameArray = (const char **) + MemoryContextAlloc(TopMemoryContext, + QueryDescExtensionNamesAllocated + * sizeof(char *)); + } + + /* If there's an array but it's currently full, expand it. */ + if (QueryDescExtensionNamesAssigned >= QueryDescExtensionNamesAllocated) + { + int i = pg_nextpower2_32(QueryDescExtensionNamesAssigned + 1); + + QueryDescExtensionNameArray = + repalloc_array(QueryDescExtensionNameArray, const char *, i); + QueryDescExtensionNamesAllocated = i; + } + + /* Assign and return new ID. The name is copied, to be on the safe side. */ + QueryDescExtensionNameArray[QueryDescExtensionNamesAssigned] = + MemoryContextStrdup(TopMemoryContext, extension_name); + return QueryDescExtensionNamesAssigned++; +} + +/* + * Get extension-specific state from a QueryDesc, or NULL if none was set. + */ +void * +GetQueryDescExtensionState(QueryDesc *qdesc, int extension_id) +{ + Assert(extension_id >= 0); + + if (extension_id >= qdesc->extension_state_allocated) + return NULL; + + return qdesc->extension_state[extension_id]; +} + +/* + * Store extension-specific state into a QueryDesc. + * + * The array of pointers is allocated in the memory context of the QueryDesc, + * and is freed by FreeQueryDesc(). + */ +void +SetQueryDescExtensionState(QueryDesc *qdesc, int extension_id, void *opaque) +{ + Assert(extension_id >= 0); + + /* If there is no array yet, create one. */ + if (qdesc->extension_state == NULL) + { + qdesc->extension_state_allocated = + Max(4, pg_nextpower2_32(extension_id + 1)); + qdesc->extension_state = + MemoryContextAllocZero(GetMemoryChunkContext(qdesc), + qdesc->extension_state_allocated + * sizeof(void *)); + } + + /* If there's an array but it's currently full, expand it. */ + if (extension_id >= qdesc->extension_state_allocated) + { + int i = pg_nextpower2_32(extension_id + 1); + + qdesc->extension_state = + repalloc0_array(qdesc->extension_state, void *, + qdesc->extension_state_allocated, i); + qdesc->extension_state_allocated = i; + } + + qdesc->extension_state[extension_id] = opaque; +} + /* * ProcessQuery diff --git a/src/include/executor/execdesc.h b/src/include/executor/execdesc.h index 37c2576e4bc..81ea65b9332 100644 --- a/src/include/executor/execdesc.h +++ b/src/include/executor/execdesc.h @@ -55,6 +55,10 @@ typedef struct QueryDesc /* This field is allocated by ExecutorStart if needed */ struct Instrumentation *query_instr; /* query level instrumentation */ + + /* Private state of extensions; see GetQueryDescExtensionId() */ + void **extension_state; + int extension_state_allocated; } QueryDesc; /* in pquery.c */ @@ -69,4 +73,9 @@ extern QueryDesc *CreateQueryDesc(PlannedStmt *plannedstmt, extern void FreeQueryDesc(QueryDesc *qdesc); +extern int GetQueryDescExtensionId(const char *extension_name); +extern void *GetQueryDescExtensionState(QueryDesc *qdesc, int extension_id); +extern void SetQueryDescExtensionState(QueryDesc *qdesc, int extension_id, + void *opaque); + #endif /* EXECDESC_H */