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

   1 package keepclient
   2
   3 import (
   4         "crypto/md5"
   5         "crypto/tls"
   6         "encoding/json"
   7         "errors"
   8         "fmt"
   9         "io"
  10         "io/ioutil"
  11         "log"
  12         "net/http"
  13         "os"
  14         "sort"
  15         "strconv"
  16 )
  17
  18 // A Keep "block" is 64MB.
  19 const BLOCKSIZE = 64 * 1024 * 1024
  20
  21 type KeepClient struct {
  22         ApiServer     string
  23         ApiToken      string
  24         ApiInsecure   bool
  25         Service_roots []string
  26         Want_replicas int
  27         Client        *http.Client
  28 }
  29
  30 type KeepDisk struct {
  31         Hostname string `json:"service_host"`
  32         Port     int    `json:"service_port"`
  33         SSL      bool   `json:"service_ssl_flag"`
  34 }
  35
  36 func MakeKeepClient() (kc KeepClient, err error) {
  37         tr := &http.Transport{
  38                 TLSClientConfig: &tls.Config{InsecureSkipVerify: kc.ApiInsecure},
  39         }
  40
  41         kc = KeepClient{
  42                 ApiServer:     os.Getenv("ARVADOS_API_HOST"),
  43                 ApiToken:      os.Getenv("ARVADOS_API_TOKEN"),
  44                 ApiInsecure:   (os.Getenv("ARVADOS_API_HOST_INSECURE") != ""),
  45                 Want_replicas: 2,
  46                 Client:        &http.Client{Transport: tr}}
  47
  48         err = (&kc).DiscoverKeepDisks()
  49
  50         return kc, err
  51 }
  52
  53 func (this *KeepClient) DiscoverKeepDisks() error {
  54         // Construct request of keep disk list
  55         var req *http.Request
  56         var err error
  57         if req, err = http.NewRequest("GET", fmt.Sprintf("https://%s/arvados/v1/keep_disks", this.ApiServer), nil); err != nil {
  58                 return err
  59         }
  60
  61         // Add api token header
  62         req.Header.Add("Authorization", fmt.Sprintf("OAuth2 %s", this.ApiToken))
  63
  64         // Make the request
  65         var resp *http.Response
  66         if resp, err = this.Client.Do(req); err != nil {
  67                 return err
  68         }
  69
  70         type SvcList struct {
  71                 Items []KeepDisk `json:"items"`
  72         }
  73
  74         // Decode json reply
  75         dec := json.NewDecoder(resp.Body)
  76         var m SvcList
  77         if err := dec.Decode(&m); err != nil {
  78                 return err
  79         }
  80
  81         listed := make(map[string]bool)
  82         this.Service_roots = make([]string, 0, len(m.Items))
  83
  84         for _, element := range m.Items {
  85                 n := ""
  86                 if element.SSL {
  87                         n = "s"
  88                 }
  89
  90                 // Construct server URL
  91                 url := fmt.Sprintf("http%s://%s:%d", n, element.Hostname, element.Port)
  92
  93                 // Skip duplicates
  94                 if !listed[url] {
  95                         listed[url] = true
  96                         this.Service_roots = append(this.Service_roots, url)
  97                 }
  98         }
  99
 100         // Must be sorted for ShuffledServiceRoots() to produce consistent
 101         // results.
 102         sort.Strings(this.Service_roots)
 103
 104         return nil
 105 }
 106
 107 func (this KeepClient) ShuffledServiceRoots(hash string) (pseq []string) {
 108         // Build an ordering with which to query the Keep servers based on the
 109         // contents of the hash.  "hash" is a hex-encoded number at least 8
 110         // digits (32 bits) long
 111
 112         // seed used to calculate the next keep server from 'pool' to be added
 113         // to 'pseq'
 114         seed := hash
 115
 116         // Keep servers still to be added to the ordering
 117         pool := make([]string, len(this.Service_roots))
 118         copy(pool, this.Service_roots)
 119
 120         // output probe sequence
 121         pseq = make([]string, 0, len(this.Service_roots))
 122
 123         // iterate while there are servers left to be assigned
 124         for len(pool) > 0 {
 125
 126                 if len(seed) < 8 {
 127                         // ran out of digits in the seed
 128                         if len(pseq) < (len(hash) / 4) {
 129                                 // the number of servers added to the probe
 130                                 // sequence is less than the number of 4-digit
 131                                 // slices in 'hash' so refill the seed with the
 132                                 // last 4 digits.
 133                                 seed = hash[len(hash)-4:]
 134                         }
 135                         seed += hash
 136                 }
 137
 138                 // Take the next 8 digits (32 bytes) and interpret as an integer,
 139                 // then modulus with the size of the remaining pool to get the next
 140                 // selected server.
 141                 probe, _ := strconv.ParseUint(seed[0:8], 16, 32)
 142                 probe %= uint64(len(pool))
 143
 144                 // Append the selected server to the probe sequence and remove it
 145                 // from the pool.
 146                 pseq = append(pseq, pool[probe])
 147                 pool = append(pool[:probe], pool[probe+1:]...)
 148
 149                 // Remove the digits just used from the seed
 150                 seed = seed[8:]
 151         }
 152         return pseq
 153 }
 154
 155 type ReaderSlice struct {
 156         slice        []byte
 157         reader_error error
 158 }
 159
 160 // Read repeatedly from the reader into the specified buffer, and report each
 161 // read to channel 'c'.  Completes when Reader 'r' reports on the error channel
 162 // and closes channel 'c'.
 163 func ReadIntoBuffer(buffer []byte, r io.Reader, slices chan<- ReaderSlice) {
 164         defer close(slices)
 165
 166         // Initially use entire buffer as scratch space
 167         ptr := buffer[:]
 168         for {
 169                 var n int
 170                 var err error
 171                 if len(ptr) > 0 {
 172                         // Read into the scratch space
 173                         n, err = r.Read(ptr)
 174                 } else {
 175                         // Ran out of scratch space, try reading one more byte
 176                         var b [1]byte
 177                         n, err = r.Read(b[:])
 178
 179                         if n > 0 {
 180                                 // Reader has more data but we have nowhere to
 181                                 // put it, so we're stuffed
 182                                 slices <- ReaderSlice{nil, io.ErrShortBuffer}
 183                         } else {
 184                                 // Return some other error (hopefully EOF)
 185                                 slices <- ReaderSlice{nil, err}
 186                         }
 187                         return
 188                 }
 189
 190                 // End on error (includes EOF)
 191                 if err != nil {
 192                         slices <- ReaderSlice{nil, err}
 193                         return
 194                 }
 195
 196                 if n > 0 {
 197                         // Make a slice with the contents of the read
 198                         slices <- ReaderSlice{ptr[:n], nil}
 199
 200                         // Adjust the scratch space slice
 201                         ptr = ptr[n:]
 202                 }
 203         }
 204 }
 205
 206 // A read request to the Transfer() function
 207 type ReadRequest struct {
 208         offset  int
 209         maxsize int
 210         result  chan<- ReadResult
 211 }
 212
 213 // A read result from the Transfer() function
 214 type ReadResult struct {
 215         slice []byte
 216         err   error
 217 }
 218
 219 // Reads from the buffer managed by the Transfer()
 220 type BufferReader struct {
 221         offset    *int
 222         requests  chan<- ReadRequest
 223         responses chan ReadResult
 224 }
 225
 226 func MakeBufferReader(requests chan<- ReadRequest) BufferReader {
 227         return BufferReader{new(int), requests, make(chan ReadResult)}
 228 }
 229
 230 // Reads from the buffer managed by the Transfer()
 231 func (this BufferReader) Read(p []byte) (n int, err error) {
 232         this.requests <- ReadRequest{*this.offset, len(p), this.responses}
 233         rr, valid := <-this.responses
 234         if valid {
 235                 *this.offset += len(rr.slice)
 236                 return copy(p, rr.slice), rr.err
 237         } else {
 238                 return 0, io.ErrUnexpectedEOF
 239         }
 240 }
 241
 242 func (this BufferReader) WriteTo(dest io.Writer) (written int64, err error) {
 243         // Record starting offset in order to correctly report the number of bytes sent
 244         starting_offset := *this.offset
 245         for {
 246                 this.requests <- ReadRequest{*this.offset, 32 * 1024, this.responses}
 247                 rr, valid := <-this.responses
 248                 if valid {
 249                         log.Printf("WriteTo slice %v %d %v", *this.offset, len(rr.slice), rr.err)
 250                         *this.offset += len(rr.slice)
 251                         if rr.err != nil {
 252                                 if rr.err == io.EOF {
 253                                         // EOF is not an error.
 254                                         return int64(*this.offset - starting_offset), nil
 255                                 } else {
 256                                         return int64(*this.offset - starting_offset), rr.err
 257                                 }
 258                         } else {
 259                                 dest.Write(rr.slice)
 260                         }
 261                 } else {
 262                         return int64(*this.offset), io.ErrUnexpectedEOF
 263                 }
 264         }
 265 }
 266
 267 // Close the responses channel
 268 func (this BufferReader) Close() error {
 269         close(this.responses)
 270         return nil
 271 }
 272
 273 // Handle a read request.  Returns true if a response was sent, and false if
 274 // the request should be queued.
 275 func HandleReadRequest(req ReadRequest, body []byte, complete bool) bool {
 276         log.Printf("HandleReadRequest %d %d %d", req.offset, req.maxsize, len(body))
 277         if req.offset < len(body) {
 278                 var end int
 279                 if req.offset+req.maxsize < len(body) {
 280                         end = req.offset + req.maxsize
 281                 } else {
 282                         end = len(body)
 283                 }
 284                 req.result <- ReadResult{body[req.offset:end], nil}
 285                 return true
 286         } else if complete && req.offset >= len(body) {
 287                 req.result <- ReadResult{nil, io.EOF}
 288                 return true
 289         } else {
 290                 return false
 291         }
 292 }
 293
 294 // If 'source_reader' is not nil, reads data from 'source_reader' and stores it
 295 // in the provided buffer.  Otherwise, use the contents of 'buffer' as is.
 296 // Accepts read requests on the buffer on the 'requests' channel.  Completes
 297 // when 'requests' channel is closed.
 298 func Transfer(source_buffer []byte, source_reader io.Reader, requests <-chan ReadRequest, reader_error chan error) {
 299         // currently buffered data
 300         var body []byte
 301
 302         // for receiving slices from ReadIntoBuffer
 303         var slices chan ReaderSlice = nil
 304
 305         // indicates whether the buffered data is complete
 306         var complete bool = false
 307
 308         if source_reader != nil {
 309                 // 'body' is the buffer slice representing the body content read so far
 310                 body = source_buffer[:0]
 311
 312                 // used to communicate slices of the buffer as they are
 313                 // ReadIntoBuffer will close 'slices' when it is done with it
 314                 slices = make(chan ReaderSlice)
 315
 316                 // Spin it off
 317                 go ReadIntoBuffer(source_buffer, source_reader, slices)
 318         } else {
 319                 // use the whole buffer
 320                 body = source_buffer[:]
 321
 322                 // buffer is complete
 323                 complete = true
 324         }
 325
 326         pending_requests := make([]ReadRequest, 0)
 327
 328         for {
 329                 select {
 330                 case req, valid := <-requests:
 331                         // Handle a buffer read request
 332                         if valid {
 333                                 if !HandleReadRequest(req, body, complete) {
 334                                         pending_requests = append(pending_requests, req)
 335                                 }
 336                         } else {
 337                                 // closed 'requests' channel indicates we're done
 338                                 return
 339                         }
 340
 341                 case bk, valid := <-slices:
 342                         // Got a new slice from the reader
 343                         if valid {
 344                                 if bk.reader_error != nil {
 345                                         reader_error <- bk.reader_error
 346                                         if bk.reader_error == io.EOF {
 347                                                 // EOF indicates the reader is done
 348                                                 // sending, so our buffer is complete.
 349                                                 complete = true
 350                                         } else {
 351                                                 // some other reader error
 352                                                 return
 353                                         }
 354                                 }
 355
 356                                 if bk.slice != nil {
 357                                         // adjust body bounds now that another slice has been read
 358                                         body = source_buffer[0 : len(body)+len(bk.slice)]
 359                                 }
 360
 361                                 // handle pending reads
 362                                 n := 0
 363                                 for n < len(pending_requests) {
 364                                         if HandleReadRequest(pending_requests[n], body, complete) {
 365
 366                                                 // move the element from the
 367                                                 // back of the slice to
 368                                                 // position 'n', then shorten
 369                                                 // the slice by one element
 370                                                 pending_requests[n] = pending_requests[len(pending_requests)-1]
 371                                                 pending_requests = pending_requests[0 : len(pending_requests)-1]
 372                                         } else {
 373
 374                                                 // Request wasn't handled, so keep it in the request slice
 375                                                 n += 1
 376                                         }
 377                                 }
 378                         } else {
 379                                 if complete {
 380                                         // no more reads
 381                                         slices = nil
 382                                 } else {
 383                                         // reader channel closed without signaling EOF
 384                                         reader_error <- io.ErrUnexpectedEOF
 385                                         return
 386                                 }
 387                         }
 388                 }
 389         }
 390 }
 391
 392 type UploadStatus struct {
 393         Err        error
 394         Url        string
 395         StatusCode int
 396 }
 397
 398 func (this KeepClient) uploadToKeepServer(host string, hash string, body io.ReadCloser,
 399         upload_status chan<- UploadStatus, expectedLength int64) {
 400
 401         log.Printf("Uploading to %s", host)
 402
 403         var req *http.Request
 404         var err error
 405         var url = fmt.Sprintf("%s/%s", host, hash)
 406         if req, err = http.NewRequest("PUT", url, nil); err != nil {
 407                 upload_status <- UploadStatus{err, url, 0}
 408                 return
 409         }
 410
 411         if expectedLength > 0 {
 412                 req.ContentLength = expectedLength
 413         }
 414
 415         req.Header.Add("Authorization", fmt.Sprintf("OAuth2 %s", this.ApiToken))
 416         req.Header.Add("Content-Type", "application/octet-stream")
 417         req.Body = body
 418
 419         var resp *http.Response
 420         if resp, err = this.Client.Do(req); err != nil {
 421                 upload_status <- UploadStatus{err, url, 0}
 422                 return
 423         }
 424
 425         if resp.StatusCode == http.StatusOK {
 426                 upload_status <- UploadStatus{nil, url, resp.StatusCode}
 427         } else {
 428                 upload_status <- UploadStatus{errors.New(resp.Status), url, resp.StatusCode}
 429         }
 430 }
 431
 432 var InsufficientReplicasError = errors.New("Could not write sufficient replicas")
 433
 434 func (this KeepClient) putReplicas(
 435         hash string,
 436         requests chan ReadRequest,
 437         reader_status chan error,
 438         expectedLength int64) (replicas int, err error) {
 439
 440         // Calculate the ordering for uploading to servers
 441         sv := this.ShuffledServiceRoots(hash)
 442
 443         // The next server to try contacting
 444         next_server := 0
 445
 446         // The number of active writers
 447         active := 0
 448
 449         // Used to communicate status from the upload goroutines
 450         upload_status := make(chan UploadStatus)
 451         defer close(upload_status)
 452
 453         // Desired number of replicas
 454         remaining_replicas := this.Want_replicas
 455
 456         for remaining_replicas > 0 {
 457                 for active < remaining_replicas {
 458                         // Start some upload requests
 459                         if next_server < len(sv) {
 460                                 go this.uploadToKeepServer(sv[next_server], hash, MakeBufferReader(requests), upload_status, expectedLength)
 461                                 next_server += 1
 462                                 active += 1
 463                         } else {
 464                                 return (this.Want_replicas - remaining_replicas), InsufficientReplicasError
 465                         }
 466                 }
 467
 468                 // Now wait for something to happen.
 469                 select {
 470                 case status := <-reader_status:
 471                         if status == io.EOF {
 472                                 // good news!
 473                         } else {
 474                                 // bad news
 475                                 return (this.Want_replicas - remaining_replicas), status
 476                         }
 477                 case status := <-upload_status:
 478                         if status.StatusCode == 200 {
 479                                 // good news!
 480                                 remaining_replicas -= 1
 481                         } else {
 482                                 // writing to keep server failed for some reason
 483                                 log.Printf("Keep server put to %v failed with '%v'",
 484                                         status.Url, status.Err)
 485                         }
 486                         active -= 1
 487                         log.Printf("Upload status %v %v %v", status.StatusCode, remaining_replicas, active)
 488                 }
 489         }
 490
 491         return (this.Want_replicas - remaining_replicas), nil
 492 }
 493
 494 var OversizeBlockError = errors.New("Block too big")
 495
 496 func (this KeepClient) PutHR(hash string, r io.Reader, expectedLength int64) (replicas int, err error) {
 497
 498         // Buffer for reads from 'r'
 499         var buffer []byte
 500         if expectedLength > 0 {
 501                 if expectedLength > BLOCKSIZE {
 502                         return 0, OversizeBlockError
 503                 }
 504                 buffer = make([]byte, expectedLength)
 505         } else {
 506                 buffer = make([]byte, BLOCKSIZE)
 507         }
 508
 509         // Read requests on Transfer() buffer
 510         requests := make(chan ReadRequest)
 511         defer close(requests)
 512
 513         // Reporting reader error states
 514         reader_status := make(chan error)
 515         defer close(reader_status)
 516
 517         // Start the transfer goroutine
 518         go Transfer(buffer, r, requests, reader_status)
 519
 520         return this.putReplicas(hash, requests, reader_status, expectedLength)
 521 }
 522
 523 func (this KeepClient) PutHB(hash string, buffer []byte) (replicas int, err error) {
 524         // Read requests on Transfer() buffer
 525         requests := make(chan ReadRequest)
 526         defer close(requests)
 527
 528         // Start the transfer goroutine
 529         go Transfer(buffer, nil, requests, nil)
 530
 531         return this.putReplicas(hash, requests, nil, int64(len(buffer)))
 532 }
 533
 534 func (this KeepClient) PutB(buffer []byte) (hash string, replicas int, err error) {
 535         hash = fmt.Sprintf("%x", md5.Sum(buffer))
 536         replicas, err = this.PutHB(hash, buffer)
 537         return hash, replicas, err
 538 }
 539
 540 func (this KeepClient) PutR(r io.Reader) (hash string, replicas int, err error) {
 541         if buffer, err := ioutil.ReadAll(r); err != nil {
 542                 return "", 0, err
 543         } else {
 544                 return this.PutB(buffer)
 545         }
 546 }
 547
 548 var BlockNotFound = errors.New("Block not found")
 549
 550 func (this KeepClient) Get(hash string) (reader io.ReadCloser,
 551         contentLength int64, url string, err error) {
 552         return this.AuthorizedGet(hash, "", "")
 553 }
 554
 555 func (this KeepClient) AuthorizedGet(hash string,
 556         signature string,
 557         timestamp string) (reader io.ReadCloser,
 558         contentLength int64, url string, err error) {
 559
 560         // Calculate the ordering for asking servers
 561         sv := this.ShuffledServiceRoots(hash)
 562
 563         for _, host := range sv {
 564                 var req *http.Request
 565                 var err error
 566                 var url string
 567                 if signature != "" {
 568                         url = fmt.Sprintf("%s/%s+A%s@%s", host, hash,
 569                                 signature, timestamp)
 570                 } else {
 571                         url = fmt.Sprintf("%s/%s", host, hash)
 572                 }
 573                 if req, err = http.NewRequest("GET", url, nil); err != nil {
 574                         continue
 575                 }
 576
 577                 req.Header.Add("Authorization", fmt.Sprintf("OAuth2 %s", this.ApiToken))
 578
 579                 var resp *http.Response
 580                 if resp, err = this.Client.Do(req); err != nil {
 581                         continue
 582                 }
 583
 584                 if resp.StatusCode == http.StatusOK {
 585                         return resp.Body, resp.ContentLength, url, nil
 586                 }
 587         }
 588
 589         return nil, 0, "", BlockNotFound
 590 }
 591
 592 func (this KeepClient) Ask(hash string) (contentLength int64, url string, err error) {
 593         return this.AuthorizedAsk(hash, "", "")
 594 }
 595
 596 func (this KeepClient) AuthorizedAsk(hash string, signature string,
 597         timestamp string) (contentLength int64, url string, err error) {
 598         // Calculate the ordering for asking servers
 599         sv := this.ShuffledServiceRoots(hash)
 600
 601         for _, host := range sv {
 602                 var req *http.Request
 603                 var err error
 604                 if signature != "" {
 605                         url = fmt.Sprintf("%s/%s+A%s@%s", host, hash,
 606                                 signature, timestamp)
 607                 } else {
 608                         url = fmt.Sprintf("%s/%s", host, hash)
 609                 }
 610
 611                 if req, err = http.NewRequest("HEAD", url, nil); err != nil {
 612                         continue
 613                 }
 614
 615                 req.Header.Add("Authorization", fmt.Sprintf("OAuth2 %s", this.ApiToken))
 616
 617                 var resp *http.Response
 618                 if resp, err = this.Client.Do(req); err != nil {
 619                         continue
 620                 }
 621
 622                 if resp.StatusCode == http.StatusOK {
 623                         return resp.ContentLength, url, nil
 624                 }
 625         }
 626
 627         return 0, "", BlockNotFound
 628
 629 }