tools/keep-rsync/keep-rsync.go

   1 package main
   2
   3 import (
   4         "bytes"
   5         "errors"
   6         "flag"
   7         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
   8         "git.curoverse.com/arvados.git/sdk/go/keepclient"
   9         "io/ioutil"
  10         "log"
  11         "regexp"
  12         "strings"
  13         "time"
  14 )
  15
  16 // keep-rsync arguments
  17 var (
  18         srcConfig           arvadosclient.APIConfig
  19         dstConfig           arvadosclient.APIConfig
  20         blobSigningKey      string
  21         srcKeepServicesJSON string
  22         dstKeepServicesJSON string
  23         replications        int
  24         prefix              string
  25 )
  26
  27 var srcConfigFile string
  28 var dstConfigFile string
  29
  30 func main() {
  31         flag.StringVar(
  32                 &srcConfigFile,
  33                 "src-config-file",
  34                 "",
  35                 "Source configuration filename with full path that contains "+
  36                         "an ARVADOS_API_TOKEN which is a valid datamanager token recognized by the source keep servers, "+
  37                         "ARVADOS_API_HOST, ARVADOS_API_HOST_INSECURE, ARVADOS_EXTERNAL_CLIENT and ARVADOS_BLOB_SIGNING_KEY.")
  38
  39         flag.StringVar(
  40                 &dstConfigFile,
  41                 "dst-config-file",
  42                 "",
  43                 "Destination configuration filename with full path that contains "+
  44                         "an ARVADOS_API_TOKEN which is a valid datamanager token recognized by the destination keep servers, "+
  45                         "ARVADOS_API_HOST, ARVADOS_API_HOST_INSECURE, ARVADOS_EXTERNAL_CLIENT and ARVADOS_BLOB_SIGNING_KEY.")
  46
  47         flag.StringVar(
  48                 &srcKeepServicesJSON,
  49                 "src-keep-services-json",
  50                 "",
  51                 "An optional list of available source keepservices. "+
  52                         "If not provided, this list is obtained from api server configured in src-config-file.")
  53
  54         flag.StringVar(
  55                 &dstKeepServicesJSON,
  56                 "dst-keep-services-json",
  57                 "",
  58                 "An optional list of available destination keepservices. "+
  59                         "If not provided, this list is obtained from api server configured in dst-config-file.")
  60
  61         flag.IntVar(
  62                 &replications,
  63                 "replications",
  64                 0,
  65                 "Number of replications to write to the destination.")
  66
  67         flag.StringVar(
  68                 &prefix,
  69                 "prefix",
  70                 "",
  71                 "Index prefix")
  72
  73         flag.Parse()
  74
  75         var err error
  76
  77         err = loadConfig()
  78         if err != nil {
  79                 log.Fatal("Error loading configuration from files: %s", err.Error())
  80         }
  81
  82         // Initialize keep-rsync
  83         err = initializeKeepRsync()
  84         if err != nil {
  85                 log.Fatal("Error configuring keep-rsync: %s", err.Error())
  86         }
  87
  88         // Copy blocks not found in dst from src
  89         err = performKeepRsync()
  90         if err != nil {
  91                 log.Fatal("Error while syncing data: %s", err.Error())
  92         }
  93 }
  94
  95 // Load src and dst config from given files
  96 func loadConfig() error {
  97         if srcConfigFile == "" {
  98                 return errors.New("-src-config-file must be specified")
  99         }
 100
 101         var err error
 102
 103         srcConfig, err = readConfigFromFile(srcConfigFile)
 104         if err != nil {
 105                 log.Printf("Error reading source configuration: %s", err.Error())
 106                 return err
 107         }
 108
 109         if dstConfigFile == "" {
 110                 return errors.New("-dst-config-file must be specified")
 111         }
 112         dstConfig, err = readConfigFromFile(dstConfigFile)
 113         if err != nil {
 114                 log.Printf("Error reading destination configuration: %s", err.Error())
 115         }
 116
 117         return err
 118 }
 119
 120 var matchTrue = regexp.MustCompile("^(?i:1|yes|true)$")
 121
 122 // Reads config from file
 123 func readConfigFromFile(filename string) (arvadosclient.APIConfig, error) {
 124         var config arvadosclient.APIConfig
 125
 126         content, err := ioutil.ReadFile(filename)
 127         if err != nil {
 128                 return config, err
 129         }
 130
 131         lines := strings.Split(string(content), "\n")
 132         for _, line := range lines {
 133                 if line == "" {
 134                         continue
 135                 }
 136                 kv := strings.Split(line, "=")
 137
 138                 switch kv[0] {
 139                 case "ARVADOS_API_TOKEN":
 140                         config.APIToken = kv[1]
 141                 case "ARVADOS_API_HOST":
 142                         config.APIHost = kv[1]
 143                 case "ARVADOS_API_HOST_INSECURE":
 144                         config.APIHostInsecure = matchTrue.MatchString(kv[1])
 145                 case "ARVADOS_EXTERNAL_CLIENT":
 146                         config.ExternalClient = matchTrue.MatchString(kv[1])
 147                 case "ARVADOS_BLOB_SIGNING_KEY":
 148                         blobSigningKey = kv[1]
 149                 }
 150         }
 151         return config, nil
 152 }
 153
 154 // keep-rsync source and destination clients
 155 var (
 156         arvSrc arvadosclient.ArvadosClient
 157         arvDst arvadosclient.ArvadosClient
 158         kcSrc  *keepclient.KeepClient
 159         kcDst  *keepclient.KeepClient
 160 )
 161
 162 // Initializes keep-rsync using the config provided
 163 func initializeKeepRsync() (err error) {
 164         // arvSrc from srcConfig
 165         arvSrc, err = arvadosclient.New(srcConfig)
 166         if err != nil {
 167                 return
 168         }
 169
 170         // arvDst from dstConfig
 171         arvDst, err = arvadosclient.New(dstConfig)
 172         if err != nil {
 173                 return
 174         }
 175
 176         // Get default replications value from destination, if it is not already provided
 177         if replications == 0 {
 178                 value, err := arvDst.Discovery("defaultCollectionReplication")
 179                 if err == nil {
 180                         replications = int(value.(float64))
 181                 } else {
 182                         replications = 2
 183                 }
 184         }
 185
 186         // if srcKeepServicesJSON is provided, use it to load services; else, use DiscoverKeepServers
 187         if srcKeepServicesJSON == "" {
 188                 kcSrc, err = keepclient.MakeKeepClient(&arvSrc)
 189                 if err != nil {
 190                         return
 191                 }
 192         } else {
 193                 kcSrc, err = keepclient.MakeKeepClientFromJSON(&arvSrc, srcKeepServicesJSON)
 194                 if err != nil {
 195                         return
 196                 }
 197         }
 198
 199         // if dstKeepServicesJSON is provided, use it to load services; else, use DiscoverKeepServers
 200         if dstKeepServicesJSON == "" {
 201                 kcDst, err = keepclient.MakeKeepClient(&arvDst)
 202                 if err != nil {
 203                         return
 204                 }
 205         } else {
 206                 kcDst, err = keepclient.MakeKeepClientFromJSON(&arvDst, dstKeepServicesJSON)
 207                 if err != nil {
 208                         return
 209                 }
 210         }
 211         kcDst.Want_replicas = replications
 212
 213         return
 214 }
 215
 216 // Get unique block locators from src and dst
 217 // Copy any blocks missing in dst
 218 func performKeepRsync() 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         err = copyBlocksToDst(toBeCopied)
 236
 237         return err
 238 }
 239
 240 // Get list of unique locators from the specified cluster
 241 func getUniqueLocators(kc *keepclient.KeepClient, indexPrefix string) (map[string]bool, error) {
 242         var indexBytes []byte
 243
 244         for uuid := range kc.LocalRoots() {
 245                 reader, err := kc.GetIndex(uuid, indexPrefix)
 246                 if err != nil {
 247                         return nil, err
 248                 }
 249
 250                 var readBytes []byte
 251                 readBytes, err = ioutil.ReadAll(reader)
 252                 if err != nil {
 253                         return nil, err
 254                 }
 255
 256                 indexBytes = append(indexBytes, readBytes...)
 257         }
 258
 259         // Got index; Now dedup it
 260         locators := bytes.Split(indexBytes, []byte("\n"))
 261
 262         uniqueLocators := map[string]bool{}
 263         for _, loc := range locators {
 264                 if len(loc) == 0 {
 265                         continue
 266                 }
 267
 268                 locator := string(bytes.Split(loc, []byte(" "))[0])
 269                 if _, ok := uniqueLocators[locator]; !ok {
 270                         uniqueLocators[locator] = true
 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 map[string]bool, 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) error {
 289         done := 0
 290         total := len(toBeCopied)
 291
 292         for _, locator := range toBeCopied {
 293                 log.Printf("Getting block %d of %d", done+1, total)
 294
 295                 log.Printf("Getting block: %v", locator)
 296
 297                 getLocator := locator
 298                 expiresAt := time.Now().AddDate(0, 0, 1)
 299                 if blobSigningKey != "" {
 300                         getLocator = keepclient.SignLocator(getLocator, arvSrc.ApiToken, expiresAt, []byte(blobSigningKey))
 301                 }
 302
 303                 reader, _, _, err := kcSrc.Get(getLocator)
 304                 if err != nil {
 305                         log.Printf("Error getting block: %q %v", locator, err)
 306                         return err
 307                 }
 308                 data, err := ioutil.ReadAll(reader)
 309                 if err != nil {
 310                         log.Printf("Error reading block data: %q %v", locator, err)
 311                         return err
 312                 }
 313
 314                 log.Printf("Copying block: %q", locator)
 315                 _, _, err = kcDst.PutB(data)
 316                 if err != nil {
 317                         log.Printf("Error putting block data: %q %v", locator, err)
 318                         return err
 319                 }
 320
 321                 done++
 322                 log.Printf("%.2f%% done", float64(done)/float64(total)*100)
 323         }
 324
 325         log.Printf("Successfully copied to destination %d blocks.", total)
 326         return nil
 327 }