From 94b2daf9e053e5e1effa3ae6a610d99d3e43d8b1 Mon Sep 17 00:00:00 2001 From: Hongxiang Jiang Date: Tue, 4 Nov 2025 14:09:04 -0500 Subject: 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 kokoro-CI: kokoro Reviewed-by: Nicholas Husin LUCI-TryBot-Result: Go LUCI Reviewed-by: Michael Pratt --- doc/config.md | 5 ++- internal/config/serverconfig/config.go | 32 ++++++++++++------ internal/config/serverconfig/config_test.go | 51 +++++++++++++++++++++-------- 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) } } } -- cgit v1.3