@@ -123,6 +123,7 @@ static int ldapServiceLookup(const char *purl, PQconninfoOption *options,
123123#define DefaultChannelBinding "disable"
124124#endif
125125#define DefaultTargetSessionAttrs "any"
126+ #define DefaultLoadBalanceHosts "disable"
126127#ifdef USE_SSL
127128#define DefaultSSLMode "prefer"
128129#define DefaultSSLCertMode "allow"
@@ -351,6 +352,11 @@ static const internalPQconninfoOption PQconninfoOptions[] = {
351352"Target-Session-Attrs" , "" , 15 , /* sizeof("prefer-standby") = 15 */
352353offsetof(struct pg_conn , target_session_attrs )},
353354
355+ {"load_balance_hosts" , "PGLOADBALANCEHOSTS" ,
356+ DefaultLoadBalanceHosts , NULL ,
357+ "Load-Balance-Hosts" , "" , 8 , /* sizeof("disable") = 8 */
358+ offsetof(struct pg_conn , load_balance_hosts )},
359+
354360/* Terminating entry --- MUST BE LAST */
355361{NULL , NULL , NULL , NULL ,
356362NULL , NULL , 0 }
@@ -435,6 +441,8 @@ static void pgpassfileWarning(PGconn *conn);
435441static void default_threadlock (int acquire );
436442static bool sslVerifyProtocolVersion (const char * version );
437443static bool sslVerifyProtocolRange (const char * min , const char * max );
444+ static bool parse_int_param (const char * value , int * result , PGconn * conn ,
445+ const char * context );
438446
439447
440448/* global variable because fe-auth.c needs to access it */
@@ -1020,6 +1028,31 @@ parse_comma_separated_list(char **startptr, bool *more)
10201028return p ;
10211029}
10221030
1031+ /*
1032+ * Initializes the prng_state field of the connection. We want something
1033+ * unpredictable, so if possible, use high-quality random bits for the
1034+ * seed. Otherwise, fall back to a seed based on the connection address,
1035+ * timestamp and PID.
1036+ */
1037+ static void
1038+ libpq_prng_init (PGconn * conn )
1039+ {
1040+ uint64 rseed ;
1041+ struct timeval tval = {0 };
1042+
1043+ if (pg_prng_strong_seed (& conn -> prng_state ))
1044+ return ;
1045+
1046+ gettimeofday (& tval , NULL );
1047+
1048+ rseed = ((uint64 ) conn ) ^
1049+ ((uint64 ) getpid ()) ^
1050+ ((uint64 ) tval .tv_usec ) ^
1051+ ((uint64 ) tval .tv_sec );
1052+
1053+ pg_prng_seed (& conn -> prng_state , rseed );
1054+ }
1055+
10231056/*
10241057* connectOptions2
10251058*
@@ -1619,6 +1652,49 @@ connectOptions2(PGconn *conn)
16191652else
16201653conn -> target_server_type = SERVER_TYPE_ANY ;
16211654
1655+ /*
1656+ * validate load_balance_hosts option, and set load_balance_type
1657+ */
1658+ if (conn -> load_balance_hosts )
1659+ {
1660+ if (strcmp (conn -> load_balance_hosts , "disable" ) == 0 )
1661+ conn -> load_balance_type = LOAD_BALANCE_DISABLE ;
1662+ else if (strcmp (conn -> load_balance_hosts , "random" ) == 0 )
1663+ conn -> load_balance_type = LOAD_BALANCE_RANDOM ;
1664+ else
1665+ {
1666+ conn -> status = CONNECTION_BAD ;
1667+ libpq_append_conn_error (conn , "invalid %s value: \"%s\"" ,
1668+ "load_balance_hosts" ,
1669+ conn -> load_balance_hosts );
1670+ return false;
1671+ }
1672+ }
1673+ else
1674+ conn -> load_balance_type = LOAD_BALANCE_DISABLE ;
1675+
1676+ if (conn -> load_balance_type == LOAD_BALANCE_RANDOM )
1677+ {
1678+ libpq_prng_init (conn );
1679+
1680+ /*
1681+ * This is the "inside-out" variant of the Fisher-Yates shuffle
1682+ * algorithm. Notionally, we append each new value to the array and
1683+ * then swap it with a randomly-chosen array element (possibly
1684+ * including itself, else we fail to generate permutations with the
1685+ * last integer last). The swap step can be optimized by combining it
1686+ * with the insertion.
1687+ */
1688+ for (i = 1 ; i < conn -> nconnhost ; i ++ )
1689+ {
1690+ int j = pg_prng_uint64_range (& conn -> prng_state , 0 , i );
1691+ pg_conn_host temp = conn -> connhost [j ];
1692+
1693+ conn -> connhost [j ] = conn -> connhost [i ];
1694+ conn -> connhost [i ] = temp ;
1695+ }
1696+ }
1697+
16221698/*
16231699* Resolve special "auto" client_encoding from the locale
16241700*/
@@ -2626,6 +2702,32 @@ PQconnectPoll(PGconn *conn)
26262702if (ret )
26272703goto error_return ; /* message already logged */
26282704
2705+ /*
2706+ * If random load balancing is enabled we shuffle the addresses.
2707+ */
2708+ if (conn -> load_balance_type == LOAD_BALANCE_RANDOM )
2709+ {
2710+ /*
2711+ * This is the "inside-out" variant of the Fisher-Yates shuffle
2712+ * algorithm. Notionally, we append each new value to the array
2713+ * and then swap it with a randomly-chosen array element (possibly
2714+ * including itself, else we fail to generate permutations with
2715+ * the last integer last). The swap step can be optimized by
2716+ * combining it with the insertion.
2717+ *
2718+ * We don't need to initialize conn->prng_state here, because that
2719+ * already happened in connectOptions2.
2720+ */
2721+ for (int i = 1 ; i < conn -> naddr ; i ++ )
2722+ {
2723+ int j = pg_prng_uint64_range (& conn -> prng_state , 0 , i );
2724+ AddrInfo temp = conn -> addr [j ];
2725+
2726+ conn -> addr [j ] = conn -> addr [i ];
2727+ conn -> addr [i ] = temp ;
2728+ }
2729+ }
2730+
26292731reset_connection_state_machine = true;
26302732conn -> try_next_host = false;
26312733}
@@ -4320,6 +4422,7 @@ freePGconn(PGconn *conn)
43204422free (conn -> outBuffer );
43214423free (conn -> rowBuf );
43224424free (conn -> target_session_attrs );
4425+ free (conn -> load_balance_hosts );
43234426termPQExpBuffer (& conn -> errorMessage );
43244427termPQExpBuffer (& conn -> workBuffer );
43254428
0 commit comments