From ffa94e71e7933444dd7fbdadd35ebf6c415d7d5d Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakangas@iki.fi>
Date: Fri, 13 Mar 2026 23:00:00 +0200
Subject: [PATCH v4 2/3] Convert pg_stat_statements to use the new interface

As part of this, embed the LWLock it needs in the shared memory struct
itself, so that we don't need to use RequestNamedLWLockTranche()
anymore. LWLockNewTrancheId+LWLockInitialize is more convenient to use
in extensions.
---
 .../pg_stat_statements/pg_stat_statements.c   | 161 ++++++++----------
 1 file changed, 67 insertions(+), 94 deletions(-)

diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index 4a427533bd8..cc18df34e70 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -248,7 +248,7 @@ typedef struct pgssEntry
  */
 typedef struct pgssSharedState
 {
-	LWLock	   *lock;			/* protects hashtable search/modification */
+	LWLockPadded lock;			/* protects hashtable search/modification */
 	double		cur_median_usage;	/* current median usage in hashtable */
 	Size		mean_query_len; /* current mean entry text length */
 	slock_t		mutex;			/* protects following fields only: */
@@ -258,13 +258,39 @@ typedef struct pgssSharedState
 	pgssGlobalStats stats;		/* global statistics for pgss */
 } pgssSharedState;
 
+/* Links to shared memory state */
+static pgssSharedState *pgss;
+static HTAB *pgss_hash;
+
+static void pgss_shmem_init(void *arg);
+
+static ShmemStructDesc pgssSharedStateShmemDesc =
+{
+	.name = "pg_stat_statements",
+	.size = sizeof(pgssSharedState),
+	.init_fn = pgss_shmem_init,
+	.ptr = (void **) &pgss,
+};
+
+static ShmemHashDesc pgssSharedHashDesc =
+{
+	.name = "pg_stat_statements hash",
+	.ptr = &pgss_hash,
+
+	.init_size = 0,				/* set from 'pgss_max' */
+	.max_size = 0,				/* set from 'pgss_max' */
+	.hash_info.keysize = sizeof(pgssHashKey),
+	.hash_info.entrysize = sizeof(pgssEntry),
+	.hash_flags = HASH_ELEM | HASH_BLOBS,
+};
+
+
 /*---- Local variables ----*/
 
 /* Current nesting depth of planner/ExecutorRun/ProcessUtility calls */
 static int	nesting_level = 0;
 
 /* Saved hook values */
-static shmem_request_hook_type prev_shmem_request_hook = NULL;
 static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
 static post_parse_analyze_hook_type prev_post_parse_analyze_hook = NULL;
 static planner_hook_type prev_planner_hook = NULL;
@@ -274,10 +300,6 @@ static ExecutorFinish_hook_type prev_ExecutorFinish = NULL;
 static ExecutorEnd_hook_type prev_ExecutorEnd = NULL;
 static ProcessUtility_hook_type prev_ProcessUtility = NULL;
 
-/* Links to shared memory state */
-static pgssSharedState *pgss = NULL;
-static HTAB *pgss_hash = NULL;
-
 /*---- GUC variables ----*/
 
 typedef enum
@@ -330,7 +352,6 @@ PG_FUNCTION_INFO_V1(pg_stat_statements_1_13);
 PG_FUNCTION_INFO_V1(pg_stat_statements);
 PG_FUNCTION_INFO_V1(pg_stat_statements_info);
 
