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