services/keepstore/pull_worker.go

   1 package main
   2
   3 import (
   4         "context"
   5         "crypto/rand"
   6         "fmt"
   7         "git.curoverse.com/arvados.git/sdk/go/keepclient"
   8         "io"
   9         "io/ioutil"
  10         "time"
  11
  12         log "github.com/Sirupsen/logrus"
  13 )
  14
  15 // RunPullWorker is used by Keepstore to initiate pull worker channel goroutine.
  16 //      The channel will process pull list.
  17 //              For each (next) pull request:
  18 //                      For each locator listed, execute Pull on the server(s) listed
  19 //                      Skip the rest of the servers if no errors
  20 //              Repeat
  21 //
  22 func RunPullWorker(pullq *WorkQueue, keepClient *keepclient.KeepClient) {
  23         nextItem := pullq.NextItem
  24         for item := range nextItem {
  25                 pullRequest := item.(PullRequest)
  26                 err := PullItemAndProcess(item.(PullRequest), GenerateRandomAPIToken(), keepClient)
  27                 pullq.DoneItem <- struct{}{}
  28                 if err == nil {
  29                         log.Printf("Pull %s success", pullRequest)
  30                 } else {
  31                         log.Printf("Pull %s error: %s", pullRequest, err)
  32                 }
  33         }
  34 }
  35
  36 // PullItemAndProcess pulls items from PullQueue and processes them.
  37 //      For each Pull request:
  38 //              Generate a random API token.
  39 //              Generate a permission signature using this token, timestamp ~60 seconds in the future, and desired block hash.
  40 //              Using this token & signature, retrieve the given block.
  41 //              Write to storage
  42 //
  43 func PullItemAndProcess(pullRequest PullRequest, token string, keepClient *keepclient.KeepClient) (err error) {
  44         keepClient.Arvados.ApiToken = token
  45
  46         serviceRoots := make(map[string]string)
  47         for _, addr := range pullRequest.Servers {
  48                 serviceRoots[addr] = addr
  49         }
  50         keepClient.SetServiceRoots(serviceRoots, nil, nil)
  51
  52         // Generate signature with a random token
  53         expiresAt := time.Now().Add(60 * time.Second)
  54         signedLocator := SignLocator(pullRequest.Locator, token, expiresAt)
  55
  56         reader, contentLen, _, err := GetContent(signedLocator, keepClient)
  57         if err != nil {
  58                 return
  59         }
  60         if reader == nil {
  61                 return fmt.Errorf("No reader found for : %s", signedLocator)
  62         }
  63         defer reader.Close()
  64
  65         readContent, err := ioutil.ReadAll(reader)
  66         if err != nil {
  67                 return err
  68         }
  69
  70         if (readContent == nil) || (int64(len(readContent)) != contentLen) {
  71                 return fmt.Errorf("Content not found for: %s", signedLocator)
  72         }
  73
  74         err = PutContent(readContent, pullRequest.Locator)
  75         return
  76 }
  77
  78 // Fetch the content for the given locator using keepclient.
  79 var GetContent = func(signedLocator string, keepClient *keepclient.KeepClient) (
  80         reader io.ReadCloser, contentLength int64, url string, err error) {
  81         reader, blocklen, url, err := keepClient.Get(signedLocator)
  82         return reader, blocklen, url, err
  83 }
  84
  85 const alphaNumeric = "0123456789abcdefghijklmnopqrstuvwxyz"
  86
  87 // GenerateRandomAPIToken generates a random api token
  88 func GenerateRandomAPIToken() string {
  89         var bytes = make([]byte, 36)
  90         rand.Read(bytes)
  91         for i, b := range bytes {
  92                 bytes[i] = alphaNumeric[b%byte(len(alphaNumeric))]
  93         }
  94         return (string(bytes))
  95 }
  96
  97 // Put block
  98 var PutContent = func(content []byte, locator string) (err error) {
  99         _, err = PutBlock(context.Background(), content, locator)
 100         return
 101 }