PG-624: pg_stat_monitor: Possible server crash when running pgbench with pg_stat_monitor loaded (#396)

PG-624: pg_stat_monitor: Possible server crash when running pgbench
with pg_stat_monitor loaded

It appears that this issue was being caused by improper handling of
dynamic number of buckets. This commit resolves the issue.

Also, as part of a larger cleanup, memory context has been moved to
local space from shared storage. Also, some unwanted
and no-longer-needed variables are removed.

Co-authored-by: Muhammad Usama <muhammad.usama@percona.com>

hash_query.c

@@ -23,6 +23,9 @@ static PGSM_HASH_TABLE_HANDLE pgsm_create_bucket_hash(pgsmSharedState * pgsm, ds
 static Size pgsm_get_shared_area_size(void);
 static void InitializeSharedState(pgsmSharedState * pgsm);
 
+#define PGSM_BUCKET_INFO_SIZE	(sizeof(TimestampTz) * pgsm_max_buckets)
+#define PGSM_SHARED_STATE_SIZE	(sizeof(pgsmSharedState) + PGSM_BUCKET_INFO_SIZE)
+
 #if USE_DYNAMIC_HASH
 /* parameter for the shared hash */
 static dshash_parameters dsh_params = {
@@ -56,7 +59,7 @@ pgsm_query_area_size(void)
 Size
 pgsm_ShmemSize(void)
 {
-	Size		sz = MAXALIGN(sizeof(pgsmSharedState));
+	Size		sz = MAXALIGN(PGSM_SHARED_STATE_SIZE);
 
 	sz = add_size(sz, MAX_QUERY_BUF);
 #if USE_DYNAMIC_HASH
@@ -114,7 +117,7 @@ pgsm_startup(void)
 		SpinLockInit(&pgsm->mutex);
 		InitializeSharedState(pgsm);
 		/* the allocation of pgsmSharedState itself */
-		p += MAXALIGN(sizeof(pgsmSharedState));
+		p += MAXALIGN(PGSM_SHARED_STATE_SIZE);
 		pgsm->raw_dsa_area = p;
 		dsa = dsa_create_in_place(pgsm->raw_dsa_area,
 								  pgsm_query_area_size(),
@@ -138,6 +141,10 @@ pgsm_startup(void)
 		 * references.
 		 */
 		dsa_detach(dsa);
+
+	    pgsmStateLocal.pgsm_mem_cxt = AllocSetContextCreate(TopMemoryContext,
+                                                            "pg_stat_monitor local store",
+                                                            ALLOCSET_DEFAULT_SIZES);
 	}
 
 #ifdef BENCHMARK
@@ -158,10 +165,6 @@ InitializeSharedState(pgsmSharedState * pgsm)
 {
 	pg_atomic_init_u64(&pgsm->current_wbucket, 0);
 	pg_atomic_init_u64(&pgsm->prev_bucket_sec, 0);
-	memset(&pgsm->bucket_entry, 0, MAX_BUCKETS * sizeof(uint64));
-	pgsm->pgsm_mem_cxt = AllocSetContextCreate(TopMemoryContext,
-											   "pg_stat_monitor local store",
-											   ALLOCSET_DEFAULT_SIZES);
 }
 
 
@@ -235,6 +238,11 @@ pgsm_attach_shmem(void)
 	MemoryContextSwitchTo(oldcontext);
 }
 
+MemoryContext GetPgsmMemoryContext(void)
+{
+	return pgsmStateLocal.pgsm_mem_cxt;
+}
+
 dsa_area *
 get_dsa_area_for_query_text(void)
 {
@@ -290,7 +298,6 @@ hash_entry_alloc(pgsmSharedState * pgsm, pgsmHashKey * key, int encoding)
 		elog(DEBUG1, "[pg_stat_monitor] hash_entry_alloc: OUT OF MEMORY.");
 	else if (!found)
 	{
-		pgsm->bucket_entry[pg_atomic_read_u64(&pgsm->current_wbucket)]++;
 		/* New entry, initialize it */
 		/* reset the statistics */
 		memset(&entry->counters, 0, sizeof(Counters));

pg_stat_monitor.c

@@ -1590,7 +1590,7 @@ static void
 pgsm_add_to_list(pgsmEntry * entry, char *query_text, int query_len)
 {
 	/* Switch to pgsm memory context */
-	MemoryContext oldctx = MemoryContextSwitchTo(pgsm_get_ss()->pgsm_mem_cxt);
+	MemoryContext oldctx = MemoryContextSwitchTo(GetPgsmMemoryContext());
 
 	entry->query_text.query_pointer = pnstrdup(query_text, query_len);
 	lentries = lappend(lentries, entry);
@@ -1645,7 +1645,8 @@ static void
 pgsm_cleanup_callback(void *arg)
 {
 	/* Reset the memory context holding the list */
-	MemoryContextReset(pgsm_get_ss()->pgsm_mem_cxt);
+	MemoryContextReset(GetPgsmMemoryContext());
+
 	lentries = NIL;
 	callback_setup = false;
 }
@@ -1666,7 +1667,7 @@ pgsm_create_hash_entry(uint64 bucket_id, uint64 queryid, PlanInfo * plan_info)
 	char	   *username = NULL;
 
 	/* Create an entry in the pgsm memory context */
-	oldctx = MemoryContextSwitchTo(pgsm_get_ss()->pgsm_mem_cxt);
+	oldctx = MemoryContextSwitchTo(GetPgsmMemoryContext());
 	entry = palloc0(sizeof(pgsmEntry));
 
 	/*

pg_stat_monitor.h

@@ -77,7 +77,6 @@
 #define HISTOGRAM_MAX_TIME		50000000
 #define MAX_RESPONSE_BUCKET 50
 #define INVALID_BUCKET_ID	-1
-#define MAX_BUCKETS			10
 #define TEXT_LEN			255
 #define ERROR_MESSAGE_LEN	100
 #define REL_TYPENAME_LEN	64
@@ -388,8 +387,6 @@ typedef struct pgsmSharedState
 	slock_t		mutex;			/* protects following fields only: */
 	pg_atomic_uint64 current_wbucket;
 	pg_atomic_uint64 prev_bucket_sec;
-	uint64		bucket_entry[MAX_BUCKETS];
-	TimestampTz bucket_start_time[MAX_BUCKETS]; /* start time of the bucket */
 	int			hash_tranche_id;
 	void	   *raw_dsa_area;	/* DSA area pointer to store query texts.
 								 * dshash also lives in this memory when
@@ -400,16 +397,9 @@ typedef struct pgsmSharedState
 	 * hash table handle. can be either classic shared memory hash or dshash
 	 * (if we are using USE_DYNAMIC_HASH)
 	 */
-	MemoryContext pgsm_mem_cxt;
 
-	/*
-	 * context to store stats in local memory until they are pushed to shared
-	 * hash
-	 */
-	int			resp_calls[MAX_RESPONSE_BUCKET + 2];	/* execution time's in
-														 * msec; including 2
-														 * outlier buckets */
 	bool		pgsm_oom;
+	TimestampTz bucket_start_time[]; /* start time of the bucket */
 }			pgsmSharedState;
 
 typedef struct pgsmLocalState
@@ -418,6 +408,8 @@ typedef struct pgsmLocalState
 	dsa_area   *dsa;			/* local dsa area for backend attached to the
 								 * dsa area created by postmaster at startup. */
 	PGSM_HASH_TABLE *shared_hash;
+    MemoryContext pgsm_mem_cxt;
+
 }			pgsmLocalState;
 
 #if PG_VERSION_NUM < 140000
@@ -479,6 +471,7 @@ void		pgsm_startup(void);
 
 /* hash_query.c */
 void		pgsm_startup(void);
+MemoryContext GetPgsmMemoryContext(void);
 
 /* guc.c */
 void		init_guc(void);