9550: Allow overriding keep services discovery with ARVADOS_KEEP_SERVICES env var.

diff --git a/sdk/go/arvadosclient/arvadosclient.go b/sdk/go/arvadosclient/arvadosclient.go
index 8cdfa484bd96df3f479c7a40c8bb6692eddc1da5..2befcd55106c7aef184f8916415d542ba8910d1d 100644 (file)
--- a/sdk/go/arvadosclient/arvadosclient.go
+++ b/sdk/go/arvadosclient/arvadosclient.go
@@ -86,6 +86,10 @@ type ArvadosClient struct {
        // the client is outside the cluster.
        External bool
 
+       // Use provided list of keep services, instead of using the
+       // API to discover available services.
+       KeepServiceURIs []string
+
        // Discovery document
        DiscoveryDoc Dict
 
@@ -113,6 +117,10 @@ func MakeArvadosClient() (ac ArvadosClient, err error) {
                External: external,
                Retries:  2}
 
+       if s := os.Getenv("ARVADOS_KEEP_SERVICES"); s != "" {
+               ac.KeepServiceURIs = strings.Split(s, " ")
+       }
+
        if ac.ApiServer == "" {
                return ac, MissingArvadosApiHost
        }

diff --git a/sdk/go/keepclient/discover.go b/sdk/go/keepclient/discover.go
index f039c2181055543bfb0a85ee14710ddd0cf5d1ad..66eb3166656fb42be75a534b7098d5241444515e 100644 (file)
--- a/sdk/go/keepclient/discover.go
+++ b/sdk/go/keepclient/discover.go
@@ -12,8 +12,25 @@ import (
        "time"
 )
 
-// DiscoverKeepServers gets list of available keep services from api server
+// DiscoverKeepServers gets list of available keep services from the
+// API server.
+//
+// If a list of services is provided in the arvadosclient (e.g., from
+// an environment variable or local config), that list is used
+// instead.
 func (this *KeepClient) DiscoverKeepServers() error {
+       if this.Arvados.KeepServiceURIs != nil {
+               this.foundNonDiskSvc = true
+               this.replicasPerService = 0
+               this.setClientSettingsNonDisk()
+               roots := make(map[string]string)
+               for i, uri := range this.Arvados.KeepServiceURIs {
+                       roots[fmt.Sprintf("00000-bi6l4-%015d", i)] = uri
+               }
+               this.SetServiceRoots(roots, roots, roots)
+               return nil
+       }
+
        var list svcList
 
        // Get keep services from api server

diff --git a/sdk/go/keepclient/discover_test.go b/sdk/go/keepclient/discover_test.go
index 43a984e8c42f8696e116c15a2471325d20d6e332..ab7b2da3aa423230453619d6a6f0281ec49f0506 100644 (file)
--- a/sdk/go/keepclient/discover_test.go
+++ b/sdk/go/keepclient/discover_test.go
@@ -1,10 +1,15 @@
 package keepclient
 
 import (
+       "crypto/md5"
        "fmt"
+       "gopkg.in/check.v1"
+       "net/http"
+       "os"
        "time"
 
        "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
+       "git.curoverse.com/arvados.git/sdk/go/arvadostest"
 )
 
 func ExampleRefreshServices() {
@@ -19,3 +24,42 @@ func ExampleRefreshServices() {
        go kc.RefreshServices(5*time.Minute, 3*time.Second)
        fmt.Printf("LocalRoots: %#v\n", kc.LocalRoots())
 }
+
+func (s *ServerRequiredSuite) TestOverrideDiscovery(c *check.C) {
+       defer os.Setenv("ARVADOS_KEEP_SERVICES", "")
+
+       hash := fmt.Sprintf("%x+3", md5.Sum([]byte("TestOverrideDiscovery")))
+       st := StubGetHandler{
+               c,
+               hash,
+               arvadostest.ActiveToken,
+               http.StatusOK,
+               []byte("TestOverrideDiscovery")}
+       ks := RunSomeFakeKeepServers(st, 2)
+
+       os.Setenv("ARVADOS_KEEP_SERVICES", "")
+       arv1, err := arvadosclient.MakeArvadosClient()
+       c.Assert(err, check.IsNil)
+       arv1.ApiToken = arvadostest.ActiveToken
+
+       os.Setenv("ARVADOS_KEEP_SERVICES", ks[0].url+" "+ks[1].url)
+       arv2, err := arvadosclient.MakeArvadosClient()
+       c.Assert(err, check.IsNil)
+       arv2.ApiToken = arvadostest.ActiveToken
+
+       // ARVADOS_KEEP_SERVICES was empty when we created arv1, but
+       // it pointed to our stub servers when we created
+       // arv2. Regardless of what it's set to now, a keepclient for
+       // arv2 should use our stub servers, but one created for arv1
+       // should not.
+
+       kc1, err := MakeKeepClient(&arv1)
+       c.Assert(err, check.IsNil)
+       kc2, err := MakeKeepClient(&arv2)
+       c.Assert(err, check.IsNil)
+
+       _, _, _, err = kc1.Get(hash)
+       c.Check(err, check.NotNil)
+       _, _, _, err = kc2.Get(hash)
+       c.Check(err, check.IsNil)
+}

diff --git a/sdk/go/keepclient/keepclient.go b/sdk/go/keepclient/keepclient.go
index 26aa7177e05c3aa64d3e956f183379614093d341..58f3ffb8348ff7b5f9d9588e6455ae7c9e9ff18a 100644 (file)
--- a/sdk/go/keepclient/keepclient.go
+++ b/sdk/go/keepclient/keepclient.go
@@ -352,7 +352,7 @@ func (kc *KeepClient) WritableLocalRoots() map[string]string {
 // caller can reuse/modify them after SetServiceRoots returns, but
 // they should not be modified by any other goroutine while
 // SetServiceRoots is running.
-func (kc *KeepClient) SetServiceRoots(newLocals, newWritableLocals map[string]string, newGateways map[string]string) {
+func (kc *KeepClient) SetServiceRoots(newLocals, newWritableLocals, newGateways map[string]string) {
        locals := make(map[string]string)
        for uuid, root := range newLocals {
                locals[uuid] = root