From 71bce9441ed2ba4311cbbcb0fa51f5d89766b701 Mon Sep 17 00:00:00 2001
From: Jelte Fennema <jelte.fennema@microsoft.com>
Date: Fri, 13 Jan 2023 17:11:00 +0100
Subject: [PATCH v7 4/4] Share prng state between all PGconns in process

Instead of having a PRNG state per connection this adds a process wide
prng state that is protected by a mutex.
---
 src/interfaces/libpq/fe-connect.c | 37 ++++++++++++++++++++++++++-----
 src/interfaces/libpq/libpq-int.h  |  1 -
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index 69ed891703a..ec974ac6b84 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -442,6 +442,10 @@ static bool parse_int_param(const char *value, int *result, PGconn *conn,
 /* global variable because fe-auth.c needs to access it */
 pgthreadlock_t pg_g_threadlock = default_threadlock;
 
+static pglock_t prng_lock = PGLOCK_INITIALIZER;
+static bool prng_seed_set = false;
+static pg_prng_state prng_state;
+
 
 /*
  *		pqDropConnection
@@ -1028,16 +1032,25 @@ parse_comma_separated_list(char **startptr, bool *more)
 static bool
 libpq_prng_init(PGconn *conn)
 {
+	if (!pglock(&prng_lock))
+		return false;
+
+	if (likely(prng_seed_set))
+		return pgunlock(&prng_lock);
+
 	if (unlikely(conn->randomseed))
 	{
 		int			rseed;
 
 		if (!parse_int_param(conn->randomseed, &rseed, conn, "random_seed"))
+		{
+			pgunlock(&prng_lock);
 			return false;
+		}
 
-		pg_prng_seed(&conn->prng_state, rseed);
+		pg_prng_seed(&prng_state, rseed);
 	}
-	else if (unlikely(!pg_prng_strong_seed(&conn->prng_state)))
+	else if (unlikely(!pg_prng_strong_seed(&prng_state)))
 	{
 		uint64		rseed;
 		time_t		now = time(NULL);
@@ -1053,8 +1066,12 @@ libpq_prng_init(PGconn *conn)
 			((uint64) now << 12) ^
 			((uint64) now >> 20);
 
-		pg_prng_seed(&conn->prng_state, rseed);
+		pg_prng_seed(&prng_state, rseed);
 	}
+
+	prng_seed_set = true;
+	if (!pgunlock(&prng_lock))
+		return false;
 	return true;
 }
 
@@ -1477,6 +1494,9 @@ connectOptions2(PGconn *conn)
 		if (!libpq_prng_init(conn))
 			return false;
 
+		if (!pglock(&prng_lock))
+			return false;
+
 		/*
 		 * Shuffle connhost with a Durstenfeld/Knuth version of the
 		 * Fisher-Yates shuffle. Source:
@@ -1484,12 +1504,14 @@ connectOptions2(PGconn *conn)
 		 */
 		for (i = conn->nconnhost - 1; i > 0; i--)
 		{
-			int			j = pg_prng_uint64_range(&conn->prng_state, 0, i);
+			int			j = pg_prng_uint64_range(&prng_state, 0, i);
 			pg_conn_host temp = conn->connhost[j];
 
 			conn->connhost[j] = conn->connhost[i];
 			conn->connhost[i] = temp;
 		}
+		if (!pgunlock(&prng_lock))
+			return false;
 	}
 
 	/*
@@ -4175,6 +4197,9 @@ store_conn_addrinfo(PGconn *conn, struct addrinfo *addrlist)
 
 	if (conn->load_balance_type == LOAD_BALANCE_RANDOM)
 	{
+		if (!pglock(&prng_lock))
+			return false;
+
 		/*
 		 * Shuffle addr with a Durstenfeld/Knuth version of the Fisher-Yates
 		 * shuffle. Source:
@@ -4185,12 +4210,14 @@ store_conn_addrinfo(PGconn *conn, struct addrinfo *addrlist)
 		 */
 		for (int i = conn->naddr - 1; i > 0; i--)
 		{
-			int			j = pg_prng_uint64_range(&conn->prng_state, 0, i);
+			int			j = pg_prng_uint64_range(&prng_state, 0, i);
 			AddrInfo	temp = conn->addr[j];
 
 			conn->addr[j] = conn->addr[i];
 			conn->addr[i] = temp;
 		}
+		if (!pgunlock(&prng_lock))
+			return false;
 	}
 	return true;
 }
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 01dd4190f33..0ecac062090 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -492,7 +492,6 @@ struct pg_conn
 	PGVerbosity verbosity;		/* error/notice message verbosity */
 	PGContextVisibility show_context;	/* whether to show CONTEXT field */
 	PGlobjfuncs *lobjfuncs;		/* private state for large-object access fns */
-	pg_prng_state prng_state;	/* prng state for load balancing connections */
 
 
 	/* Buffer for data received from backend and not yet processed */
-- 
2.34.1