Merge branch '7167-keep-rsync-test-setup' into 7167-keep-rsync
[arvados.git] / tools / keep-rsync / keep-rsync.go
index 3761cc7e9c9d154b4a5d061c5e98c85b2b0260e1..f33c8167579c39a6191460cd580f1f3f14148b5e 100644 (file)
@@ -1,18 +1,21 @@
 package main
 
 import (
+       "bytes"
        "flag"
        "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
        "git.curoverse.com/arvados.git/sdk/go/keepclient"
        "io/ioutil"
        "log"
+       "regexp"
        "strings"
 )
 
 // keep-rsync arguments
 var (
-       srcConfig           map[string]string
-       dstConfig           map[string]string
+       srcConfig           arvadosclient.APIConfig
+       dstConfig           arvadosclient.APIConfig
+       blobSigningKey      string
        srcKeepServicesJSON string
        dstKeepServicesJSON string
        replications        int
@@ -29,7 +32,7 @@ func main() {
                "",
                "Source configuration filename with full path that contains "+
                        "an ARVADOS_API_TOKEN which is a valid datamanager token recognized by the source keep servers, "+
-                       "ARVADOS_API_HOST, ARVADOS_API_HOST_INSECURE, and ARVADOS_BLOB_SIGNING_KEY.")
+                       "ARVADOS_API_HOST, ARVADOS_API_HOST_INSECURE, ARVADOS_EXTERNAL_CLIENT and ARVADOS_BLOB_SIGNING_KEY.")
 
        flag.StringVar(
                &dstConfigFile,
@@ -37,7 +40,7 @@ func main() {
                "",
                "Destination configuration filename with full path that contains "+
                        "an ARVADOS_API_TOKEN which is a valid datamanager token recognized by the destination keep servers, "+
-                       "ARVADOS_API_HOST, ARVADOS_API_HOST_INSECURE, and ARVADOS_BLOB_SIGNING_KEY.")
+                       "ARVADOS_API_HOST, ARVADOS_API_HOST_INSECURE, ARVADOS_EXTERNAL_CLIENT and ARVADOS_BLOB_SIGNING_KEY.")
 
        flag.StringVar(
                &srcKeepServicesJSON,
@@ -69,6 +72,7 @@ func main() {
 
        var err error
 
+       // Load config
        if srcConfigFile == "" {
                log.Fatal("-src-config-file must be specified.")
        }
@@ -85,27 +89,46 @@ func main() {
                log.Fatal("Error reading destination configuration: %s", err.Error())
        }
 
+       // Initialize keep-rsync
        err = initializeKeepRsync()
        if err != nil {
                log.Fatal("Error configurating keep-rsync: %s", err.Error())
        }
+
+       // Copy blocks not found in dst from src
+       performKeepRsync()
 }
 
+var matchTrue = regexp.MustCompile("^(?i:1|yes|true)$")
+
 // Reads config from file
-func readConfigFromFile(filename string) (map[string]string, error) {
+func readConfigFromFile(filename string) (arvadosclient.APIConfig, error) {
+       var config arvadosclient.APIConfig
+
        content, err := ioutil.ReadFile(filename)
        if err != nil {
-               return nil, err
+               return config, err
        }
 
-       config := make(map[string]string)
        lines := strings.Split(string(content), "\n")
        for _, line := range lines {
                if line == "" {
                        continue
                }
                kv := strings.Split(line, "=")
-               config[kv[0]] = kv[1]
+
+               switch kv[0] {
+               case "ARVADOS_API_TOKEN":
+                       config.APIToken = kv[1]
+               case "ARVADOS_API_HOST":
+                       config.APIHost = kv[1]
+               case "ARVADOS_API_HOST_INSECURE":
+                       config.APIHostInsecure = matchTrue.MatchString(kv[1])
+               case "ARVADOS_EXTERNAL_CLIENT":
+                       config.ExternalClient = matchTrue.MatchString(kv[1])
+               case "ARVADOS_BLOB_SIGNING_KEY":
+                       blobSigningKey = kv[1]
+               }
        }
        return config, nil
 }
@@ -121,13 +144,13 @@ var (
 // Initializes keep-rsync using the config provided
 func initializeKeepRsync() (err error) {
        // arvSrc from srcConfig
-       arvSrc, err = arvadosclient.MakeArvadosClientWithConfig(srcConfig)
+       arvSrc, err = arvadosclient.New(srcConfig)
        if err != nil {
                return
        }
 
        // arvDst from dstConfig
-       arvDst, err = arvadosclient.MakeArvadosClientWithConfig(dstConfig)
+       arvDst, err = arvadosclient.New(dstConfig)
        if err != nil {
                return
        }
@@ -157,6 +180,123 @@ func initializeKeepRsync() (err error) {
                        return
                }
        }
+       kcDst.Want_replicas = replications
 
        return
 }
+
+// Get unique block locators from src and dst
+// Copy any blocks missing in dst
+func performKeepRsync() error {
+       // Get unique locators from src
+       srcIndex, err := getUniqueLocators(kcSrc, prefix)
+       if err != nil {
+               return err
+       }
+
+       // Get unique locators from dst
+       dstIndex, err := getUniqueLocators(kcDst, prefix)
+       if err != nil {
+               return err
+       }
+
+       // Get list of locators found in src, but missing in dst
+       toBeCopied := getMissingLocators(srcIndex, dstIndex)
+
+       // Copy each missing block to dst
+       copyBlocksToDst(toBeCopied)
+
+       return nil
+}
+
+// Get list of unique locators from the specified cluster
+func getUniqueLocators(kc *keepclient.KeepClient, prefix string) (map[string]bool, error) {
+       var indexBytes []byte
+
+       for uuid := range kc.LocalRoots() {
+               reader, err := kc.GetIndex(uuid, prefix)
+               if err != nil {
+                       return nil, err
+               }
+
+               var readBytes []byte
+               readBytes, err = ioutil.ReadAll(reader)
+               if err != nil {
+                       return nil, err
+               }
+
+               indexBytes = append(indexBytes, readBytes...)
+       }
+
+       // Got index; Now dedup it
+       locators := bytes.Split(indexBytes, []byte("\n"))
+
+       uniqueLocators := map[string]bool{}
+       for _, loc := range locators {
+               if len(loc) == 0 {
+                       continue
+               }
+
+               locator := string(bytes.Split(loc, []byte(" "))[0])
+               if _, ok := uniqueLocators[locator]; !ok {
+                       uniqueLocators[locator] = true
+               }
+       }
+       return uniqueLocators, nil
+}
+
+// Get list of locators that are in src but not in dst
+func getMissingLocators(srcLocators map[string]bool, dstLocators map[string]bool) []string {
+       var missingLocators []string
+       for locator := range srcLocators {
+               if _, ok := dstLocators[locator]; !ok {
+                       missingLocators = append(missingLocators, locator)
+               }
+       }
+       return missingLocators
+}
+
+// Copy blocks from src to dst; only those that are missing in dst are copied
+func copyBlocksToDst(toBeCopied []string) {
+       done := 0
+       total := len(toBeCopied)
+       var failed []string
+
+       for _, locator := range toBeCopied {
+               log.Printf("Getting block %d of %d", done+1, total)
+
+               log.Printf("Getting block: %v", locator)
+
+               reader, _, _, err := kcSrc.Get(locator)
+               if err != nil {
+                       log.Printf("Error getting block: %q %v", locator, err)
+                       failed = append(failed, locator)
+                       continue
+               }
+               data, err := ioutil.ReadAll(reader)
+               if err != nil {
+                       log.Printf("Error reading block data: %q %v", locator, err)
+                       failed = append(failed, locator)
+                       continue
+               }
+
+               log.Printf("Copying block: %q", locator)
+               _, rep, err := kcDst.PutB(data)
+               if err != nil {
+                       log.Printf("Error putting block data: %q %v", locator, err)
+                       failed = append(failed, locator)
+                       continue
+               }
+               if rep != replications {
+                       log.Printf("Failed to put enough number of replicas. Wanted: %d; Put: %d", replications, rep)
+                       failed = append(failed, locator)
+                       continue
+               }
+
+               done++
+               log.Printf("%.2f%% done", float64(done)/float64(total)*100)
+       }
+
+       log.Printf("Successfully copied to destination %d and failed %d out of a total of %d", done, len(failed), total)
+       log.Printf("Failed blocks %v", failed)
+}