services/keepstore/pull_worker.go

   1 package main
   2
   3 import (
   4         "crypto/rand"
   5         "errors"
   6         "fmt"
   7         "git.curoverse.com/arvados.git/sdk/go/keepclient"
   8         "io"
   9         "io/ioutil"
  10         "log"
  11         "time"
  12 )
  13
  14 /*
  15         Keepstore initiates 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 /*
  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         service_roots := make(map[string]string)
  47         for _, addr := range pullRequest.Servers {
  48                 service_roots[addr] = addr
  49         }
  50         keepClient.SetServiceRoots(service_roots, nil, nil)
  51
  52         // Generate signature with a random token
  53         expires_at := time.Now().Add(60 * time.Second)
  54         signedLocator := SignLocator(pullRequest.Locator, token, expires_at)
  55
  56         reader, contentLen, _, err := GetContent(signedLocator, keepClient)
  57         if err != nil {
  58                 return
  59         }
  60         if reader == nil {
  61                 return errors.New(fmt.Sprintf("No reader found for : %s", signedLocator))
  62         }
  63         defer reader.Close()
  64
  65         read_content, err := ioutil.ReadAll(reader)
  66         if err != nil {
  67                 return err
  68         }
  69
  70         if (read_content == nil) || (int64(len(read_content)) != contentLen) {
  71                 return errors.New(fmt.Sprintf("Content not found for: %s", signedLocator))
  72         }
  73
  74         err = PutContent(read_content, 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 ALPHA_NUMERIC = "0123456789abcdefghijklmnopqrstuvwxyz"
  86
  87 func GenerateRandomApiToken() string {
  88         var bytes = make([]byte, 36)
  89         rand.Read(bytes)
  90         for i, b := range bytes {
  91                 bytes[i] = ALPHA_NUMERIC[b%byte(len(ALPHA_NUMERIC))]
  92         }
  93         return (string(bytes))
  94 }
  95
  96 // Put block
  97 var PutContent = func(content []byte, locator string) (err error) {
  98         err = PutBlock(content, locator)
  99         return
 100 }