From 793eafad8e988b6754c9d89e0ea14b64b07eef81 Mon Sep 17 00:00:00 2001
From: Atsushi Torikoshi<torikoshia@oss.nttdata.com>
Date: Fri, 31 Jul 2020 17:52:14 +0900
Subject: [PATCH] Previously the number of custom and generic plans are
 recoreded only in pg_prepared_statements, meaning we could only track them
 regarding current session. This patch records them in pg_stat_statements and
 it enables to track them regarding all sessions of the PostgreSQL instance.

---
 .../pg_stat_statements--1.6--1.7.sql          |  5 ++-
 .../pg_stat_statements--1.7--1.8.sql          |  1 +
 .../pg_stat_statements/pg_stat_statements.c   | 44 +++++++++++++++----
 src/backend/utils/cache/plancache.c           |  2 +
 src/include/utils/plancache.h                 |  1 +
 5 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql b/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql
index 6fc3fed4c9..fd7aa05c92 100644
--- a/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql
+++ b/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql
@@ -12,11 +12,12 @@ DROP FUNCTION pg_stat_statements_reset();
 /* Now redefine */
 CREATE FUNCTION pg_stat_statements_reset(IN userid Oid DEFAULT 0,
 	IN dbid Oid DEFAULT 0,
-	IN queryid bigint DEFAULT 0
+	IN queryid bigint DEFAULT 0,
+	IN generic_plan int DEFAULT -1
 )
 RETURNS void
 AS 'MODULE_PATHNAME', 'pg_stat_statements_reset_1_7'
 LANGUAGE C STRICT PARALLEL SAFE;
 
 -- Don't want this to be available to non-superusers.
-REVOKE ALL ON FUNCTION pg_stat_statements_reset(Oid, Oid, bigint) FROM PUBLIC;
+REVOKE ALL ON FUNCTION pg_stat_statements_reset(Oid, Oid, bigint, int) FROM PUBLIC;
diff --git a/contrib/pg_stat_statements/pg_stat_statements--1.7--1.8.sql b/contrib/pg_stat_statements/pg_stat_statements--1.7--1.8.sql
index 0f63f08f7e..0d7c4e7343 100644
--- a/contrib/pg_stat_statements/pg_stat_statements--1.7--1.8.sql
+++ b/contrib/pg_stat_statements/pg_stat_statements--1.7--1.8.sql
@@ -15,6 +15,7 @@ DROP FUNCTION pg_stat_statements(boolean);
 CREATE FUNCTION pg_stat_statements(IN showtext boolean,
     OUT userid oid,
     OUT dbid oid,
+    OUT generic_plan bool,
     OUT queryid bigint,
     OUT query text,
     OUT plans int8,
diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index 6b91c62c31..14c580a95e 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -78,9 +78,11 @@
 #include "storage/ipc.h"
 #include "storage/spin.h"
 #include "tcop/utility.h"
+#include "tcop/pquery.h"
 #include "utils/acl.h"
 #include "utils/builtins.h"
 #include "utils/memutils.h"
+#include "utils/plancache.h"
 
 PG_MODULE_MAGIC;
 
@@ -156,6 +158,7 @@ typedef struct pgssHashKey
 	Oid			userid;			/* user OID */
 	Oid			dbid;			/* database OID */
 	uint64		queryid;		/* query identifier */
+	bool			is_generic_plan;
 } pgssHashKey;
 
 /*
@@ -266,6 +269,9 @@ static int	exec_nested_level = 0;
 /* Current nesting depth of planner calls */
 static int	plan_nested_level = 0;
 
+/* Current plan type */
+static bool	is_generic_plan = false;
+
 /* Saved hook values in case of unload */
 static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
 static post_parse_analyze_hook_type prev_post_parse_analyze_hook = NULL;
@@ -367,7 +373,7 @@ static char *qtext_fetch(Size query_offset, int query_len,
 						 char *buffer, Size buffer_size);
 static bool need_gc_qtexts(void);
 static void gc_qtexts(void);
-static void entry_reset(Oid userid, Oid dbid, uint64 queryid);
+static void entry_reset(Oid userid, Oid dbid, uint64 queryid, int generic_plan);
 static void AppendJumble(pgssJumbleState *jstate,
 						 const unsigned char *item, Size size);
 static void JumbleQuery(pgssJumbleState *jstate, Query *query);
@@ -1013,6 +1019,16 @@ pgss_ExecutorStart(QueryDesc *queryDesc, int eflags)
 	 */
 	if (pgss_enabled(exec_nested_level) && queryDesc->plannedstmt->queryId != UINT64CONST(0))
 	{
+		/*
+		 * ActivePortal is not available at ExecutorEnd. We preserve
+		 * the plan type from ActivePortal here.
+		 */
+		Assert(ActivePortal);
+		if(ActivePortal->cplan != NULL && ActivePortal->cplan->is_generic)
+			is_generic_plan = true;
+		else
+			is_generic_plan = false;
+
 		/*
 		 * Set up to track total elapsed time in ExecutorRun.  Make sure the
 		 * space is allocated in the per-query context so it will go away at
@@ -1100,6 +1116,8 @@ pgss_ExecutorEnd(QueryDesc *queryDesc)
 				   &queryDesc->totaltime->walusage,
 				   NULL);
 	}
+	/* reset the plan type. */
+	is_generic_plan = false;
 
 	if (prev_ExecutorEnd)
 		prev_ExecutorEnd(queryDesc);
@@ -1307,9 +1325,11 @@ pgss_store(const char *query, uint64 queryId,
 	}
 
 	/* Set up key for hashtable search */
+	memset(&key, 0, sizeof(key));
 	key.userid = GetUserId();
 	key.dbid = MyDatabaseId;
 	key.queryid = queryId;
+	key.is_generic_plan = is_generic_plan;
 
 	/* Lookup the hash table entry with shared lock. */
 	LWLockAcquire(pgss->lock, LW_SHARED);
@@ -1463,12 +1483,14 @@ pg_stat_statements_reset_1_7(PG_FUNCTION_ARGS)
 	Oid			userid;
 	Oid			dbid;
 	uint64		queryid;
+	int		generic_plan;
 
 	userid = PG_GETARG_OID(0);
 	dbid = PG_GETARG_OID(1);
 	queryid = (uint64) PG_GETARG_INT64(2);
+	generic_plan = (uint32) PG_GETARG_INT32(3);
 
-	entry_reset(userid, dbid, queryid);
+	entry_reset(userid, dbid, queryid, generic_plan);
 
 	PG_RETURN_VOID();
 }
