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