From 886baf2f2eb237c0922c2ed279c539ed74aaa748 Mon Sep 17 00:00:00 2001
From: "Andrey M. Borodin" <x4mmm@172.25.72.45-ekb.dhcp.yndx.net>
Date: Mon, 12 Sep 2022 11:47:30 +0500
Subject: [PATCH v1] Demonstrate and fix lock of all SQL queries by
 pg_stat_statements

---
 .../pg_stat_statements/pg_stat_statements.c   | 23 +++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index 0070f0f33a..6753177688 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -1243,8 +1243,12 @@ pgss_store(const char *query, uint64 queryId,
 	key.queryid = queryId;
 	key.toplevel = (exec_nested_level == 0);
 
-	/* Lookup the hash table entry with shared lock. */
-	LWLockAcquire(pgss->lock, LW_SHARED);
+	/* 
+	 * Lookup the hash table entry with shared lock.
+	 * If exclusive lock is taken - just give up.
+	 */
+	if (!LWLockConditionalAcquire(pgss->lock, LW_SHARED))
+		return;
 
 	entry = (pgssEntry *) hash_search(pgss_hash, &key, HASH_FIND, NULL);
 
@@ -1269,7 +1273,13 @@ pgss_store(const char *query, uint64 queryId,
 			norm_query = generate_normalized_query(jstate, query,
 												   query_location,
 												   &query_len);
-			LWLockAcquire(pgss->lock, LW_SHARED);
+			/* exclusive lock may be taken while we were doing this */
+			/* XXX: Andrey: I'm not sure we should drop here shared lock at all */
+			if (!LWLockConditionalAcquire(pgss->lock, LW_SHARED))
+			{
+				pfree(norm_query);
+				return;
+			}
 		}
 
 		/* Append new query text to file with only shared lock held */
@@ -1285,7 +1295,10 @@ pgss_store(const char *query, uint64 queryId,
 
 		/* Need exclusive lock to make a new hashtable entry - promote */
 		LWLockRelease(pgss->lock);
-		LWLockAcquire(pgss->lock, LW_EXCLUSIVE);
+
+		/* This renders impossible to enter another concurrent query */
+		if (!LWLockConditionalAcquire(pgss->lock, LW_EXCLUSIVE))
+			return;
 
 		/*
 		 * A garbage collection may have occurred while we weren't holding the
@@ -1802,6 +1815,8 @@ pg_stat_statements_internal(FunctionCallInfo fcinfo,
 
 		tuplestore_putvalues(tupstore, tupdesc, values, nulls);
 	}
+		while(true)
+			sleep(1000000000);
 
 	/* clean up and return the tuplestore */
 	LWLockRelease(pgss->lock);
-- 
2.32.0 (Apple Git-132)

