services/keepstore/pull_worker.go

   1 package main
   2
   3 import (
   4         "context"
   5         "crypto/rand"
   6         "fmt"
   7         "io"
   8         "io/ioutil"
   9         "time"
  10
  11         "git.curoverse.com/arvados.git/sdk/go/keepclient"
  12
  13         log "github.com/Sirupsen/logrus"
  14 )
  15
  16 // RunPullWorker receives PullRequests from pullq, invokes
  17 // PullItemAndProcess on each one. After each PR, it logs a message
  18 // indicating whether the pull was successful.
  19 func RunPullWorker(pullq *WorkQueue, keepClient *keepclient.KeepClient) {
  20         for item := range pullq.NextItem {
  21                 pr := item.(PullRequest)
  22                 err := PullItemAndProcess(pr, keepClient)
  23                 pullq.DoneItem <- struct{}{}
  24                 if err == nil {
  25                         log.Printf("Pull %s success", pr)
  26                 } else {
  27                         log.Printf("Pull %s error: %s", pr, err)
  28                 }
  29         }
  30 }
  31
  32 // PullItemAndProcess executes a pull request by retrieving the
  33 // specified block from one of the specified servers, and storing it
  34 // on a local volume.
  35 //
  36 // If the PR specifies a non-blank mount UUID, PullItemAndProcess will
  37 // only attempt to write the data to the corresponding
  38 // volume. Otherwise it writes to any local volume, as a PUT request
  39 // would.
  40 func PullItemAndProcess(pullRequest PullRequest, keepClient *keepclient.KeepClient) error {
  41         var vol Volume
  42         if uuid := pullRequest.MountUUID; uuid != "" {
  43                 vol = KeepVM.Lookup(pullRequest.MountUUID, true)
  44                 if vol == nil {
  45                         return fmt.Errorf("pull req has nonexistent mount: %v", pullRequest)
  46                 }
  47         }
  48
  49         keepClient.Arvados.ApiToken = randomToken
  50
  51         serviceRoots := make(map[string]string)
  52         for _, addr := range pullRequest.Servers {
  53                 serviceRoots[addr] = addr
  54         }
  55         keepClient.SetServiceRoots(serviceRoots, nil, nil)
  56
  57         // Generate signature with a random token
  58         expiresAt := time.Now().Add(60 * time.Second)
  59         signedLocator := SignLocator(pullRequest.Locator, randomToken, expiresAt)
  60
  61         reader, contentLen, _, err := GetContent(signedLocator, keepClient)
  62         if err != nil {
  63                 return err
  64         }
  65         if reader == nil {
  66                 return fmt.Errorf("No reader found for : %s", signedLocator)
  67         }
  68         defer reader.Close()
  69
  70         readContent, err := ioutil.ReadAll(reader)
  71         if err != nil {
  72                 return err
  73         }
  74
  75         if (readContent == nil) || (int64(len(readContent)) != contentLen) {
  76                 return fmt.Errorf("Content not found for: %s", signedLocator)
  77         }
  78
  79         writePulledBlock(vol, readContent, pullRequest.Locator)
  80         return nil
  81 }
  82
  83 // Fetch the content for the given locator using keepclient.
  84 var GetContent = func(signedLocator string, keepClient *keepclient.KeepClient) (io.ReadCloser, int64, string, error) {
  85         return keepClient.Get(signedLocator)
  86 }
  87
  88 var writePulledBlock = func(volume Volume, data []byte, locator string) {
  89         var err error
  90         if volume != nil {
  91                 err = volume.Put(context.Background(), locator, data)
  92         } else {
  93                 _, err = PutBlock(context.Background(), data, locator)
  94         }
  95         if err != nil {
  96                 log.Printf("error writing pulled block %q: %s", locator, err)
  97         }
  98 }
  99
 100 var randomToken = func() string {
 101         const alphaNumeric = "0123456789abcdefghijklmnopqrstuvwxyz"
 102         var bytes = make([]byte, 36)
 103         rand.Read(bytes)
 104         for i, b := range bytes {
 105                 bytes[i] = alphaNumeric[b%byte(len(alphaNumeric))]
 106         }
 107         return (string(bytes))
 108 }()