Skip to content

Commit 589ad43

Browse files
arp242arp242
committed
Implement load_balance_hosts=random
1 parent 9541c34 commit 589ad43

3 files changed

Lines changed: 80 additions & 15 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ newer. Previously PostgreSQL 8.4 and newer were supported.
3636
- Support `hostaddr` and `$PGHOSTADDR` ([#1243]).
3737

3838
- Support multiple values in `host`, `port`, and `hostaddr`, which are each
39-
tried in order ([#1246]).
39+
tried in order, or randomly if `load_balance_hosts=random` is set ([#1246]).
4040

4141
- Support `target_session_attrs` connection parameter ([#1246]).
4242

‎connector.go‎

Lines changed: 56 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"context"
55
"database/sql/driver"
66
"fmt"
7+
"math/rand"
78
"net"
89
"net/netip"
910
neturl "net/url"
@@ -29,6 +30,9 @@ type (
2930

3031
// TargetSessionAttrs is a target_session_attrs setting.
3132
TargetSessionAttrs string
33+
34+
// LoadBalanceHosts is a load_balance_hosts setting.
35+
LoadBalanceHosts string
3236
)
3337

3438
// Values for [SSLMode] that pq supports.
@@ -89,6 +93,23 @@ var targetSessionAttrs = []TargetSessionAttrs{TargetSessionAttrsAny,
8993
TargetSessionAttrsReadWrite, TargetSessionAttrsReadOnly, TargetSessionAttrsPrimary,
9094
TargetSessionAttrsStandby, TargetSessionAttrsPreferStandby}
9195

96+
// Values for [LoadBalanceHosts] that pq supports.
97+
const (
98+
// Don't load balance; try hosts in the order in which they're provided.
99+
// This is the default.
100+
LoadBalanceHostsDisable = LoadBalanceHosts("disable")
101+
102+
// Hosts are tried in random order to balance connections across multiple
103+
// PostgreSQL servers.
104+
//
105+
// When using this value it's recommended to also configure a reasonable
106+
// value for connect_timeout. Because then, if one of the nodes that are
107+
// used for load balancing is not responding, a new node will be tried.
108+
LoadBalanceHostsRandom = LoadBalanceHosts("random")
109+
)
110+
111+
var loadBalanceHosts = []LoadBalanceHosts{LoadBalanceHostsDisable, LoadBalanceHostsRandom}
112+
92113
// Connector represents a fixed configuration for the pq driver with a given
93114
// dsn. Connector satisfies the [database/sql/driver.Connector] interface and
94115
// can be used to create any number of DB Conn's via [sql.OpenDB].
@@ -137,9 +158,10 @@ type Config struct {
137158
// for unix domain sockets. Defaults to localhost.
138159
//
139160
// A comma-separated list of host names is also accepted, in which case each
140-
// host name in the list is tried in order; an empty item selects the
141-
// default of localhost. The target_session_attrs option controls properties
142-
// the host must have to be considered acceptable.
161+
// host name in the list is tried in order or randomly if load_balance_hosts
162+
// is set. An empty item selects the default of localhost. The
163+
// target_session_attrs option controls properties the host must have to be
164+
// considered acceptable.
143165
Host string `postgres:"host" env:"PGHOST"`
144166

145167
// IPv4 or IPv6 address to connect to. Using hostaddr allows the application
@@ -161,10 +183,11 @@ type Config struct {
161183
// host name.
162184
//
163185
// A comma-separated list of hostaddr values is also accepted, in which case
164-
// each host in the list is tried in order. An empty item causes the
165-
// corresponding host name to be used, or the default host name if that is
166-
// empty as well. The target_session_attrs option controls properties the
167-
// host must have to be considered acceptable.
186+
// each host in the list is tried in order or randonly if load_balance_hosts
187+
// is set. An empty item causes the corresponding host name to be used, or
188+
// the default host name if that is empty as well. The target_session_attrs
189+
// option controls properties the host must have to be considered
190+
// acceptable.
168191
Hostaddr netip.Addr `postgres:"hostaddr" env:"PGHOSTADDR"`
169192

170193
// The port to connect to. Defaults to 5432.
@@ -264,6 +287,22 @@ type Config struct {
264287
// Default mode for the genetic query optimizer.
265288
Geqo string `postgres:"geqo" env:"PGGEQO"`
266289

290+
// Determine whether the session must have certain properties to be
291+
// acceptable. It's typically used in combination with multiple host names
292+
// to select the first acceptable alternative among several hosts.
293+
TargetSessionAttrs TargetSessionAttrs `postgres:"target_session_attrs" env:"PGTARGETSESSIONATTRS"`
294+
295+
// Controls the order in which the client tries to connect to the available
296+
// hosts. Once a connection attempt is successful no other hosts will be
297+
// tried. This parameter is typically used in combination with multiple host
298+
// names.
299+
//
300+
// This parameter can be used in combination with target_session_attrs to,
301+
// for example, load balance over standby servers only. Once successfully
302+
// connected, subsequent queries on the returned connection will all be sent
303+
// to the same server.
304+
LoadBalanceHosts LoadBalanceHosts `postgres:"load_balance_hosts" env:"PGLOADBALANCEHOSTS"`
305+
267306
// Runtime parameters: any unrecognized parameter in the DSN will be added
268307
// to this and sent to PostgreSQL during startup.
269308
Runtime map[string]string `postgres:"-" env:"-"`
@@ -273,11 +312,6 @@ type Config struct {
273312
// additional ones (if any) are available here.
274313
Multi []ConfigMultihost
275314

276-
// Determine whether the session must have certain properties to be
277-
// acceptable. It's typically used in combination with multiple host names
278-
// to select the first acceptable alternative among several hosts.
279-
TargetSessionAttrs TargetSessionAttrs `postgres:"target_session_attrs" env:"PGTARGETSESSIONATTRS"`
280-
281315
// Record which parameters were given, so we can distinguish between an
282316
// empty string "not given at all".
283317
//
@@ -370,6 +404,11 @@ func (cfg Config) hosts() []Config {
370404
c.Host, c.Hostaddr, c.Port = m.Host, m.Hostaddr, m.Port
371405
cfgs = append(cfgs, c)
372406
}
407+
408+
if cfg.LoadBalanceHosts == LoadBalanceHostsRandom {
409+
rand.Shuffle(len(cfgs), func(i, j int) { cfgs[i], cfgs[j] = cfgs[j], cfgs[i] })
410+
}
411+
373412
return cfgs
374413
}
375414

@@ -475,8 +514,7 @@ func (cfg *Config) fromEnv(env []string) error {
475514
case "PGREQUIREAUTH", "PGCHANNELBINDING", "PGSERVICE", "PGSERVICEFILE", "PGREALM",
476515
"PGSSLCERTMODE", "PGSSLCOMPRESSION", "PGREQUIRESSL", "PGSSLCRL", "PGREQUIREPEER",
477516
"PGSYSCONFDIR", "PGLOCALEDIR", "PGSSLCRLDIR", "PGSSLMINPROTOCOLVERSION", "PGSSLMAXPROTOCOLVERSION",
478-
"PGGSSENCMODE", "PGGSSDELEGATION", "PGLOADBALANCEHOSTS", "PGMINPROTOCOLVERSION",
479-
"PGMAXPROTOCOLVERSION", "PGGSSLIB":
517+
"PGGSSENCMODE", "PGGSSDELEGATION", "PGMINPROTOCOLVERSION", "PGMAXPROTOCOLVERSION", "PGGSSLIB":
480518
return fmt.Errorf("pq: environment variable $%s is not supported", k)
481519
case "PGKRBSRVNAME":
482520
if newGss == nil {
@@ -615,6 +653,7 @@ func (cfg *Config) setFromTag(o map[string]string, tag string) error {
615653
sslmode = (tag == "postgres" && k == "sslmode") || (tag == "env" && k == "PGSSLMODE")
616654
sslnegotiation = (tag == "postgres" && k == "sslnegotiation") || (tag == "env" && k == "PGSSLNEGOTIATION")
617655
targetsessionattrs = (tag == "postgres" && k == "target_session_attrs") || (tag == "env" && k == "PGTARGETSESSIONATTRS")
656+
loadbalancehosts = (tag == "postgres" && k == "load_balance_hosts") || (tag == "env" && k == "PGLOADBALANCEHOSTS")
618657
)
619658
if k == "" || k == "-" {
620659
continue
@@ -664,6 +703,9 @@ func (cfg *Config) setFromTag(o map[string]string, tag string) error {
664703
if targetsessionattrs && !slices.Contains(targetSessionAttrs, TargetSessionAttrs(v)) {
665704
return fmt.Errorf(f+`%q is not supported; supported values are %s`, k, v, pqutil.Join(targetSessionAttrs))
666705
}
706+
if loadbalancehosts && !slices.Contains(loadBalanceHosts, LoadBalanceHosts(v)) {
707+
return fmt.Errorf(f+`%q is not supported; supported values are %s`, k, v, pqutil.Join(loadBalanceHosts))
708+
}
667709
if host {
668710
vv := strings.Split(v, ",")
669711
v = vv[0]

‎connector_test.go‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"net"
1010
"os"
1111
"reflect"
12+
"slices"
1213
"strings"
1314
"testing"
1415
"time"
@@ -608,6 +609,28 @@ func TestConnectMulti(t *testing.T) {
608609
}
609610
})
610611
}
612+
613+
t.Run("load_balance_hosts=random", func(t *testing.T) {
614+
hosts := [3]int{}
615+
for i := 0; i < 25; i++ {
616+
connectedTo = [3]bool{}
617+
db := pqtest.MustDB(t, fmt.Sprintf(
618+
"host=%s,%s,%s port=%s,%s,%s load_balance_hosts=random", f1.Host(), f2.Host(), f3.Host(), f1.Port(), f2.Port(), f3.Port(),
619+
))
620+
err := db.Ping()
621+
if err != nil {
622+
t.Fatal(err)
623+
}
624+
if n := strings.Count(fmt.Sprintf("%v", connectedTo), "true"); n != 1 {
625+
t.Fatal(connectedTo)
626+
}
627+
628+
hosts[slices.Index(connectedTo[:], true)]++
629+
}
630+
if slices.Index(hosts[:], 0) != -1 {
631+
t.Fatal(hosts)
632+
}
633+
})
611634
}
612635

613636
func TestConnectionTargetSessionAttrs(t *testing.T) {

0 commit comments

Comments
 (0)