Arvados-DCO-1.1-Signed-off-by: Radhika Chippada <radhika@curoverse.com>
[arvados.git] / services / keepstore / pull_worker.go
index 9f0b96fa35b1c32af4b06be92ea02632547ef249..f821fb548350587528c29bc9fab53528b424eaab 100644 (file)
@@ -1,46 +1,56 @@
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
 package main
 
 import (
+       "context"
        "crypto/rand"
-       "errors"
        "fmt"
-       "git.curoverse.com/arvados.git/sdk/go/keepclient"
        "io"
        "io/ioutil"
-       "log"
        "time"
+
+       "git.curoverse.com/arvados.git/sdk/go/keepclient"
+
+       log "github.com/Sirupsen/logrus"
 )
 
-// RunPullWorker is used by Keepstore to initiate pull worker channel goroutine.
-//     The channel will process pull list.
-//             For each (next) pull request:
-//                     For each locator listed, execute Pull on the server(s) listed
-//                     Skip the rest of the servers if no errors
-//             Repeat
-//
+// RunPullWorker receives PullRequests from pullq, invokes
+// PullItemAndProcess on each one. After each PR, it logs a message
+// indicating whether the pull was successful.
 func RunPullWorker(pullq *WorkQueue, keepClient *keepclient.KeepClient) {
-       nextItem := pullq.NextItem
-       for item := range nextItem {
-               pullRequest := item.(PullRequest)
-               err := PullItemAndProcess(item.(PullRequest), GenerateRandomAPIToken(), keepClient)
+       for item := range pullq.NextItem {
+               pr := item.(PullRequest)
+               err := PullItemAndProcess(pr, keepClient)
                pullq.DoneItem <- struct{}{}
                if err == nil {
-                       log.Printf("Pull %s success", pullRequest)
+                       log.Printf("Pull %s success", pr)
                } else {
-                       log.Printf("Pull %s error: %s", pullRequest, err)
+                       log.Printf("Pull %s error: %s", pr, err)
                }
        }
 }
 
-// PullItemAndProcess pulls items from PullQueue and processes them.
-//     For each Pull request:
-//             Generate a random API token.
-//             Generate a permission signature using this token, timestamp ~60 seconds in the future, and desired block hash.
-//             Using this token & signature, retrieve the given block.
-//             Write to storage
+// PullItemAndProcess executes a pull request by retrieving the
+// specified block from one of the specified servers, and storing it
+// on a local volume.
 //
-func PullItemAndProcess(pullRequest PullRequest, token string, keepClient *keepclient.KeepClient) (err error) {
-       keepClient.Arvados.ApiToken = token
+// If the PR specifies a non-blank mount UUID, PullItemAndProcess will
+// only attempt to write the data to the corresponding
+// volume. Otherwise it writes to any local volume, as a PUT request
+// would.
+func PullItemAndProcess(pullRequest PullRequest, keepClient *keepclient.KeepClient) error {
+       var vol Volume
+       if uuid := pullRequest.MountUUID; uuid != "" {
+               vol = KeepVM.Lookup(pullRequest.MountUUID, true)
+               if vol == nil {
+                       return fmt.Errorf("pull req has nonexistent mount: %v", pullRequest)
+               }
+       }
+
+       keepClient.Arvados.ApiToken = randomToken
 
        serviceRoots := make(map[string]string)
        for _, addr := range pullRequest.Servers {
@@ -50,14 +60,14 @@ func PullItemAndProcess(pullRequest PullRequest, token string, keepClient *keepc
 
        // Generate signature with a random token
        expiresAt := time.Now().Add(60 * time.Second)
-       signedLocator := SignLocator(pullRequest.Locator, token, expiresAt)
+       signedLocator := SignLocator(pullRequest.Locator, randomToken, expiresAt)
 
        reader, contentLen, _, err := GetContent(signedLocator, keepClient)
        if err != nil {
-               return
+               return err
        }
        if reader == nil {
-               return errors.New(fmt.Sprintf("No reader found for : %s", signedLocator))
+               return fmt.Errorf("No reader found for : %s", signedLocator)
        }
        defer reader.Close()
 
@@ -67,34 +77,36 @@ func PullItemAndProcess(pullRequest PullRequest, token string, keepClient *keepc
        }
 
        if (readContent == nil) || (int64(len(readContent)) != contentLen) {
-               return errors.New(fmt.Sprintf("Content not found for: %s", signedLocator))
+               return fmt.Errorf("Content not found for: %s", signedLocator)
        }
 
-       err = PutContent(readContent, pullRequest.Locator)
-       return
+       writePulledBlock(vol, readContent, pullRequest.Locator)
+       return nil
 }
 
 // Fetch the content for the given locator using keepclient.
-var GetContent = func(signedLocator string, keepClient *keepclient.KeepClient) (
-       reader io.ReadCloser, contentLength int64, url string, err error) {
-       reader, blocklen, url, err := keepClient.Get(signedLocator)
-       return reader, blocklen, url, err
+var GetContent = func(signedLocator string, keepClient *keepclient.KeepClient) (io.ReadCloser, int64, string, error) {
+       return keepClient.Get(signedLocator)
 }
 
-const alphaNumeric = "0123456789abcdefghijklmnopqrstuvwxyz"
+var writePulledBlock = func(volume Volume, data []byte, locator string) {
+       var err error
+       if volume != nil {
+               err = volume.Put(context.Background(), locator, data)
+       } else {
+               _, err = PutBlock(context.Background(), data, locator)
+       }
+       if err != nil {
+               log.Printf("error writing pulled block %q: %s", locator, err)
+       }
+}
 
-// GenerateRandomAPIToken generates a random api token
-func GenerateRandomAPIToken() string {
+var randomToken = func() string {
+       const alphaNumeric = "0123456789abcdefghijklmnopqrstuvwxyz"
        var bytes = make([]byte, 36)
        rand.Read(bytes)
        for i, b := range bytes {
                bytes[i] = alphaNumeric[b%byte(len(alphaNumeric))]
        }
        return (string(bytes))
-}
-
-// Put block
-var PutContent = func(content []byte, locator string) (err error) {
-       err = PutBlock(content, locator)
-       return
-}
+}()