X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/afd38976a33c46b618ff9317192388b8f3f1a195..80f547e339c452e5b03be5beee00e845d56d8e18:/services/keepstore/pull_worker.go

diff --git a/services/keepstore/pull_worker.go b/services/keepstore/pull_worker.go
index b68e933130..2626d4bf68 100644
--- a/services/keepstore/pull_worker.go
+++ b/services/keepstore/pull_worker.go
@@ -4,121 +4,97 @@ 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.
-		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
-*/
-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())
-	}
-
+// 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
+//
+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
-			}
-		}
-	}
-}
-
-/*
-	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
-*/
-func Pull(addr string, locator string) (err error) {
-	log.Printf("Pull %s/%s starting", addr, locator)
-
-	defer func() {
+		pullRequest := item.(PullRequest)
+		err := PullItemAndProcess(item.(PullRequest), GenerateRandomAPIToken(), keepClient)
+		pullq.DoneItem <- struct{}{}
 		if err == nil {
-			log.Printf("Pull %s/%s success", addr, locator)
+			log.Printf("Pull %s success", pullRequest)
 		} else {
-			log.Printf("Pull %s/%s error: %s", addr, locator, err)
+			log.Printf("Pull %s error: %s", pullRequest, err)
 		}
-	}()
-
-	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
 	}
-
-	err = PutContent(read_content, locator)
-	return
 }
 
-// Fetch the content for the given locator using keepclient.
-var GetContent = func(addr string, locator string) ([]byte, error) {
+// 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
+//
+func PullItemAndProcess(pullRequest PullRequest, token string, keepClient *keepclient.KeepClient) (err error) {
+	keepClient.Arvados.ApiToken = token
+
+	serviceRoots := make(map[string]string)
+	for _, addr := range pullRequest.Servers {
+		serviceRoots[addr] = addr
+	}
+	keepClient.SetServiceRoots(serviceRoots, nil, nil)
+
 	// 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)
+	expiresAt := time.Now().Add(60 * time.Second)
+	signedLocator := SignLocator(pullRequest.Locator, token, expiresAt)
 
-	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)
+	readContent, err := ioutil.ReadAll(reader)
 	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 (readContent == nil) || (int64(len(readContent)) != contentLen) {
+		return errors.New(fmt.Sprintf("Content not found for: %s", signedLocator))
 	}
 
-	return read_content, nil
+	err = PutContent(readContent, 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"
+const alphaNumeric = "0123456789abcdefghijklmnopqrstuvwxyz"
 
-func GenerateRandomApiToken() string {
+// GenerateRandomAPIToken generates a random api token
+func GenerateRandomAPIToken() string {
 	var bytes = make([]byte, 36)
 	rand.Read(bytes)
 	for i, b := range bytes {
-		bytes[i] = ALPHA_NUMERIC[b%byte(len(ALPHA_NUMERIC))]
+		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)
+	_, err = PutBlock(content, locator)
 	return
 }