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 is used by Keepstore to initiate pull worker channel goroutine.
  17 //      The channel will process pull list.
  18 //              For each (next) pull request:
  19 //                      For each locator listed, execute Pull on the server(s) listed
  20 //                      Skip the rest of the servers if no errors
  21 //              Repeat
  22 //
  23 func RunPullWorker(pullq *WorkQueue, keepClient *keepclient.KeepClient) {
  24         nextItem := pullq.NextItem
  25         for item := range nextItem {
  26                 pullRequest := item.(PullRequest)
  27                 err := PullItemAndProcess(item.(PullRequest), keepClient)
  28                 pullq.DoneItem <- struct{}{}
  29                 if err == nil {
  30                         log.Printf("Pull %s success", pullRequest)
  31                 } else {
  32                         log.Printf("Pull %s error: %s", pullRequest, err)
  33                 }
  34         }
  35 }
  36
  37 // PullItemAndProcess pulls items from PullQueue and processes them.
  38 //      For each Pull request:
  39 //              Generate a random API token.
  40 //              Generate a permission signature using this token, timestamp ~60 seconds in the future, and desired block hash.
  41 //              Using this token & signature, retrieve the given block.
  42 //              Write to storage
  43 //
  44 func PullItemAndProcess(pullRequest PullRequest, keepClient *keepclient.KeepClient) error {
  45         var vol Volume
  46         if uuid := pullRequest.MountUUID; uuid != "" {
  47                 vol = KeepVM.Lookup(pullRequest.MountUUID, true)
  48                 if vol == nil {
  49                         return fmt.Errorf("pull req has nonexistent mount: %v", pullRequest)
  50                 }
  51         }
  52
  53         keepClient.Arvados.ApiToken = randomToken
  54
  55         serviceRoots := make(map[string]string)
  56         for _, addr := range pullRequest.Servers {
  57                 serviceRoots[addr] = addr
  58         }
  59         keepClient.SetServiceRoots(serviceRoots, nil, nil)
  60
  61         // Generate signature with a random token
  62         expiresAt := time.Now().Add(60 * time.Second)
  63         signedLocator := SignLocator(pullRequest.Locator, randomToken, expiresAt)
  64
  65         reader, contentLen, _, err := GetContent(signedLocator, keepClient)
  66         if err != nil {
  67                 return err
  68         }
  69         if reader == nil {
  70                 return fmt.Errorf("No reader found for : %s", signedLocator)
  71         }
  72         defer reader.Close()
  73
  74         readContent, err := ioutil.ReadAll(reader)
  75         if err != nil {
  76                 return err
  77         }
  78
  79         if (readContent == nil) || (int64(len(readContent)) != contentLen) {
  80                 return fmt.Errorf("Content not found for: %s", signedLocator)
  81         }
  82
  83         writePulledBlock(vol, readContent, pullRequest.Locator)
  84         return nil
  85 }
  86
  87 // Fetch the content for the given locator using keepclient.
  88 var GetContent = func(signedLocator string, keepClient *keepclient.KeepClient) (io.ReadCloser, int64, string, error) {
  89         return keepClient.Get(signedLocator)
  90 }
  91
  92 var writePulledBlock = func(volume Volume, data []byte, locator string) {
  93         var err error
  94         if volume != nil {
  95                 err = volume.Put(context.Background(), locator, data)
  96         } else {
  97                 _, err = PutBlock(context.Background(), data, locator)
  98         }
  99         if err != nil {
 100                 log.Printf("error writing pulled block %q: %s", locator, err)
 101         }
 102 }
 103
 104 var randomToken = func() string {
 105         const alphaNumeric = "0123456789abcdefghijklmnopqrstuvwxyz"
 106         var bytes = make([]byte, 36)
 107         rand.Read(bytes)
 108         for i, b := range bytes {
 109                 bytes[i] = alphaNumeric[b%byte(len(alphaNumeric))]
 110         }
 111         return (string(bytes))
 112 }()