v7-0002-Support-load-balancing-in-libpq.patch

From c92a730bff697d7e2ce694013ca677bcd947e3c6 Mon Sep 17 00:00:00 2001
From: Jelte Fennema <github-tech@jeltef.nl>
Date: Mon, 12 Sep 2022 09:44:06 +0200
Subject: [PATCH v7 2/4] Support load balancing in libpq

This adds support for load balancing to libpq using the newly added
load_balance_hosts parameter. When setting the load_balance_hosts
parameter to random, hosts and addresses will be connected to in a
random order. This then results in load balancing across these
hosts/addresses if multiple clients do this at the same time.

This patch implements two levels of random load balancing:
1. The given hosts are randomly shuffled, before resolving them
    one-by-one.
2. Once a host its addresses get resolved, those addresses are shuffled,
    before trying to connect to them one-by-one.
---
 .cirrus.yml                               |  14 ++
 doc/src/sgml/libpq.sgml                   |  69 +++++++
 src/include/libpq/pqcomm.h                |   6 +
 src/interfaces/libpq/fe-connect.c         | 215 ++++++++++++++++++----
 src/interfaces/libpq/libpq-int.h          |  21 ++-
 src/interfaces/libpq/meson.build          |   1 +
 src/interfaces/libpq/t/003_loadbalance.pl | 167 +++++++++++++++++
 7 files changed, 460 insertions(+), 33 deletions(-)
 create mode 100644 src/interfaces/libpq/t/003_loadbalance.pl

diff --git a/.cirrus.yml b/.cirrus.yml
index 69837bcd5ad..f71e93451fa 100644
--- a/.cirrus.yml
+++ b/.cirrus.yml
@@ -308,6 +308,14 @@ task:
     mkdir -m 770 /tmp/cores
     chown root:postgres /tmp/cores
     sysctl kernel.core_pattern='/tmp/cores/%e-%s-%p.core'
+
+  setup_hosts_file_script: |
+    cat >> /etc/hosts <<-EOF
+      127.0.0.1 pg-loadbalancetest
+      127.0.0.2 pg-loadbalancetest
+      127.0.0.3 pg-loadbalancetest
+    EOF
+
   setup_additional_packages_script: |
     #apt-get update
     #DEBIAN_FRONTEND=noninteractive apt-get -y install ...
@@ -557,6 +565,12 @@ task:
   setup_additional_packages_script: |
     REM choco install -y --no-progress ...
 
+  setup_hosts_file_script: |
+    echo 127.0.0.1 pg-loadbalancetest >> c:\Windows\System32\Drivers\etc\hosts
+    echo 127.0.0.2 pg-loadbalancetest >> c:\Windows\System32\Drivers\etc\hosts
+    echo 127.0.0.3 pg-loadbalancetest >> c:\Windows\System32\Drivers\etc\hosts
+    type c:\Windows\System32\Drivers\etc\hosts
+
   # Use /DEBUG:FASTLINK to avoid high memory usage during linking
   configure_script: |
     vcvarsall x64
diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index 0e7ae70c706..c540554ade7 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -1964,6 +1964,75 @@ postgresql://%2Fvar%2Flib%2Fpostgresql/dbname
        </para>
       </listitem>
      </varlistentry>
