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