services/keepstore/pull_worker.go

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