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 <michael@paquier.xyz>
Reviewed-by: Sami Imseih <samimseih.pg@gmail.com>
Discussion: https://postgr.es/m/CAEcQ1bZrgwKsauC=03otebkLk77Zc_hYNjhTP3OUdpXs++u+BA@mail.gmail.com

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

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  */