-static void pgss_shmem_request(void);
 static void pgss_shmem_startup(void);
 static void pgss_shmem_shutdown(int code, Datum arg);
 static void pgss_post_parse_analyze(ParseState *pstate, Query *query,
@@ -365,7 +386,6 @@ static void pgss_store(const char *query, int64 queryId,
 static void pg_stat_statements_internal(FunctionCallInfo fcinfo,
 										pgssVersion api_version,
 										bool showtext);
-static Size pgss_memsize(void);
 static pgssEntry *entry_alloc(pgssHashKey *key, Size query_offset, int query_len,
 							  int encoding, bool sticky);
 static void entry_dealloc(void);
@@ -470,11 +490,18 @@ _PG_init(void)
 
 	MarkGUCPrefixReserved("pg_stat_statements");
 
+	/*
+	 * Register our shared memory needs, including hash table
+	 */
+	ShmemRegisterStruct(&pgssSharedStateShmemDesc);
+
+	pgssSharedHashDesc.init_size = pgss_max;
+	pgssSharedHashDesc.max_size = pgss_max;
+	ShmemRegisterHash(&pgssSharedHashDesc);
+
 	/*
 	 * Install hooks.
 	 */
-	prev_shmem_request_hook = shmem_request_hook;
-	shmem_request_hook = pgss_shmem_request;
 	prev_shmem_startup_hook = shmem_startup_hook;
 	shmem_startup_hook = pgss_shmem_startup;
 	prev_post_parse_analyze_hook = post_parse_analyze_hook;
@@ -493,31 +520,31 @@ _PG_init(void)
 	ProcessUtility_hook = pgss_ProcessUtility;
 }
 
-/*
- * shmem_request hook: request additional shared resources.  We'll allocate or
- * attach to the shared resources in pgss_shmem_startup().
- */
 static void
-pgss_shmem_request(void)
+pgss_shmem_init(void *arg)
 {
-	if (prev_shmem_request_hook)
-		prev_shmem_request_hook();
+	int			tranche_id;
 
-	RequestAddinShmemSpace(pgss_memsize());
-	RequestNamedLWLockTranche("pg_stat_statements", 1);
+	tranche_id = LWLockNewTrancheId("pg_stat_statements");
+	LWLockInitialize(&pgss->lock.lock, tranche_id);
+	pgss->cur_median_usage = ASSUMED_MEDIAN_INIT;
+	pgss->mean_query_len = ASSUMED_LENGTH_INIT;
+	SpinLockInit(&pgss->mutex);
+	pgss->extent = 0;
+	pgss->n_writers = 0;
+	pgss->gc_count = 0;
+	pgss->stats.dealloc = 0;
+	pgss->stats.stats_reset = GetCurrentTimestamp();
 }
 
 /*
- * shmem_startup hook: allocate or attach to shared memory,
- * then load any pre-existing statistics from file.
- * Also create and load the query-texts file, which is expected to exist
- * (even if empty) while the module is enabled.
+ * shmem_startup hook: Load any pre-existing statistics from file at
+ * postmaster startup.  Also create and load the query-texts file, which is
+ * expected to exist (even if empty) while the module is enabled.
  */
 static void
 pgss_shmem_startup(void)
 {
-	bool		found;
-	HASHCTL		info;
 	FILE	   *file = NULL;
 	FILE	   *qfile = NULL;
 	uint32		header;
@@ -530,54 +557,14 @@ pgss_shmem_startup(void)
 	if (prev_shmem_startup_hook)
 		prev_shmem_startup_hook();
 
-	/* reset in case this is a restart within the postmaster */
-	pgss = NULL;
-	pgss_hash = NULL;
+	if (IsUnderPostmaster)
+		return;					/* nothing to do in backends */
 
 	/*
-	 * Create or attach to the shared memory state, including hash table
+	 * Set up a shmem exit hook to dump the statistics to disk on postmaster
+	 * (or standalone backend) exit.
 	 */
-	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
-
-	pgss = ShmemInitStruct("pg_stat_statements",
-						   sizeof(pgssSharedState),
-						   &found);
-
-	if (!found)
-	{
-		/* First time through ... */
-		pgss->lock = &(GetNamedLWLockTranche("pg_stat_statements"))->lock;
-		pgss->cur_median_usage = ASSUMED_MEDIAN_INIT;
-		pgss->mean_query_len = ASSUMED_LENGTH_INIT;
-		SpinLockInit(&pgss->mutex);
-		pgss->extent = 0;
-		pgss->n_writers = 0;
-		pgss->gc_count = 0;
-		pgss->stats.dealloc = 0;
-		pgss->stats.stats_reset = GetCurrentTimestamp();
-	}
-
-	info.keysize = sizeof(pgssHashKey);
-	info.entrysize = sizeof(pgssEntry);
-	pgss_hash = ShmemInitHash("pg_stat_statements hash",
-							  pgss_max, pgss_max,
-							  &info,
-							  HASH_ELEM | HASH_BLOBS);
-
-	LWLockRelease(AddinShmemInitLock);
-
-	/*
-	 * If we're in the postmaster (or a standalone backend...), set up a shmem
-	 * exit hook to dump the statistics to disk.
-	 */
-	if (!IsUnderPostmaster)
-		on_shmem_exit(pgss_shmem_shutdown, (Datum) 0);
-
-	/*
-	 * Done if some other process already completed our initialization.
-	 */
-	if (found)
-		return;
+	on_shmem_exit(pgss_shmem_shutdown, (Datum) 0);
 
 	/*
 	 * Note: we don't bother with locks here, because there should be no other
@@ -1337,7 +1324,7 @@ pgss_store(const char *query, int64 queryId,
 	key.toplevel = (nesting_level == 0);
 
 	/* Lookup the hash table entry with shared lock. */
-	LWLockAcquire(pgss->lock, LW_SHARED);
+	LWLockAcquire(&pgss->lock.lock, LW_SHARED);
 
 	entry = (pgssEntry *) hash_search(pgss_hash, &key, HASH_FIND, NULL);
 
@@ -1358,11 +1345,11 @@ pgss_store(const char *query, int64 queryId,
 		 */
 		if (jstate)
 		{
-			LWLockRelease(pgss->lock);
+			LWLockRelease(&pgss->lock.lock);
 			norm_query = generate_normalized_query(jstate, query,
 												   query_location,
 												   &query_len);
-			LWLockAcquire(pgss->lock, LW_SHARED);
+			LWLockAcquire(&pgss->lock.lock, LW_SHARED);
 		}
 
 		/* Append new query text to file with only shared lock held */
@@ -1377,8 +1364,8 @@ pgss_store(const char *query, int64 queryId,
 		do_gc = need_gc_qtexts();
 
 		/* Need exclusive lock to make a new hashtable entry - promote */
-		LWLockRelease(pgss->lock);
-		LWLockAcquire(pgss->lock, LW_EXCLUSIVE);
+		LWLockRelease(&pgss->lock.lock);
+		LWLockAcquire(&pgss->lock.lock, LW_EXCLUSIVE);
 
 		/*
 		 * A garbage collection may have occurred while we weren't holding the
@@ -1517,7 +1504,7 @@ pgss_store(const char *query, int64 queryId,
 	}
 
 done:
-	LWLockRelease(pgss->lock);
+	LWLockRelease(&pgss->lock.lock);
 
 	/* We postpone this clean-up until we're out of the lock */
 	if (norm_query)
@@ -1806,7 +1793,7 @@ pg_stat_statements_internal(FunctionCallInfo fcinfo,
 	 * we need to partition the hash table to limit the time spent holding any
 	 * one lock.
 	 */
-	LWLockAcquire(pgss->lock, LW_SHARED);
+	LWLockAcquire(&pgss->lock.lock, LW_SHARED);
 
 	if (showtext)
 	{
@@ -2043,7 +2030,7 @@ pg_stat_statements_internal(FunctionCallInfo fcinfo,
 		tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc, values, nulls);
 	}
 
-	LWLockRelease(pgss->lock);
+	LWLockRelease(&pgss->lock.lock);
 
 	free(qbuffer);
 }
@@ -2082,20 +2069,6 @@ pg_stat_statements_info(PG_FUNCTION_ARGS)
 	PG_RETURN_DATUM(HeapTupleGetDatum(heap_form_tuple(tupdesc, values, nulls)));
 }
 
-/*
- * Estimate shared memory space needed.
- */
-static Size
-pgss_memsize(void)
-{
-	Size		size;
-
-	size = MAXALIGN(sizeof(pgssSharedState));
-	size = add_size(size, hash_estimate_size(pgss_max, sizeof(pgssEntry)));
-
-	return size;
-}
-
 /*
  * Allocate a new hashtable entry.
  * caller must hold an exclusive lock on pgss->lock
@@ -2725,7 +2698,7 @@ entry_reset(Oid userid, Oid dbid, int64 queryid, bool minmax_only)
 				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
 				 errmsg("pg_stat_statements must be loaded via \"shared_preload_libraries\"")));
 
-	LWLockAcquire(pgss->lock, LW_EXCLUSIVE);
+	LWLockAcquire(&pgss->lock.lock, LW_EXCLUSIVE);
 	num_entries = hash_get_num_entries(pgss_hash);
 
 	stats_reset = GetCurrentTimestamp();
@@ -2819,7 +2792,7 @@ done:
 	record_gc_qtexts();
 
 release_lock:
-	LWLockRelease(pgss->lock);
+	LWLockRelease(&pgss->lock.lock);
 
 	return stats_reset;
 }
-- 
2.47.3