+
+     <varlistentry id="libpq-connect-load-balance-hosts" xreflabel="load_balance_hosts">
+      <term><literal>load_balance_hosts</literal></term>
+      <listitem>
+       <para>
+        Controls the order in which the client tries to connect to the available
+        hosts and addresses. It's typically used in combination with multiple
+        host names or a DNS record that returns multiple IPs. This parameter be
+        used in combination with <xref linkend="libpq-connect-target-session-attrs"/>
+        to, for example, load balance over stanby servers only. Once successfully
+        connected, subsequent queries on the returned connection will all be
+        sent to the same server. There are currently two modes:
+        <variablelist>
+         <varlistentry>
+          <term><literal>disable</literal> (default)</term>
+          <listitem>
+           <para>
+            Hosts are tried in the order in which they are provided and
+            addresses are tried in the order they are received from DNS or a
+            hosts file.
+           </para>
+          </listitem>
+         </varlistentry>
+
+         <varlistentry>
+          <term><literal>random</literal></term>
+          <listitem>
+           <para>
+            The provided hosts and the addresses that they resolve to are
+            tried in random order. This value is mostly useful when opening
+            multiple connections at the same time, possibly from different
+            machines. This way connections can be load balanced across multiple
+            Postgres servers.
+           </para>
+           <para>
+            This algorithm uses two levels of random choices: First the hosts
+            will be resolved in random order. Then before resolving the next
+            host, all resolved addresses for the current host will be tried in
+            random order. This behaviour can lead to non-uniform address
+            selection in certain cases, for instance when some hosts resolve to
+            more addresses than others. So if you want uniform load balancing,
+            this is something to keep in mind. However, non-uniform load
+            balancing also can be used to your advantage, e.g. by providing the
+            hostname of a larger server multiple times in the host string so it
+            gets more connections.
+           </para>
+           <para>
+            When using this value it's recommended to also configure a reasonable
+            value for <xref linkend="libpq-connect-connect-timeout"/>. Because then,
+            if one of the nodes that are used for load balancing is not responding,
+            a new node will be tried.
+           </para>
+          </listitem>
+         </varlistentry>
+        </variablelist>
+       </para>
+      </listitem>
+     </varlistentry>
+
+     <varlistentry id="libpq-random-seed" xreflabel="random_seed">
+      <term><literal>random_seed</literal></term>
+      <listitem>
+       <para>
+        Sets the random seed that is used by <xref linkend="libpq-connect-load-balance-hosts"/>
+        to randomize the host order. This option is mostly useful when running
+        tests that require a stable random order.
+       </para>
+      </listitem>
+     </varlistentry>
     </variablelist>
    </para>
   </sect2>
diff --git a/src/include/libpq/pqcomm.h b/src/include/libpq/pqcomm.h
index 66ba359390f..ee28e223bd7 100644
--- a/src/include/libpq/pqcomm.h
+++ b/src/include/libpq/pqcomm.h
@@ -27,6 +27,12 @@ typedef struct
 	socklen_t	salen;
 } SockAddr;
 
+typedef struct
+{
+	int			family;
+	SockAddr	addr;
+} AddrInfo;
+
 /* Configure the UNIX socket location for the well known port. */
 
 #define UNIXSOCK_PATH(path, port, sockdir) \
diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index 773e9e1f3a2..18a07d810dc 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -123,6 +123,7 @@ static int	ldapServiceLookup(const char *purl, PQconninfoOption *options,
 #define DefaultChannelBinding	"disable"
 #endif
 #define DefaultTargetSessionAttrs	"any"
+#define DefaultLoadBalanceHosts	"disable"
 #ifdef USE_SSL
 #define DefaultSSLMode "prefer"
 #else
@@ -341,6 +342,15 @@ static const internalPQconninfoOption PQconninfoOptions[] = {
 		"Target-Session-Attrs", "", 15, /* sizeof("prefer-standby") = 15 */
 	offsetof(struct pg_conn, target_session_attrs)},
 
+	{"load_balance_hosts", "PGLOADBALANCEHOSTS",
+		DefaultLoadBalanceHosts, NULL,
+		"Load-Balance-Hosts", "", 8,	/* sizeof("disable") = 8 */
+	offsetof(struct pg_conn, load_balance_hosts)},
+
+	{"random_seed", NULL, NULL, NULL,
+		"Random-Seed", "", 10,	/* strlen(INT32_MAX) == 10 */
+	offsetof(struct pg_conn, randomseed)},
+
 	/* Terminating entry --- MUST BE LAST */
 	{NULL, NULL, NULL, NULL,
 	NULL, NULL, 0}
@@ -379,6 +389,7 @@ static bool fillPGconn(PGconn *conn, PQconninfoOption *connOptions);
 static void freePGconn(PGconn *conn);
 static void closePGconn(PGconn *conn);
 static void release_conn_addrinfo(PGconn *conn);
+static bool store_conn_addrinfo(PGconn *conn, struct addrinfo *addrlist);
 static void sendTerminateConn(PGconn *conn);
 static PQconninfoOption *conninfo_init(PQExpBuffer errorMessage);
 static PQconninfoOption *parse_connection_string(const char *connstr,
@@ -424,6 +435,8 @@ static void pgpassfileWarning(PGconn *conn);
 static void default_threadlock(int acquire);
 static bool sslVerifyProtocolVersion(const char *version);
 static bool sslVerifyProtocolRange(const char *min, const char *max);
+static bool parse_int_param(const char *value, int *result, PGconn *conn,
+							const char *context);
 
 
 /* global variable because fe-auth.c needs to access it */
@@ -1007,6 +1020,44 @@ parse_comma_separated_list(char **startptr, bool *more)
 	return p;
 }
 
+/*
+ * Initializes the prng_state field of the connection. We want something
+ * unpredictable, so if possible, use high-quality random bits for the
+ * seed. Otherwise, fall back to a seed based on timestamp and PID.
+ */
+static bool
+libpq_prng_init(PGconn *conn)
+{
+	if (unlikely(conn->randomseed))
+	{
+		int			rseed;
+
+		if (!parse_int_param(conn->randomseed, &rseed, conn, "random_seed"))
+			return false;
+
+		pg_prng_seed(&conn->prng_state, rseed);
+	}
+	else if (unlikely(!pg_prng_strong_seed(&conn->prng_state)))
+	{
+		uint64		rseed;
+		time_t		now = time(NULL);
+
+		/*
+		 * Since PIDs and timestamps tend to change more frequently in their
+		 * least significant bits, shift the timestamp left to allow a larger
+		 * total number of seeds in a given time period.  Since that would
+		 * leave only 20 bits of the timestamp that cycle every ~1 second,
+		 * also mix in some higher bits.
+		 */
+		rseed = ((uint64) getpid()) ^
+			((uint64) now << 12) ^
+			((uint64) now >> 20);
+
+		pg_prng_seed(&conn->prng_state, rseed);
+	}
+	return true;
+}
+
 /*
  *		connectOptions2
  *
@@ -1400,6 +1451,47 @@ connectOptions2(PGconn *conn)
 	else
 		conn->target_server_type = SERVER_TYPE_ANY;
 
+	/*
+	 * validate load_balance_hosts option, and set load_balance_type
+	 */
+	if (conn->load_balance_hosts)
+	{
+		if (strcmp(conn->load_balance_hosts, "disable") == 0)
+			conn->load_balance_type = LOAD_BALANCE_DISABLE;
+		else if (strcmp(conn->load_balance_hosts, "random") == 0)
+			conn->load_balance_type = LOAD_BALANCE_RANDOM;
+		else
+		{
+			conn->status = CONNECTION_BAD;
+			libpq_append_conn_error(conn, "invalid %s value: \"%s\"",
+									"load_balance_hosts",
+									conn->load_balance_hosts);
+			return false;
+		}
+	}
+	else
+		conn->load_balance_type = LOAD_BALANCE_DISABLE;
+
+	if (conn->load_balance_type == LOAD_BALANCE_RANDOM)
+	{
+		if (!libpq_prng_init(conn))
+			return false;
+
+		/*
+		 * Shuffle connhost with a Durstenfeld/Knuth version of the
+		 * Fisher-Yates shuffle. Source:
+		 * https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#The_modern_algorithm
+		 */
+		for (i = conn->nconnhost - 1; i > 0; i--)
+		{
+			int			j = pg_prng_uint64_range(&conn->prng_state, 0, i);
+			pg_conn_host temp = conn->connhost[j];
+
+			conn->connhost[j] = conn->connhost[i];
+			conn->connhost[i] = temp;
+		}
+	}
+
 	/*
 	 * Resolve special "auto" client_encoding from the locale
 	 */
@@ -2077,7 +2169,7 @@ connectDBComplete(PGconn *conn)
 	time_t		finish_time = ((time_t) -1);
 	int			timeout = 0;
 	int			last_whichhost = -2;	/* certainly different from whichhost */
-	struct addrinfo *last_addr_cur = NULL;
+	int			last_whichaddr = -2;	/* certainly different from whichaddr */
 
 	if (conn == NULL || conn->status == CONNECTION_BAD)
 		return 0;
@@ -2121,11 +2213,11 @@ connectDBComplete(PGconn *conn)
 		if (flag != PGRES_POLLING_OK &&
 			timeout > 0 &&
 			(conn->whichhost != last_whichhost ||
-			 conn->addr_cur != last_addr_cur))
+			 conn->whichaddr != last_whichaddr))
 		{
 			finish_time = time(NULL) + timeout;
 			last_whichhost = conn->whichhost;
-			last_addr_cur = conn->addr_cur;
+			last_whichaddr = conn->whichaddr;
 		}
 
 		/*
@@ -2272,9 +2364,9 @@ keep_going:						/* We will come back to here until there is
 	/* Time to advance to next address, or next host if no more addresses? */
 	if (conn->try_next_addr)
 	{
-		if (conn->addr_cur && conn->addr_cur->ai_next)
+		if (conn->whichaddr < conn->naddr)
 		{
-			conn->addr_cur = conn->addr_cur->ai_next;
+			conn->whichaddr++;
 			reset_connection_state_machine = true;
 		}
 		else
@@ -2287,6 +2379,7 @@ keep_going:						/* We will come back to here until there is
 	{
 		pg_conn_host *ch;
 		struct addrinfo hint;
+		struct addrinfo *addrlist;
 		int			thisport;
 		int			ret;
 		char		portstr[MAXPGPATH];
@@ -2327,7 +2420,7 @@ keep_going:						/* We will come back to here until there is
 		/* Initialize hint structure */
 		MemSet(&hint, 0, sizeof(hint));
 		hint.ai_socktype = SOCK_STREAM;
-		conn->addrlist_family = hint.ai_family = AF_UNSPEC;
+		hint.ai_family = AF_UNSPEC;
 
 		/* Figure out the port number we're going to use. */
 		if (ch->port == NULL || ch->port[0] == '\0')
@@ -2350,8 +2443,8 @@ keep_going:						/* We will come back to here until there is
 		{
 			case CHT_HOST_NAME:
 				ret = pg_getaddrinfo_all(ch->host, portstr, &hint,
-										 &conn->addrlist);
-				if (ret || !conn->addrlist)
+										 &addrlist);
+				if (ret || !addrlist)
 				{
 					libpq_append_conn_error(conn, "could not translate host name \"%s\" to address: %s",
 											ch->host, gai_strerror(ret));
@@ -2362,8 +2455,8 @@ keep_going:						/* We will come back to here until there is
 			case CHT_HOST_ADDRESS:
 				hint.ai_flags = AI_NUMERICHOST;
 				ret = pg_getaddrinfo_all(ch->hostaddr, portstr, &hint,
-										 &conn->addrlist);
-				if (ret || !conn->addrlist)
+										 &addrlist);
+				if (ret || !addrlist)
 				{
 					libpq_append_conn_error(conn, "could not parse network address \"%s\": %s",
 											ch->hostaddr, gai_strerror(ret));
@@ -2372,7 +2465,7 @@ keep_going:						/* We will come back to here until there is
 				break;
 
 			case CHT_UNIX_SOCKET:
-				conn->addrlist_family = hint.ai_family = AF_UNIX;
+				hint.ai_family = AF_UNIX;
 				UNIXSOCK_PATH(portstr, thisport, ch->host);
 				if (strlen(portstr) >= UNIXSOCK_PATH_BUFLEN)
 				{
@@ -2387,8 +2480,8 @@ keep_going:						/* We will come back to here until there is
 				 * name as a Unix-domain socket path.
 				 */
 				ret = pg_getaddrinfo_all(NULL, portstr, &hint,
-										 &conn->addrlist);
-				if (ret || !conn->addrlist)
+										 &addrlist);
+				if (ret || !addrlist)
 				{
 					libpq_append_conn_error(conn, "could not translate Unix-domain socket path \"%s\" to address: %s",
 											portstr, gai_strerror(ret));
@@ -2397,8 +2490,14 @@ keep_going:						/* We will come back to here until there is
 				break;
 		}
 
-		/* OK, scan this addrlist for a working server address */
-		conn->addr_cur = conn->addrlist;
+		if (!store_conn_addrinfo(conn, addrlist))
+		{
+			pg_freeaddrinfo_all(hint.ai_family, addrlist);
+			libpq_append_conn_error(conn, "out of memory");
+			goto error_return;
+		}
+		pg_freeaddrinfo_all(hint.ai_family, addrlist);
+
 		reset_connection_state_machine = true;
 		conn->try_next_host = false;
 	}
