tools/keep-rsync/keep-rsync.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         "bufio"
   9         "crypto/tls"
  10         "errors"
  11         "flag"
  12         "fmt"
  13         "io/ioutil"
  14         "log"
  15         "net/http"
  16         "os"
  17         "strings"
  18         "time"
  19
  20         "git.arvados.org/arvados.git/lib/cmd"
  21         "git.arvados.org/arvados.git/sdk/go/arvadosclient"
  22         "git.arvados.org/arvados.git/sdk/go/keepclient"
  23 )
  24
  25 var version = "dev"
  26
  27 func main() {
  28         err := doMain()
  29         if err != nil {
  30                 log.Fatalf("%v", err)
  31         }
  32 }
  33
  34 func doMain() error {
  35         flags := flag.NewFlagSet("keep-rsync", flag.ExitOnError)
  36
  37         srcConfigFile := flags.String(
  38                 "src",
  39                 "",
  40                 "Source configuration filename. May be either a pathname to a config file, or (for example) 'foo' as shorthand for $HOME/.config/arvados/foo.conf file. This file is expected to specify the values for ARVADOS_API_TOKEN, ARVADOS_API_HOST, ARVADOS_API_HOST_INSECURE, and ARVADOS_BLOB_SIGNING_KEY for the source.")
  41
  42         dstConfigFile := flags.String(
  43                 "dst",
  44                 "",
  45                 "Destination configuration filename. May be either a pathname to a config file, or (for example) 'foo' as shorthand for $HOME/.config/arvados/foo.conf file. This file is expected to specify the values for ARVADOS_API_TOKEN, ARVADOS_API_HOST, and ARVADOS_API_HOST_INSECURE for the destination.")
  46
  47         srcKeepServicesJSON := flags.String(
  48                 "src-keep-services-json",
  49                 "",
  50                 "An optional list of available source keepservices. "+
  51                         "If not provided, this list is obtained from api server configured in src-config-file.")
  52
  53         dstKeepServicesJSON := flags.String(
  54                 "dst-keep-services-json",
  55                 "",
  56                 "An optional list of available destination keepservices. "+
  57                         "If not provided, this list is obtained from api server configured in dst-config-file.")
  58
  59         replications := flags.Int(
  60                 "replications",
  61                 0,
  62                 "Number of replications to write to the destination. If replications not specified, "+
  63                         "default replication level configured on destination server will be used.")
  64
  65         prefix := flags.String(
  66                 "prefix",
  67                 "",
  68                 "Index prefix")
  69
  70         srcBlobSignatureTTLFlag := flags.Duration(
  71                 "src-blob-signature-ttl",
  72                 0,
  73                 "Lifetime of blob permission signatures on source keepservers. If not provided, this will be retrieved from the API server's discovery document.")
  74
  75         getVersion := flags.Bool(
  76                 "version",
  77                 false,
  78                 "Print version information and exit.")
  79
  80         if ok, code := cmd.ParseFlags(flags, os.Args[0], os.Args[1:], "", os.Stderr); !ok {
  81                 os.Exit(code)
  82         } else if *getVersion {
  83                 fmt.Printf("%s %s\n", os.Args[0], version)
  84                 os.Exit(0)
  85         }
  86
  87         srcConfig, srcBlobSigningKey, err := loadConfig(*srcConfigFile)
  88         if err != nil {
  89                 return fmt.Errorf("Error loading src configuration from file: %s", err.Error())
  90         }
  91
  92         dstConfig, _, err := loadConfig(*dstConfigFile)
  93         if err != nil {
  94                 return fmt.Errorf("Error loading dst configuration from file: %s", err.Error())
  95         }
  96
  97         // setup src and dst keepclients
  98         kcSrc, srcBlobSignatureTTL, err := setupKeepClient(srcConfig, *srcKeepServicesJSON, false, 0, *srcBlobSignatureTTLFlag)
  99         if err != nil {
 100                 return fmt.Errorf("Error configuring src keepclient: %s", err.Error())
 101         }
 102
 103         kcDst, _, err := setupKeepClient(dstConfig, *dstKeepServicesJSON, true, *replications, 0)
 104         if err != nil {
 105                 return fmt.Errorf("Error configuring dst keepclient: %s", err.Error())
 106         }
 107
 108         // Copy blocks not found in dst from src
 109         err = performKeepRsync(kcSrc, kcDst, srcBlobSignatureTTL, srcBlobSigningKey, *prefix)
 110         if err != nil {
 111                 return fmt.Errorf("Error while syncing data: %s", err.Error())
 112         }
 113
 114         return nil
 115 }
 116
 117 type apiConfig struct {
 118         APIToken        string
 119         APIHost         string
 120         APIHostInsecure bool
 121         ExternalClient  bool
 122 }
 123
 124 // Load src and dst config from given files
 125 func loadConfig(configFile string) (config apiConfig, blobSigningKey string, err error) {
 126         if configFile == "" {
 127                 return config, blobSigningKey, errors.New("config file not specified")
 128         }
 129
 130         config, blobSigningKey, err = readConfigFromFile(configFile)
 131         if err != nil {
 132                 return config, blobSigningKey, fmt.Errorf("Error reading config file: %v", err)
 133         }
 134
 135         return
 136 }
 137
 138 // Read config from file
 139 func readConfigFromFile(filename string) (config apiConfig, blobSigningKey string, err error) {
 140         if !strings.Contains(filename, "/") {
 141                 filename = os.Getenv("HOME") + "/.config/arvados/" + filename + ".conf"
 142         }
 143
 144         content, err := ioutil.ReadFile(filename)
 145
 146         if err != nil {
 147                 return config, "", err
 148         }
 149
 150         lines := strings.Split(string(content), "\n")
 151         for _, line := range lines {
 152                 if line == "" {
 153                         continue
 154                 }
 155
 156                 kv := strings.SplitN(line, "=", 2)
 157                 key := strings.TrimSpace(kv[0])
 158                 value := strings.TrimSpace(kv[1])
 159
 160                 switch key {
 161                 case "ARVADOS_API_TOKEN":
 162                         config.APIToken = value
 163                 case "ARVADOS_API_HOST":
 164                         config.APIHost = value
 165                 case "ARVADOS_API_HOST_INSECURE":
 166                         config.APIHostInsecure = arvadosclient.StringBool(value)
 167                 case "ARVADOS_EXTERNAL_CLIENT":
 168                         config.ExternalClient = arvadosclient.StringBool(value)
 169                 case "ARVADOS_BLOB_SIGNING_KEY":
 170                         blobSigningKey = value
 171                 }
 172         }
 173         return
 174 }
 175
 176 // setup keepclient using the config provided
 177 func setupKeepClient(config apiConfig, keepServicesJSON string, isDst bool, replications int, srcBlobSignatureTTL time.Duration) (kc *keepclient.KeepClient, blobSignatureTTL time.Duration, err error) {
 178         arv := arvadosclient.ArvadosClient{
 179                 ApiToken:    config.APIToken,
 180                 ApiServer:   config.APIHost,
 181                 ApiInsecure: config.APIHostInsecure,
 182                 Client: &http.Client{Transport: &http.Transport{
 183                         TLSClientConfig: &tls.Config{InsecureSkipVerify: config.APIHostInsecure}}},
 184                 External: config.ExternalClient,
 185         }
 186
 187         // If keepServicesJSON is provided, use it instead of service discovery
 188         if keepServicesJSON == "" {
 189                 kc, err = keepclient.MakeKeepClient(&arv)
 190                 if err != nil {
 191                         return nil, 0, err
 192                 }
 193         } else {
 194                 kc = keepclient.New(&arv)
 195                 err = kc.LoadKeepServicesFromJSON(keepServicesJSON)
 196                 if err != nil {
 197                         return kc, 0, err
 198                 }
 199         }
 200
 201         if isDst {
 202                 // Get default replications value from destination, if it is not already provided
 203                 if replications == 0 {
 204                         value, err := arv.Discovery("defaultCollectionReplication")
 205                         if err == nil {
 206                                 replications = int(value.(float64))
 207                         } else {
 208                                 return nil, 0, err
 209                         }
 210                 }
 211
 212                 kc.Want_replicas = replications
 213         }
 214
 215         // If srcBlobSignatureTTL is not provided, get it from API server discovery doc
 216         blobSignatureTTL = srcBlobSignatureTTL
 217         if !isDst && srcBlobSignatureTTL == 0 {
 218                 value, err := arv.Discovery("blobSignatureTtl")
 219                 if err == nil {
 220                         blobSignatureTTL = time.Duration(int(value.(float64))) * time.Second
 221                 } else {
 222                         return nil, 0, err
 223                 }
 224         }
 225
 226         return kc, blobSignatureTTL, nil
 227 }
 228
 229 // Get unique block locators from src and dst
 230 // Copy any blocks missing in dst
 231 func performKeepRsync(kcSrc, kcDst *keepclient.KeepClient, srcBlobSignatureTTL time.Duration, blobSigningKey, prefix string) error {
 232         // Get unique locators from src
 233         srcIndex, err := getUniqueLocators(kcSrc, prefix)
 234         if err != nil {
 235                 return err
 236         }
 237
 238         // Get unique locators from dst
 239         dstIndex, err := getUniqueLocators(kcDst, prefix)
 240         if err != nil {
 241                 return err
 242         }
 243
 244         // Get list of locators found in src, but missing in dst
 245         toBeCopied := getMissingLocators(srcIndex, dstIndex)
 246
 247         // Copy each missing block to dst
 248         log.Printf("Before keep-rsync, there are %d blocks in src and %d blocks in dst. Start copying %d blocks from src not found in dst.",
 249                 len(srcIndex), len(dstIndex), len(toBeCopied))
 250
 251         err = copyBlocksToDst(toBeCopied, kcSrc, kcDst, srcBlobSignatureTTL, blobSigningKey)
 252
 253         return err
 254 }
 255
 256 // Get list of unique locators from the specified cluster
 257 func getUniqueLocators(kc *keepclient.KeepClient, prefix string) (map[string]bool, error) {
 258         uniqueLocators := map[string]bool{}
 259
 260         // Get index and dedup
 261         for uuid := range kc.LocalRoots() {
 262                 reader, err := kc.GetIndex(uuid, prefix)
 263                 if err != nil {
 264                         return uniqueLocators, err
 265                 }
 266                 scanner := bufio.NewScanner(reader)
 267                 for scanner.Scan() {
 268                         uniqueLocators[strings.Split(scanner.Text(), " ")[0]] = true
 269                 }
 270         }
 271
 272         return uniqueLocators, nil
 273 }
 274
 275 // Get list of locators that are in src but not in dst
 276 func getMissingLocators(srcLocators, dstLocators map[string]bool) []string {
 277         var missingLocators []string
 278         for locator := range srcLocators {
 279                 if _, ok := dstLocators[locator]; !ok {
 280                         missingLocators = append(missingLocators, locator)
 281                 }
 282         }
 283         return missingLocators
 284 }
 285
 286 // Copy blocks from src to dst; only those that are missing in dst are copied
 287 func copyBlocksToDst(toBeCopied []string, kcSrc, kcDst *keepclient.KeepClient, srcBlobSignatureTTL time.Duration, blobSigningKey string) error {
 288         total := len(toBeCopied)
 289
 290         startedAt := time.Now()
 291         for done, locator := range toBeCopied {
 292                 if done == 0 {
 293                         log.Printf("Copying data block %d of %d (%.2f%% done): %v", done+1, total,
 294                                 float64(done)/float64(total)*100, locator)
 295                 } else {
 296                         timePerBlock := time.Since(startedAt) / time.Duration(done)
 297                         log.Printf("Copying data block %d of %d (%.2f%% done, %v est. time remaining): %v", done+1, total,
 298                                 float64(done)/float64(total)*100, timePerBlock*time.Duration(total-done), locator)
 299                 }
 300
 301                 getLocator := locator
 302                 expiresAt := time.Now().AddDate(0, 0, 1)
 303                 if blobSigningKey != "" {
 304                         getLocator = keepclient.SignLocator(getLocator, kcSrc.Arvados.ApiToken, expiresAt, srcBlobSignatureTTL, []byte(blobSigningKey))
 305                 }
 306
 307                 reader, len, _, err := kcSrc.Get(getLocator)
 308                 if err != nil {
 309                         return fmt.Errorf("Error getting block: %v %v", locator, err)
 310                 }
 311
 312                 _, _, err = kcDst.PutHR(getLocator[:32], reader, len)
 313                 if err != nil {
 314                         return fmt.Errorf("Error copying data block: %v %v", locator, err)
 315                 }
 316         }
 317
 318         log.Printf("Successfully copied to destination %d blocks.", total)
 319         return nil
 320 }