tools/keep-rsync/keep-rsync.go

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