From 416354866d1379793da8d95e754b676c933228e9 Mon Sep 17 00:00:00 2001 From: Michael Paquier Date: Fri, 23 Oct 2026 14:05:23 +0900 Subject: [PATCH v3 1/2] 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 on the extension state of 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. Author: Michael Paquier Discussion: https://postgr.es/m/ --- src/include/executor/execdesc.h | 29 ++++++++++ src/backend/tcop/pquery.c | 95 ++++++++++++++++++++++++++++++++- 2 files changed, 123 insertions(+), 1 deletion(-) diff --git a/src/include/executor/execdesc.h b/src/include/executor/execdesc.h index 37c2576e4bc9..b83efbd3315a 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,29 @@ extern QueryDesc *CreateQueryDesc(PlannedStmt *plannedstmt, extern void FreeQueryDesc(QueryDesc *qdesc); +/* + * Loadable modules can attach private state to a QueryDesc. Get an integer + * ID for the module with GetQueryDescExtensionId(), which returns the same ID + * for the same name within a backend, and use it with the functions below. + * The state lives as long as the QueryDesc, and is not touched by core + * otherwise, so the module is responsible for what it points to. + */ +extern int GetQueryDescExtensionId(const char *extension_name); +extern void SetQueryDescExtensionState(QueryDesc *qdesc, int extension_id, + void *opaque); + +/* + * Get extension-specific state from a QueryDesc, or NULL if none was set. + */ +static inline 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]; +} + #endif /* EXECDESC_H */ diff --git a/src/backend/tcop/pquery.c b/src/backend/tcop/pquery.c index fa2ef20f178a..a94550544cec 100644 --- a/src/backend/tcop/pquery.c +++ b/src/backend/tcop/pquery.c @@ -23,6 +23,7 @@ #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" @@ -97,6 +98,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 +118,98 @@ FreeQueryDesc(QueryDesc *qdesc) UnregisterSnapshot(qdesc->snapshot); UnregisterSnapshot(qdesc->crosscheck_snapshot); - /* Only the QueryDesc itself need be freed */ + /* + * Free the array of extension state pointers. What these point to is + * owned by the extensions. + */ + if (qdesc->extension_state != NULL) + pfree(qdesc->extension_state); + + /* Free the QueryDesc itself */ pfree(qdesc); } +static const char **QueryDescExtensionNameArray = NULL; +static int QueryDescExtensionNamesAssigned = 0; +static int QueryDescExtensionNamesAllocated = 0; + +/* + * 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++; +} + +/* + * 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 -- 2.56.0