sdk/go/src/arvados.org/keepclient/support.go

   1 /* Internal methods to support keepclient.go */
   2 package keepclient
   3
   4 import (
   5         "arvados.org/streamer"
   6         "encoding/json"
   7         "errors"
   8         "fmt"
   9         "io"
  10         "log"
  11         "net/http"
  12         "os"
  13         "strconv"
  14 )
  15
  16 type keepDisk struct {
  17         Hostname string `json:"service_host"`
  18         Port     int    `json:"service_port"`
  19         SSL      bool   `json:"service_ssl_flag"`
  20         SvcType  string `json:"service_type"`
  21 }
  22
  23 func (this *KeepClient) DiscoverKeepServers() error {
  24         if prx := os.Getenv("ARVADOS_KEEP_PROXY"); prx != "" {
  25                 this.SetServiceRoots([]string{prx})
  26                 this.Using_proxy = true
  27                 return nil
  28         }
  29
  30         // Construct request of keep disk list
  31         var req *http.Request
  32         var err error
  33
  34         if req, err = http.NewRequest("GET", fmt.Sprintf("https://%s/arvados/v1/keep_services/accessible?format=json", this.ApiServer), nil); err != nil {
  35                 return err
  36         }
  37
  38         // Add api token header
  39         req.Header.Add("Authorization", fmt.Sprintf("OAuth2 %s", this.ApiToken))
  40         if this.External {
  41                 req.Header.Add("X-External-Client", "1")
  42         }
  43
  44         // Make the request
  45         var resp *http.Response
  46         if resp, err = this.Client.Do(req); err != nil {
  47                 return err
  48         }
  49
  50         if resp.StatusCode != http.StatusOK {
  51                 // fall back on keep disks
  52                 if req, err = http.NewRequest("GET", fmt.Sprintf("https://%s/arvados/v1/keep_disks", this.ApiServer), nil); err != nil {
  53                         return err
  54                 }
  55                 req.Header.Add("Authorization", fmt.Sprintf("OAuth2 %s", this.ApiToken))
  56                 if resp, err = this.Client.Do(req); err != nil {
  57                         return err
  58                 }
  59                 if resp.StatusCode != http.StatusOK {
  60                         return errors.New(resp.Status)
  61                 }
  62         }
  63
  64         type svcList struct {
  65                 Items []keepDisk `json:"items"`
  66         }
  67
  68         // Decode json reply
  69         dec := json.NewDecoder(resp.Body)
  70         var m svcList
  71         if err := dec.Decode(&m); err != nil {
  72                 return err
  73         }
  74
  75         listed := make(map[string]bool)
  76         service_roots := make([]string, 0, len(m.Items))
  77
  78         for _, element := range m.Items {
  79                 n := ""
  80
  81                 if element.SSL {
  82                         n = "s"
  83                 }
  84
  85                 // Construct server URL
  86                 url := fmt.Sprintf("http%s://%s:%d", n, element.Hostname, element.Port)
  87
  88                 // Skip duplicates
  89                 if !listed[url] {
  90                         listed[url] = true
  91                         service_roots = append(service_roots, url)
  92                 }
  93                 if element.SvcType == "proxy" {
  94                         this.Using_proxy = true
  95                 }
  96         }
  97
  98         this.SetServiceRoots(service_roots)
  99
 100         return nil
 101 }
 102
 103 func (this KeepClient) shuffledServiceRoots(hash string) (pseq []string) {
 104         // Build an ordering with which to query the Keep servers based on the
 105         // contents of the hash.  "hash" is a hex-encoded number at least 8
 106         // digits (32 bits) long
 107
 108         // seed used to calculate the next keep server from 'pool' to be added
 109         // to 'pseq'
 110         seed := hash
 111
 112         // Keep servers still to be added to the ordering
 113         service_roots := this.ServiceRoots()
 114         pool := make([]string, len(service_roots))
 115         copy(pool, service_roots)
 116
 117         // output probe sequence
 118         pseq = make([]string, 0, len(service_roots))
 119
 120         // iterate while there are servers left to be assigned
 121         for len(pool) > 0 {
 122
 123                 if len(seed) < 8 {
 124                         // ran out of digits in the seed
 125                         if len(pseq) < (len(hash) / 4) {
 126                                 // the number of servers added to the probe
 127                                 // sequence is less than the number of 4-digit
 128                                 // slices in 'hash' so refill the seed with the
 129                                 // last 4 digits.
 130                                 seed = hash[len(hash)-4:]
 131                         }
 132                         seed += hash
 133                 }
 134
 135                 // Take the next 8 digits (32 bytes) and interpret as an integer,
 136                 // then modulus with the size of the remaining pool to get the next
 137                 // selected server.
 138                 probe, _ := strconv.ParseUint(seed[0:8], 16, 32)
 139                 probe %= uint64(len(pool))
 140
 141                 // Append the selected server to the probe sequence and remove it
 142                 // from the pool.
 143                 pseq = append(pseq, pool[probe])
 144                 pool = append(pool[:probe], pool[probe+1:]...)
 145
 146                 // Remove the digits just used from the seed
 147                 seed = seed[8:]
 148         }
 149         return pseq
 150 }
 151
 152 type uploadStatus struct {
 153         err             error
 154         url             string
 155         statusCode      int
 156         replicas_stored int
 157 }
 158
 159 func (this KeepClient) uploadToKeepServer(host string, hash string, body io.ReadCloser,
 160         upload_status chan<- uploadStatus, expectedLength int64) {
 161
 162         log.Printf("Uploading %s to %s", hash, host)
 163
 164         var req *http.Request
 165         var err error
 166         var url = fmt.Sprintf("%s/%s", host, hash)
 167         if req, err = http.NewRequest("PUT", url, nil); err != nil {
 168                 upload_status <- uploadStatus{err, url, 0, 0}
 169                 body.Close()
 170                 return
 171         }
 172
 173         if expectedLength > 0 {
 174                 req.ContentLength = expectedLength
 175         }
 176
 177         req.Header.Add("Authorization", fmt.Sprintf("OAuth2 %s", this.ApiToken))
 178         req.Header.Add("Content-Type", "application/octet-stream")
 179
 180         if this.Using_proxy {
 181                 req.Header.Add(X_Keep_Desired_Replicas, fmt.Sprint(this.Want_replicas))
 182         }
 183
 184         req.Body = body
 185
 186         var resp *http.Response
 187         if resp, err = this.Client.Do(req); err != nil {
 188                 upload_status <- uploadStatus{err, url, 0, 0}
 189                 body.Close()
 190                 return
 191         }
 192
 193         rep := 1
 194         if xr := resp.Header.Get(X_Keep_Replicas_Stored); xr != "" {
 195                 fmt.Sscanf(xr, "%d", &rep)
 196         }
 197
 198         if resp.StatusCode == http.StatusOK {
 199                 upload_status <- uploadStatus{nil, url, resp.StatusCode, rep}
 200         } else {
 201                 upload_status <- uploadStatus{errors.New(resp.Status), url, resp.StatusCode, rep}
 202         }
 203 }
 204
 205 func (this KeepClient) putReplicas(
 206         hash string,
 207         tr *streamer.AsyncStream,
 208         expectedLength int64) (replicas int, err error) {
 209
 210         // Calculate the ordering for uploading to servers
 211         sv := this.shuffledServiceRoots(hash)
 212
 213         // The next server to try contacting
 214         next_server := 0
 215
 216         // The number of active writers
 217         active := 0
 218
 219         // Used to communicate status from the upload goroutines
 220         upload_status := make(chan uploadStatus)
 221         defer close(upload_status)
 222
 223         // Desired number of replicas
 224
 225         remaining_replicas := this.Want_replicas
 226
 227         for remaining_replicas > 0 {
 228                 for active < remaining_replicas {
 229                         // Start some upload requests
 230                         if next_server < len(sv) {
 231                                 go this.uploadToKeepServer(sv[next_server], hash, tr.MakeStreamReader(), upload_status, expectedLength)
 232                                 next_server += 1
 233                                 active += 1
 234                         } else {
 235                                 if active == 0 {
 236                                         return (this.Want_replicas - remaining_replicas), InsufficientReplicasError
 237                                 } else {
 238                                         break
 239                                 }
 240                         }
 241                 }
 242
 243                 // Now wait for something to happen.
 244                 status := <-upload_status
 245                 if status.statusCode == 200 {
 246                         // good news!
 247                         remaining_replicas -= status.replicas_stored
 248                 } else {
 249                         // writing to keep server failed for some reason
 250                         log.Printf("Keep server put to %v failed with '%v'",
 251                                 status.url, status.err)
 252                 }
 253                 active -= 1
 254                 log.Printf("Upload to %v status code: %v remaining replicas: %v active: %v", status.url, status.statusCode, remaining_replicas, active)
 255         }
 256
 257         return this.Want_replicas, nil
 258 }