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