7179: Tighten Put requirements when overwriting existing data.

[arvados.git] / services / keepstore / pull_worker.go

diff --git a/services/keepstore/pull_worker.go b/services/keepstore/pull_worker.go
index b68e93313027c0480e8f17d6ca78b42ad3f8065d..acf861119f47fd1b765bcad461d826c369151968 100644 (file)
--- a/services/keepstore/pull_worker.go
+++ b/services/keepstore/pull_worker.go
@@ -4,18 +4,13 @@ import (
        "crypto/rand"
        "errors"
        "fmt"
-       "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
        "git.curoverse.com/arvados.git/sdk/go/keepclient"
+       "io"
        "io/ioutil"
        "log"
-       "os"
-       "strconv"
        "time"
 )
 
-var arv arvadosclient.ArvadosClient
-var keepClient keepclient.KeepClient
-
 /*
        Keepstore initiates pull worker channel goroutine.
        The channel will process pull list.
@@ -24,26 +19,16 @@ var keepClient keepclient.KeepClient
                        Skip the rest of the servers if no errors
                Repeat
 */
-func RunPullWorker(nextItem <-chan interface{}) {
-       var err error
-       arv, err = arvadosclient.MakeArvadosClient()
-       if err != nil {
-               log.Fatalf("Error setting up arvados client %s", err.Error())
-       }
-       arv.ApiToken = os.Getenv("ARVADOS_API_TOKEN")
-
-       keepClient, err = keepclient.MakeKeepClient(&arv)
-       if err != nil {
-               log.Fatalf("Error setting up keep client %s", err.Error())
-       }
-
+func RunPullWorker(pullq *WorkQueue, keepClient *keepclient.KeepClient) {
+       nextItem := pullq.NextItem
        for item := range nextItem {
-               pullReq := item.(PullRequest)
-               for _, addr := range pullReq.Servers {
-                       err := Pull(addr, pullReq.Locator)
-                       if err == nil {
-                               break
-                       }
+               pullRequest := item.(PullRequest)
+               err := PullItemAndProcess(item.(PullRequest), GenerateRandomApiToken(), keepClient)
+               pullq.DoneItem <- struct{}{}
+               if err == nil {
+                       log.Printf("Pull %s success", pullRequest)
+               } else {
+                       log.Printf("Pull %s error: %s", pullRequest, err)
                }
        }
 }
@@ -55,55 +40,46 @@ func RunPullWorker(nextItem <-chan interface{}) {
                Using this token & signature, retrieve the given block.
                Write to storage
 */
-func Pull(addr string, locator string) (err error) {
-       log.Printf("Pull %s/%s starting", addr, locator)
-
-       defer func() {
-               if err == nil {
-                       log.Printf("Pull %s/%s success", addr, locator)
-               } else {
-                       log.Printf("Pull %s/%s error: %s", addr, locator, err)
-               }
-       }()
+func PullItemAndProcess(pullRequest PullRequest, token string, keepClient *keepclient.KeepClient) (err error) {
+       keepClient.Arvados.ApiToken = token
 
        service_roots := make(map[string]string)
-       service_roots[locator] = addr
-       keepClient.SetServiceRoots(service_roots)
-
-       read_content, err := GetContent(addr, locator)
-       log.Print(read_content, err)
-       if err != nil {
-               return
+       for _, addr := range pullRequest.Servers {
+               service_roots[addr] = addr
        }
+       keepClient.SetServiceRoots(service_roots, nil, nil)
 
-       err = PutContent(read_content, locator)
-       return
-}
-
-// Fetch the content for the given locator using keepclient.
-var GetContent = func(addr string, locator string) ([]byte, error) {
        // Generate signature with a random token
-       expires_at := time.Now().Unix() + 60 // now + 1 min in seconds
-       hints := "+A" + GenerateRandomApiToken() + "@" + strconv.FormatInt(expires_at, 16)
-       signature := keepclient.MakeLocator2(locator, hints)
+       expires_at := time.Now().Add(60 * time.Second)
+       signedLocator := SignLocator(pullRequest.Locator, token, expires_at)
 
-       reader, blocklen, _, err := keepClient.AuthorizedGet(locator, signature.Signature, signature.Timestamp)
-       defer reader.Close()
+       reader, contentLen, _, err := GetContent(signedLocator, keepClient)
        if err != nil {
-               return nil, err
+               return
+       }
+       if reader == nil {
+               return errors.New(fmt.Sprintf("No reader found for : %s", signedLocator))
        }
+       defer reader.Close()
 
        read_content, err := ioutil.ReadAll(reader)
-       log.Print(read_content, err)
        if err != nil {
-               return nil, err
+               return err
        }
 
-       if (read_content == nil) || (int64(len(read_content)) != blocklen) {
-               return nil, errors.New(fmt.Sprintf("Content not found for: %s/%s", addr, locator))
+       if (read_content == nil) || (int64(len(read_content)) != contentLen) {
+               return errors.New(fmt.Sprintf("Content not found for: %s", signedLocator))
        }
 
-       return read_content, nil
+       err = PutContent(read_content, pullRequest.Locator)
+       return
+}
+
+// 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
 }
 
 const ALPHA_NUMERIC = "0123456789abcdefghijklmnopqrstuvwxyz"