internal/config/serverconfig: randomly pick two DBs from same env

Primary and secondary DB will be derived from a single env var
GO_DISCOVERY_DATABASE_HOST.

The env GO_DISCOVERY_DATABASE_SECONDARY_HOST will be obsolete.

Change-Id: I4c5276c1fef7c93f9abbe5878401f22bdbab69dd
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/717800
Reviewed-by: Nicholas Husin <nsh@golang.org>
kokoro-CI: kokoro <noreply+kokoro@google.com>
Reviewed-by: Nicholas Husin <husin@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Michael Pratt <mpratt@google.com>

3 files changed, 62 insertions, 26 deletions

diff --git a/doc/config.md b/doc/config.md
index 12a2751b..d45bf70b 100644
--- a/doc/config.md
+++ b/doc/config.md
@@ -7,10 +7,9 @@ Pkgsite uses these environment variables:
 | GO_DISCOVERY_AUTH_VALUES             | Set of values that could be set on the AuthHeader, in order to bypass checks by the cache.                                                                                                                                                                                                                                         |
 | GO_DISCOVERY_CONFIG_BUCKET           | Bucket use for dynamic configuration (gs://bucket/object) GO_DISCOVERY_CONFIG_DYNAMIC must be set if GO_DISCOVERY_CONFIG_BUCKET is set.                                                                                                                                                                                            |
 | GO_DISCOVERY_CONFIG_DYNAMIC          | File that experiments are read from. Can be set locally using devtools/cmd/create_experiment_config/main.go.                                                                                                                                                                                                                       |
-| GO_DISCOVERY_DATABASE_HOST           | Database server hostname.                                                                                                                                                                                                                                                                                                          |
+| GO_DISCOVERY_DATABASE_HOST           | Database server hostname. A primary and a secondary hosts are picked randomly. If the primary is not reachable, the system will fall back to the secondary host.                                                                                                                                                                                                                                                                                                          |
 | GO_DISCOVERY_DATABASE_NAME           | Name of database within the server.                                                                                                                                                                                                                                                                                                |
-| GO_DISCOVERY_DATABASE_PASSWORD       | Password for database.                                                                                                                                                                                                                                                                                                             |
-| GO_DISCOVERY_DATABASE_SECONDARY_HOST | If `GO_DISCOVERY_DATABASE_HOST` is unreachable, use this host. Used only by prod frontends.                                                                                                                                                                                                                                        |
+| GO_DISCOVERY_DATABASE_PASSWORD       | Password for database.                                                                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                     |
 | GO_DISCOVERY_DATABASE_USER           | Used for frontend, worker and scripts.                                                                                                                                                                                                                                                                                             |
 | GO_DISCOVERY_DISABLE_ERROR_REPORTING | Disables calls to GCP errorreporting API. Set only in dev.                                                                                                                                                                                                                                                                         |
 | GO_DISCOVERY_E2E_AUTHORIZATION       | Auth token for e2e tests.                                                                                                                                                                                                                                                                                                          |
diff --git a/internal/config/serverconfig/config.go b/internal/config/serverconfig/config.go
index 643453b3..bd98e465 100644
--- a/internal/config/serverconfig/config.go
+++ b/internal/config/serverconfig/config.go
@@ -115,6 +115,10 @@ type configOverride struct {
 // must be called before any configuration values are used.
 func Init(ctx context.Context) (_ *config.Config, err error) {
 	defer derrors.Add(&err, "config.Init(ctx)")
+
+	dbs := chooseN(GetEnv("GO_DISCOVERY_DATABASE_HOST", "localhost"), 2)
+	primaryDB, secondaryDB := dbs[0], dbs[1]
+
 	// Build a Config from the execution environment, loading some values
 	// from envvars and others from remote services.
 	cfg := &config.Config{
@@ -141,10 +145,10 @@ func Init(ctx context.Context) (_ *config.Config, err error) {
 		LocationID: GetEnv("GO_DISCOVERY_GAE_LOCATION_ID", "us-central1"),
 		// This fallback should only be used when developing locally.
 		FallbackVersionLabel: time.Now().Format(config.AppVersionFormat),
-		DBHost:               chooseOne(GetEnv("GO_DISCOVERY_DATABASE_HOST", "localhost")),
+		DBHost:               primaryDB,
 		DBUser:               GetEnv("GO_DISCOVERY_DATABASE_USER", "postgres"),
 		DBPassword:           os.Getenv("GO_DISCOVERY_DATABASE_PASSWORD"),
-		DBSecondaryHost:      chooseOne(os.Getenv("GO_DISCOVERY_DATABASE_SECONDARY_HOST")),
+		DBSecondaryHost:      secondaryDB,
 		DBPort:               GetEnv("GO_DISCOVERY_DATABASE_PORT", "5432"),
 		DBName:               GetEnv("GO_DISCOVERY_DATABASE_NAME", "discovery-db"),
 		DBSecret:             os.Getenv("GO_DISCOVERY_DATABASE_SECRET"),
@@ -333,16 +337,24 @@ func override[T comparable](ctx context.Context, name string, field *T, val T) {
 	}
 }
 
-// chooseOne selects one entry at random from a whitespace-separated
-// string. It returns the empty string if there are no elements.
-func chooseOne(configVar string) string {
+// chooseN selects N entries at random from a whitespace-separated string, and
+// returns them as a slice of size N.
+//
+// If the input string contains fewer than N entries, the returned slice is
+// padded with empty strings at the end.
+func chooseN(configVar string, n int) []string {
 	fields := strings.Fields(configVar)
-	if len(fields) == 0 {
-		return ""
+
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	r.Shuffle(len(fields), func(i, j int) {
+		fields[i], fields[j] = fields[j], fields[i]
+	})
+
+	if len(fields) <= n {
+		return append(fields, make([]string, n-len(fields))...)
+	} else {
+		return fields[:n]
 	}
-	src := rand.NewSource(time.Now().UnixNano())
-	rng := rand.New(src)
-	return fields[rng.Intn(len(fields))]
 }
 
 // gceMetadata reads a metadata value from GCE.
diff --git a/internal/config/serverconfig/config_test.go b/internal/config/serverconfig/config_test.go
index c45f6a0c..a4899205 100644
--- a/internal/config/serverconfig/config_test.go
+++ b/internal/config/serverconfig/config_test.go
@@ -32,24 +32,49 @@ func TestValidateAppVersion(t *testing.T) {
 	}
 }
 
-func TestChooseOne(t *testing.T) {
+func TestChooseN(t *testing.T) {
 	tests := []struct {
-		configVar   string
-		wantMatches string
+		configVar string
+		n         int
+		wantMatch []string
 	}{
-		{"foo", "foo"},
-		{"foo1 \n foo2", "^foo[12]$"},
-		{"foo1\nfoo2", "^foo[12]$"},
-		{"foo1 foo2", "^foo[12]$"},
+		{"foo", 2, []string{"foo", ""}},
+		{"foo1 \n foo2", 1, []string{"^foo[12]$"}},
+		{"foo1 \n foo2", 2, []string{"^foo[12]$", "^foo[12]$"}},
+		{"foo1 foo2", 4, []string{"^foo[12]$", "^foo[12]$", "", ""}},
+		{"foo1\nfoo2\nfoo3", 5, []string{"^foo[123]$", "^foo[123]$", "^foo[123]$", "", ""}},
 	}
 	for _, test := range tests {
-		got := chooseOne(test.configVar)
-		matched, err := regexp.MatchString(test.wantMatches, got)
-		if err != nil {
-			t.Fatal(err)
+		gots := chooseN(test.configVar, test.n)
+
+		if len(gots) != test.n {
+			t.Errorf("chooseN must return a slice of n(%v), got %v", test.n, len(gots))
+		}
+		seen := make(map[string]struct{}, test.n)
+
+		allMatch := true
+		allUnique := true
+		for i, got := range gots {
+			if got != "" {
+				_, ok := seen[got]
+				allUnique = allUnique && !ok
+
+				seen[got] = struct{}{}
+			}
+
+			matched, err := regexp.MatchString(test.wantMatch[i], got)
+			if err != nil {
+				t.Fatal(err)
+			}
+			allMatch = allMatch && matched
+
+			seen[got] = struct{}{}
+		}
+		if !allMatch {
+			t.Errorf("chooseN(%q, %v) = %v, want matches %v", test.configVar, test.n, gots, test.wantMatch)
 		}
-		if !matched {
-			t.Errorf("chooseOne(%q) = %q, _, want matches %q", test.configVar, got, test.wantMatches)
+		if !allUnique {
+			t.Errorf("chooseN(%q, %v) = %v, want all unique", test.configVar, test.n, gots)
 		}
 	}
 }