tools/keep-rsync/keep-rsync.go

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