@@ -2455,30 +2554,29 @@ keep_going:						/* We will come back to here until there is
 			{
 				/*
 				 * Try to initiate a connection to one of the addresses
-				 * returned by pg_getaddrinfo_all().  conn->addr_cur is the
+				 * returned by pg_getaddrinfo_all().  conn->whichaddr is the
 				 * next one to try.
 				 *
 				 * The extra level of braces here is historical.  It's not
 				 * worth reindenting this whole switch case to remove 'em.
 				 */
 				{
-					struct addrinfo *addr_cur = conn->addr_cur;
 					char		host_addr[NI_MAXHOST];
+					AddrInfo   *addr_cur;
 
 					/*
 					 * Advance to next possible host, if we've tried all of
 					 * the addresses for the current host.
 					 */
-					if (addr_cur == NULL)
+					if (conn->whichaddr == conn->naddr)
 					{
 						conn->try_next_host = true;
 						goto keep_going;
 					}
+					addr_cur = &conn->addr[conn->whichaddr];
 
 					/* Remember current address for possible use later */
-					memcpy(&conn->raddr.addr, addr_cur->ai_addr,
-						   addr_cur->ai_addrlen);
-					conn->raddr.salen = addr_cur->ai_addrlen;
+					memcpy(&conn->raddr, &addr_cur->addr, sizeof(SockAddr));
 
 					/*
 					 * Set connip, too.  Note we purposely ignore strdup
@@ -2494,7 +2592,7 @@ keep_going:						/* We will come back to here until there is
 						conn->connip = strdup(host_addr);
 
 					/* Try to create the socket */
-					conn->sock = socket(addr_cur->ai_family, SOCK_STREAM, 0);
+					conn->sock = socket(addr_cur->family, SOCK_STREAM, 0);
 					if (conn->sock == PGINVALID_SOCKET)
 					{
 						int			errorno = SOCK_ERRNO;
@@ -2505,7 +2603,7 @@ keep_going:						/* We will come back to here until there is
 						 * cases where the address list includes both IPv4 and
 						 * IPv6 but kernel only accepts one family.
 						 */
-						if (addr_cur->ai_next != NULL ||
+						if (conn->whichaddr < conn->naddr ||
 							conn->whichhost + 1 < conn->nconnhost)
 						{
 							conn->try_next_addr = true;
@@ -2531,7 +2629,7 @@ keep_going:						/* We will come back to here until there is
 					 * TCP sockets, nonblock mode, close-on-exec.  Try the
 					 * next address if any of this fails.
 					 */
-					if (addr_cur->ai_family != AF_UNIX)
+					if (addr_cur->family != AF_UNIX)
 					{
 						if (!connectNoDelay(conn))
 						{
@@ -2558,7 +2656,7 @@ keep_going:						/* We will come back to here until there is
 					}
 #endif							/* F_SETFD */
 
-					if (addr_cur->ai_family != AF_UNIX)
+					if (addr_cur->family != AF_UNIX)
 					{
 #ifndef WIN32
 						int			on = 1;
@@ -2650,8 +2748,8 @@ keep_going:						/* We will come back to here until there is
 					 * Start/make connection.  This should not block, since we
 					 * are in nonblock mode.  If it does, well, too bad.
 					 */
-					if (connect(conn->sock, addr_cur->ai_addr,
-								addr_cur->ai_addrlen) < 0)
+					if (connect(conn->sock, (struct sockaddr *) &addr_cur->addr.addr,
+								addr_cur->addr.salen) < 0)
 					{
 						if (SOCK_ERRNO == EINPROGRESS ||
 #ifdef WIN32
@@ -4035,12 +4133,68 @@ freePGconn(PGconn *conn)
 	free(conn->outBuffer);
 	free(conn->rowBuf);
 	free(conn->target_session_attrs);
+	free(conn->load_balance_hosts);
+	free(conn->randomseed);
 	termPQExpBuffer(&conn->errorMessage);
 	termPQExpBuffer(&conn->workBuffer);
 
 	free(conn);
 }
 
+/*
+ * Copies over the AddrInfos from addrlist to the PGconn.
+ */
+static bool
+store_conn_addrinfo(PGconn *conn, struct addrinfo *addrlist)
+{
+	struct addrinfo *ai = addrlist;
+
+	conn->whichaddr = 0;
+
+	conn->naddr = 0;
+	while (ai)
+	{
+		ai = ai->ai_next;
+		conn->naddr++;
+	}
+
+	conn->addr = calloc(conn->naddr, sizeof(AddrInfo));
+	if (conn->addr == NULL)
+		return false;
+
+	ai = addrlist;
+	for (int i = 0; i < conn->naddr; i++)
+	{
+		conn->addr[i].family = ai->ai_family;
+
+		memcpy(&conn->addr[i].addr.addr, ai->ai_addr,
+			   ai->ai_addrlen);
+		conn->addr[i].addr.salen = ai->ai_addrlen;
+		ai = ai->ai_next;
+	}
+
+	if (conn->load_balance_type == LOAD_BALANCE_RANDOM)
+	{
+		/*
+		 * Shuffle addr with a Durstenfeld/Knuth version of the Fisher-Yates
+		 * shuffle. Source:
+		 * https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#The_modern_algorithm
+		 *
+		 * We don't need to initialize conn->prng_state here, because that
+		 * already happened in connectOptions2.
+		 */
+		for (int i = conn->naddr - 1; i > 0; i--)
+		{
+			int			j = pg_prng_uint64_range(&conn->prng_state, 0, i);
+			AddrInfo	temp = conn->addr[j];
+
+			conn->addr[j] = conn->addr[i];
+			conn->addr[i] = temp;
+		}
+	}
+	return true;
+}
+
 /*
  * release_conn_addrinfo
  *	 - Free any addrinfo list in the PGconn.
@@ -4048,11 +4202,10 @@ freePGconn(PGconn *conn)
 static void
 release_conn_addrinfo(PGconn *conn)
 {
-	if (conn->addrlist)
+	if (conn->addr)
 	{
-		pg_freeaddrinfo_all(conn->addrlist_family, conn->addrlist);
-		conn->addrlist = NULL;
-		conn->addr_cur = NULL;	/* for safety */
+		free(conn->addr);
+		conn->addr = NULL;
 	}
 }
 
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 712d572373c..86dd1d6d405 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -82,6 +82,8 @@ typedef struct
 #endif
 #endif							/* USE_OPENSSL */
 
+#include "common/pg_prng.h"
+
 /*
  * POSTGRES backend dependent Constants.
  */
@@ -242,6 +244,13 @@ typedef enum
 	SERVER_TYPE_PREFER_STANDBY_PASS2	/* second pass - behaves same as ANY */
 } PGTargetServerType;
 
+/* Target server type (decoded value of load_balance_hosts) */
+typedef enum
+{
+	LOAD_BALANCE_DISABLE = 0,	/* Use the existing host order (default) */
+	LOAD_BALANCE_RANDOM,		/* Read-write server */
+}			PGLoadBalanceType;
+
 /* Boolean value plus a not-known state, for GUCs we might have to fetch */
 typedef enum
 {
@@ -396,6 +405,8 @@ struct pg_conn
 	char	   *ssl_min_protocol_version;	/* minimum TLS protocol version */
 	char	   *ssl_max_protocol_version;	/* maximum TLS protocol version */
 	char	   *target_session_attrs;	/* desired session properties */
+	char	   *load_balance_hosts; /* load balance over hosts */
+	char	   *randomseed;		/* seed for randomization of load balancing */
 
 	/* Optional file to write trace info to */
 	FILE	   *Pfdebug;
@@ -459,10 +470,14 @@ struct pg_conn
 
 	/* Transient state needed while establishing connection */
 	PGTargetServerType target_server_type;	/* desired session properties */
+	PGLoadBalanceType load_balance_type;	/* desired load balancing
+											 * algorithm */
 	bool		try_next_addr;	/* time to advance to next address/host? */
 	bool		try_next_host;	/* time to advance to next connhost[]? */
-	struct addrinfo *addrlist;	/* list of addresses for current connhost */
-	struct addrinfo *addr_cur;	/* the one currently being tried */
+	int			naddr;			/* number of addresses returned by getaddrinfo */
+	int			whichaddr;		/* the address currently being tried */
+	AddrInfo   *addr;			/* the array of addresses for the currently
+								 * tried host */
 	int			addrlist_family;	/* needed to know how to free addrlist */
 	bool		send_appname;	/* okay to send application_name? */
 
@@ -477,6 +492,8 @@ 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 */
 	char	   *inBuffer;		/* currently allocated buffer */
diff --git a/src/interfaces/libpq/meson.build b/src/interfaces/libpq/meson.build
index 573fd9b6ea4..52b327500c7 100644
--- a/src/interfaces/libpq/meson.build
+++ b/src/interfaces/libpq/meson.build
@@ -116,6 +116,7 @@ tests += {
     'tests': [
       't/001_uri.pl',
       't/002_api.pl',
+      't/003_loadbalance.pl',
     ],
     'env': {'with_ssl': get_option('ssl')},
   },
diff --git a/src/interfaces/libpq/t/003_loadbalance.pl b/src/interfaces/libpq/t/003_loadbalance.pl
new file mode 100644
index 00000000000..1b4e8fd8133
--- /dev/null
+++ b/src/interfaces/libpq/t/003_loadbalance.pl
@@ -0,0 +1,167 @@
+# Copyright (c) 2023, PostgreSQL Global Development Group
+use strict;
+use warnings;
+use Config;
+use PostgreSQL::Test::Utils;
+use PostgreSQL::Test::Cluster;
+use File::Spec::Functions 'catfile';
+use Test::More;
+
+# This tests two different methods of load balancing from libpq
+# 1. Load balancing by providing multiple host and port combinations in the
+#    libpq connection string.
+# 2. By using a hosts file where hostname maps to multiple different IP
+#    addresses. Regular Postgres users wouldn't usually use such a host file,
+#    but this is the easiest way to immitate behaviour of a DNS server that
+#    returns multiple IP addresses for the same DNS record.
+#
+# Testing method 1 is supported on all platforms and works out of the box. But
+# testing method 2 has some more requirements, both on the platform and on the
+# initial setup. If any of these requirements are not met, then method 2 is
+# simply not tested.
+#
+# The requirements to test method 2 are as follows:
+# 1. Windows or Linux should be used.
+# 2. The OS hosts file at /etc/hosts or c:\Windows\System32\Drivers\etc\hosts
+#    should contain the following contents:
+#
+# 127.0.0.1 pg-loadbalancetest
+# 127.0.0.2 pg-loadbalancetest
+# 127.0.0.3 pg-loadbalancetest
+#
+#
+# Windows or Linux are required to test method 2 because these OSes allow
+# binding to 127.0.0.2 and 127.0.0.3 addresess by default, but other OSes
+# don't. We need to bind to different IP addresses, so that we can use these
+# different IP addresses in the hosts file.
+#
+# The hosts file needs to be prepared before running this test. We don't do it
+# on the fly, because it requires root permissions to change the hosts file. In
+# CI we set up the previously mentioned rules in the hosts file, so that this
+# load balancing method is tested.
+
+
+# Cluster setup which is shared for testing both load balancing methods
+my $can_bind_to_127_0_0_2 = $Config{osname} eq 'linux' || $PostgreSQL::Test::Utils::windows_os;
+
+if ($can_bind_to_127_0_0_2)
+{
+	$PostgreSQL::Test::Cluster::use_tcp = 1;
+	$PostgreSQL::Test::Cluster::test_pghost = '127.0.0.1';
+}
+my $port = PostgreSQL::Test::Cluster::get_free_port();
+my $node1 = PostgreSQL::Test::Cluster->new('node1', port => $port);
+my $node2 = PostgreSQL::Test::Cluster->new('node2', port => $port, own_host => 1);
+my $node3 = PostgreSQL::Test::Cluster->new('node3', port => $port, own_host => 1);
+
+# Create a data directory with initdb
+$node1->init();
+$node2->init();
+$node3->init();
+
+# Start the PostgreSQL server
+$node1->start();
+$node2->start();
+$node3->start();
+
+# Start the tests for load balancing method 1
+my $hostlist = $node1->host . ',' . $node2->host . ',' . $node3->host;
+my $portlist = "$port,$port,$port";
+
+$node1->connect_ok("host=$hostlist port=$portlist load_balance_hosts=random random_seed=123",
+	"seed 123 selects node 1 first",
+	sql => "SELECT 'connect1'",
+	log_like => [qr/statement: SELECT 'connect1'/]);
+
+$node2->connect_ok("host=$hostlist port=$portlist load_balance_hosts=random random_seed=123",
+	"seed 123 does not select node 2 first",
+	sql => "SELECT 'connect1'",
+	log_unlike => [qr/statement: SELECT 'connect1'/]);
+
+$node3->connect_ok("host=$hostlist port=$portlist load_balance_hosts=random random_seed=123",
+	"seed 123 does not select node 3 first",
+	sql => "SELECT 'connect1'",
+	log_unlike => [qr/statement: SELECT 'connect1'/]);
+
+$node3->connect_ok("host=$hostlist port=$portlist load_balance_hosts=random random_seed=42",
+	"seed 42 selects node 3 first",
+	sql => "SELECT 'connect2'",
+	log_like => [qr/statement: SELECT 'connect2'/]);
+
+$node1->connect_ok("host=$hostlist port=$portlist load_balance_hosts=random random_seed=42",
+	"seed 42 does not select node 1 first",
+	sql => "SELECT 'connect2'",
+	log_unlike => [qr/statement: SELECT 'connect2'/]);
+
+$node2->connect_ok("host=$hostlist port=$portlist load_balance_hosts=random random_seed=42",
+	"seed 42 does not select node 2 first",
+	sql => "SELECT 'connect2'",
+	log_unlike => [qr/statement: SELECT 'connect2'/]);
+
+$node3->stop();
+
+$node1->connect_ok("host=$hostlist port=$portlist load_balance_hosts=random random_seed=42",
+	"seed 42 does select node 1 second",
+	sql => "SELECT 'connect3'",
+	log_like => [qr/statement: SELECT 'connect3'/]);
+
+$node2->connect_ok("host=$hostlist port=$portlist load_balance_hosts=random random_seed=42",
+	"seed 42 does not select node 2 second",
+	sql => "SELECT 'connect3'",
+	log_unlike => [qr/statement: SELECT 'connect3'/]);
+
+$node3->start();
+
+# Checks for the requirements for testing load balancing method 2
+if (!$can_bind_to_127_0_0_2) {
+	# The OS requirement is not met
+	done_testing();
+	exit;
+}
+
+my $hosts_path;
+if ($windows_os) {
+	$hosts_path = 'c:\Windows\System32\Drivers\etc\hosts';
+}
+else
+{
+	$hosts_path = '/etc/hosts';
+}
+
+my $hosts_content = PostgreSQL::Test::Utils::slurp_file($hosts_path);
+
+if ($hosts_content !~ m/pg-loadbalancetest/) {
+	# Host file is not prepared for this test
+	done_testing();
+	exit;
+}
+
+# Start the tests for load balancing method 2
+$node2->connect_ok("host=pg-loadbalancetest port=$port load_balance_hosts=random random_seed=44",
+	"seed 44 selects node 2 first",
+	sql => "SELECT 'connect4'",
+	log_like => [qr/statement: SELECT 'connect4'/]);
+
+$node1->connect_ok("host=pg-loadbalancetest port=$port load_balance_hosts=random random_seed=44",
+	"seed 44 does not select node 1 first",
+	sql => "SELECT 'connect4'",
+	log_unlike => [qr/statement: SELECT 'connect4'/]);
+
+$node3->connect_ok("host=pg-loadbalancetest port=$port load_balance_hosts=random random_seed=44",
+	"seed 44 does not select node 3 first",
+	sql => "SELECT 'connect4'",
+	log_unlike => [qr/statement: SELECT 'connect4'/]);
+
+$node2->stop();
+
+$node1->connect_ok("host=pg-loadbalancetest port=$port load_balance_hosts=random random_seed=44",
+	"seed 44 does select node 1 second",
+	sql => "SELECT 'connect5'",
+	log_like => [qr/statement: SELECT 'connect5'/]);
+
+$node3->connect_ok("host=pg-loadbalancetest port=$port load_balance_hosts=random random_seed=44",
+	"seed 44 does not select node 3 second",
+	sql => "SELECT 'connect5'",
+	log_unlike => [qr/statement: SELECT 'connect5'/]);
+
+done_testing();
-- 
2.34.1