@@ -1479,7 +1501,7 @@ pg_stat_statements_reset_1_7(PG_FUNCTION_ARGS)
 Datum
 pg_stat_statements_reset(PG_FUNCTION_ARGS)
 {
-	entry_reset(0, 0, 0);
+	entry_reset(0, 0, 0, -1);
 
 	PG_RETURN_VOID();
 }
@@ -1489,8 +1511,8 @@ pg_stat_statements_reset(PG_FUNCTION_ARGS)
 #define PG_STAT_STATEMENTS_COLS_V1_1	18
 #define PG_STAT_STATEMENTS_COLS_V1_2	19
 #define PG_STAT_STATEMENTS_COLS_V1_3	23
-#define PG_STAT_STATEMENTS_COLS_V1_8	32
-#define PG_STAT_STATEMENTS_COLS			32	/* maximum of above */
+#define PG_STAT_STATEMENTS_COLS_V1_8	33
+#define PG_STAT_STATEMENTS_COLS			33	/* maximum of above */
 
 /*
  * Retrieve statement statistics.
@@ -1712,6 +1734,7 @@ pg_stat_statements_internal(FunctionCallInfo fcinfo,
 
 		values[i++] = ObjectIdGetDatum(entry->key.userid);
 		values[i++] = ObjectIdGetDatum(entry->key.dbid);
+		values[i++] = BoolGetDatum(entry->key.is_generic_plan);
 
 		if (is_allowed_role || entry->key.userid == userid)
 		{
@@ -2451,7 +2474,7 @@ gc_fail:
  * Release entries corresponding to parameters passed.
  */
 static void
-entry_reset(Oid userid, Oid dbid, uint64 queryid)
+entry_reset(Oid userid, Oid dbid, uint64 queryid, int generic_plan)
 {
 	HASH_SEQ_STATUS hash_seq;
 	pgssEntry  *entry;
@@ -2468,19 +2491,21 @@ entry_reset(Oid userid, Oid dbid, uint64 queryid)
 	LWLockAcquire(pgss->lock, LW_EXCLUSIVE);
 	num_entries = hash_get_num_entries(pgss_hash);
 
-	if (userid != 0 && dbid != 0 && queryid != UINT64CONST(0))
+	if (userid != 0 && dbid != 0 && queryid != UINT64CONST(0) && generic_plan != -1)
 	{
 		/* If all the parameters are available, use the fast path. */
+		memset(&key, 0, sizeof(key));
 		key.userid = userid;
 		key.dbid = dbid;
 		key.queryid = queryid;
+		key.is_generic_plan = generic_plan;
 
 		/* Remove the key if exists */
 		entry = (pgssEntry *) hash_search(pgss_hash, &key, HASH_REMOVE, NULL);
 		if (entry)				/* found */
 			num_remove++;
 	}
-	else if (userid != 0 || dbid != 0 || queryid != UINT64CONST(0))
+	else if (userid != 0 || dbid != 0 || queryid != UINT64CONST(0) || generic_plan != -1)
 	{
 		/* Remove entries corresponding to valid parameters. */
 		hash_seq_init(&hash_seq, pgss_hash);
@@ -2488,7 +2513,8 @@ entry_reset(Oid userid, Oid dbid, uint64 queryid)
 		{
 			if ((!userid || entry->key.userid == userid) &&
 				(!dbid || entry->key.dbid == dbid) &&
-				(!queryid || entry->key.queryid == queryid))
+				(!queryid || entry->key.queryid == queryid) &&
+				(generic_plan == -1 || entry->key.is_generic_plan == generic_plan))
 			{
 				hash_search(pgss_hash, &entry->key, HASH_REMOVE, NULL);
 				num_remove++;
diff --git a/src/backend/utils/cache/plancache.c b/src/backend/utils/cache/plancache.c
index 50d6ad28b4..62e0539370 100644
--- a/src/backend/utils/cache/plancache.c
+++ b/src/backend/utils/cache/plancache.c
@@ -1219,10 +1219,12 @@ GetCachedPlan(CachedPlanSource *plansource, ParamListInfo boundParams,
 		plansource->total_custom_cost += cached_plan_cost(plan, true);
 
 		plansource->num_custom_plans++;
+		plan->is_generic = false;
 	}
 	else
 	{
 		plansource->num_generic_plans++;
+		plan->is_generic = true;
 	}
 
 	Assert(plan != NULL);
diff --git a/src/include/utils/plancache.h b/src/include/utils/plancache.h
index 4901568553..e4f076fdb3 100644
--- a/src/include/utils/plancache.h
+++ b/src/include/utils/plancache.h
@@ -158,6 +158,7 @@ typedef struct CachedPlan
 	int			generation;		/* parent's generation number for this plan */
 	int			refcount;		/* count of live references to this struct */
 	MemoryContext context;		/* context containing this CachedPlan */
+	bool		is_generic;	/* is this plan generic or not */
 } CachedPlan;
 
 /*
-- 
2.18.1

