services/keepstore/pull_worker.go

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