Merge branch 'master' into 1885-keep-proxy refs #1885

[arvados.git] / sdk / go / src / arvados.org / keepclient / support.go

diff --git a/sdk/go/src/arvados.org/keepclient/support.go b/sdk/go/src/arvados.org/keepclient/support.go
index e657a60a39e46d668253436b61c7a191504188cf..ef4a8e1133b8674dc84d696cc01b04ae1e6269b4 100644 (file)
--- a/sdk/go/src/arvados.org/keepclient/support.go
+++ b/sdk/go/src/arvados.org/keepclient/support.go
@@ -9,6 +9,7 @@ import (
        "io"
        "log"
        "net/http"
+       "os"
        "sort"
        "strconv"
 )
@@ -17,18 +18,30 @@ type keepDisk struct {
        Hostname string `json:"service_host"`
        Port     int    `json:"service_port"`
        SSL      bool   `json:"service_ssl_flag"`
+       SvcType  string `json:"service_type"`
 }
 
 func (this *KeepClient) discoverKeepServers() error {
+       if prx := os.Getenv("ARVADOS_KEEP_PROXY"); prx != "" {
+               this.Service_roots = make([]string, 1)
+               this.Service_roots[0] = prx
+               this.Using_proxy = true
+               return nil
+       }
+
        // Construct request of keep disk list
        var req *http.Request
        var err error
-       if req, err = http.NewRequest("GET", fmt.Sprintf("https://%s/arvados/v1/keep_disks", this.ApiServer), nil); err != nil {
+
+       if req, err = http.NewRequest("GET", fmt.Sprintf("https://%s/arvados/v1/keep_services/accessible?format=json", this.ApiServer), nil); err != nil {
                return err
        }
 
        // Add api token header
        req.Header.Add("Authorization", fmt.Sprintf("OAuth2 %s", this.ApiToken))
+       if this.External {
+               req.Header.Add("X-External-Client", "1")
+       }
 
        // Make the request
        var resp *http.Response
@@ -36,6 +49,17 @@ func (this *KeepClient) discoverKeepServers() error {
                return err
        }
 
+       if resp.StatusCode != 200 {
+               // fall back on keep disks
+               if req, err = http.NewRequest("GET", fmt.Sprintf("https://%s/arvados/v1/keep_disks", this.ApiServer), nil); err != nil {
+                       return err
+               }
+               req.Header.Add("Authorization", fmt.Sprintf("OAuth2 %s", this.ApiToken))
+               if resp, err = this.Client.Do(req); err != nil {
+                       return err
+               }
+       }
+
        type svcList struct {
                Items []keepDisk `json:"items"`
        }
@@ -52,6 +76,7 @@ func (this *KeepClient) discoverKeepServers() error {
 
        for _, element := range m.Items {
                n := ""
+
                if element.SSL {
                        n = "s"
                }
@@ -64,6 +89,9 @@ func (this *KeepClient) discoverKeepServers() error {
                        listed[url] = true
                        this.Service_roots = append(this.Service_roots, url)
                }
+               if element.SvcType == "proxy" {
+                       this.Using_proxy = true
+               }
        }
 
        // Must be sorted for ShuffledServiceRoots() to produce consistent
@@ -122,9 +150,10 @@ func (this KeepClient) shuffledServiceRoots(hash string) (pseq []string) {
 }
 
 type uploadStatus struct {
-       Err        error
-       Url        string
-       StatusCode int
+       err             error
+       url             string
+       statusCode      int
+       replicas_stored int
 }
 
 func (this KeepClient) uploadToKeepServer(host string, hash string, body io.ReadCloser,
@@ -136,7 +165,7 @@ func (this KeepClient) uploadToKeepServer(host string, hash string, body io.Read
        var err error
        var url = fmt.Sprintf("%s/%s", host, hash)
        if req, err = http.NewRequest("PUT", url, nil); err != nil {
-               upload_status <- uploadStatus{err, url, 0}
+               upload_status <- uploadStatus{err, url, 0, 0}
                body.Close()
                return
        }
@@ -147,18 +176,29 @@ func (this KeepClient) uploadToKeepServer(host string, hash string, body io.Read
 
        req.Header.Add("Authorization", fmt.Sprintf("OAuth2 %s", this.ApiToken))
        req.Header.Add("Content-Type", "application/octet-stream")
+
+       if this.Using_proxy {
+               req.Header.Add("X-Keep-Desired-Replicas", fmt.Sprint(this.Want_replicas))
+       }
+
        req.Body = body
 
        var resp *http.Response
        if resp, err = this.Client.Do(req); err != nil {
-               upload_status <- uploadStatus{err, url, 0}
+               upload_status <- uploadStatus{err, url, 0, 0}
+               body.Close()
                return
        }
 
+       rep := 1
+       if xr := resp.Header.Get("X-Keep-Replicas-Stored"); xr != "" {
+               fmt.Sscanf(xr, "%d", &rep)
+       }
+
        if resp.StatusCode == http.StatusOK {
-               upload_status <- uploadStatus{nil, url, resp.StatusCode}
+               upload_status <- uploadStatus{nil, url, resp.StatusCode, rep}
        } else {
-               upload_status <- uploadStatus{errors.New(resp.Status), url, resp.StatusCode}
+               upload_status <- uploadStatus{errors.New(resp.Status), url, resp.StatusCode, rep}
        }
 }
 
@@ -181,6 +221,7 @@ func (this KeepClient) putReplicas(
        defer close(upload_status)
 
        // Desired number of replicas
+
        remaining_replicas := this.Want_replicas
 
        for remaining_replicas > 0 {
@@ -191,23 +232,27 @@ func (this KeepClient) putReplicas(
                                next_server += 1
                                active += 1
                        } else {
-                               return (this.Want_replicas - remaining_replicas), InsufficientReplicasError
+                               if active == 0 {
+                                       return (this.Want_replicas - remaining_replicas), InsufficientReplicasError
+                               } else {
+                                       break
+                               }
                        }
                }
 
                // Now wait for something to happen.
                status := <-upload_status
-               if status.StatusCode == 200 {
+               if status.statusCode == 200 {
                        // good news!
-                       remaining_replicas -= 1
+                       remaining_replicas -= status.replicas_stored
                } else {
                        // writing to keep server failed for some reason
                        log.Printf("Keep server put to %v failed with '%v'",
-                               status.Url, status.Err)
+                               status.url, status.err)
                }
                active -= 1
-               log.Printf("Upload status %v %v %v", status.StatusCode, remaining_replicas, active)
+               log.Printf("Upload status code: %v remaining replicas: %v active: %v", status.statusCode, remaining_replicas, active)
        }
 
-       return (this.Want_replicas - remaining_replicas), nil
+       return this.Want_replicas, nil
 }