services/keep-balance/main.go

   1 package main
   2
   3 import (
   4         "encoding/json"
   5         "flag"
   6         "log"
   7         "os"
   8         "os/signal"
   9         "syscall"
  10         "time"
  11
  12         "git.curoverse.com/arvados.git/sdk/go/arvados"
  13         "git.curoverse.com/arvados.git/sdk/go/config"
  14         "github.com/ghodss/yaml"
  15 )
  16
  17 const defaultConfigPath = "/etc/arvados/keep-balance/keep-balance.yml"
  18
  19 // Config specifies site configuration, like API credentials and the
  20 // choice of which servers are to be balanced.
  21 //
  22 // Config is loaded from a JSON config file (see usage()).
  23 type Config struct {
  24         // Arvados API endpoint and credentials.
  25         Client arvados.Client
  26
  27         // List of service types (e.g., "disk") to balance.
  28         KeepServiceTypes []string
  29
  30         KeepServiceList arvados.KeepServiceList
  31
  32         // How often to check
  33         RunPeriod arvados.Duration
  34
  35         // Number of collections to request in each API call
  36         CollectionBatchSize int
  37
  38         // Max collections to buffer in memory (bigger values consume
  39         // more memory, but can reduce store-and-forward latency when
  40         // fetching pages)
  41         CollectionBuffers int
  42 }
  43
  44 // RunOptions controls runtime behavior. The flags/options that belong
  45 // here are the ones that are useful for interactive use. For example,
  46 // "CommitTrash" is a runtime option rather than a config item because
  47 // it invokes a troubleshooting feature rather than expressing how
  48 // balancing is meant to be done at a given site.
  49 //
  50 // RunOptions fields are controlled by command line flags.
  51 type RunOptions struct {
  52         Once        bool
  53         CommitPulls bool
  54         CommitTrash bool
  55         Logger      *log.Logger
  56         Dumper      *log.Logger
  57
  58         // SafeRendezvousState from the most recent balance operation,
  59         // or "" if unknown. If this changes from one run to the next,
  60         // we need to watch out for races. See
  61         // (*Balancer)ClearTrashLists.
  62         SafeRendezvousState string
  63 }
  64
  65 var debugf = func(string, ...interface{}) {}
  66
  67 func main() {
  68         var config Config
  69         var runOptions RunOptions
  70
  71         configPath := flag.String("config", defaultConfigPath,
  72                 "`path` of JSON or YAML configuration file")
  73         serviceListPath := flag.String("config.KeepServiceList", "",
  74                 "`path` of JSON or YAML file with list of keep services to balance, as given by \"arv keep_service list\" "+
  75                         "(default: config[\"KeepServiceList\"], or if none given, get all available services and filter by config[\"KeepServiceTypes\"])")
  76         flag.BoolVar(&runOptions.Once, "once", false,
  77                 "balance once and then exit")
  78         flag.BoolVar(&runOptions.CommitPulls, "commit-pulls", false,
  79                 "send pull requests (make more replicas of blocks that are underreplicated or are not in optimal rendezvous probe order)")
  80         flag.BoolVar(&runOptions.CommitTrash, "commit-trash", false,
  81                 "send trash requests (delete unreferenced old blocks, and excess replicas of overreplicated blocks)")
  82         dumpConfig := flag.Bool("dump-config", false, "write current configuration to stdout and exit")
  83         dumpFlag := flag.Bool("dump", false, "dump details for each block to stdout")
  84         debugFlag := flag.Bool("debug", false, "enable debug messages")
  85         flag.Usage = usage
  86         flag.Parse()
  87
  88         mustReadConfig(&config, *configPath)
  89         if *serviceListPath != "" {
  90                 mustReadConfig(&config.KeepServiceList, *serviceListPath)
  91         }
  92
  93         if *dumpConfig {
  94                 y, err := yaml.Marshal(config)
  95                 if err != nil {
  96                         log.Fatal(err)
  97                 }
  98                 os.Stdout.Write(y)
  99                 os.Exit(0)
 100         }
 101
 102         if *debugFlag {
 103                 debugf = log.Printf
 104                 if j, err := json.Marshal(config); err != nil {
 105                         log.Fatal(err)
 106                 } else {
 107                         log.Printf("config is %s", j)
 108                 }
 109         }
 110         if *dumpFlag {
 111                 runOptions.Dumper = log.New(os.Stdout, "", log.LstdFlags)
 112         }
 113         err := CheckConfig(config, runOptions)
 114         if err != nil {
 115                 // (don't run)
 116         } else if runOptions.Once {
 117                 _, err = (&Balancer{}).Run(config, runOptions)
 118         } else {
 119                 err = RunForever(config, runOptions, nil)
 120         }
 121         if err != nil {
 122                 log.Fatal(err)
 123         }
 124 }
 125
 126 func mustReadConfig(dst interface{}, path string) {
 127         if err := config.LoadFile(dst, path); err != nil {
 128                 log.Fatal(err)
 129         }
 130 }
 131
 132 // RunForever runs forever, or (for testing purposes) until the given
 133 // stop channel is ready to receive.
 134 func RunForever(config Config, runOptions RunOptions, stop <-chan interface{}) error {
 135         if runOptions.Logger == nil {
 136                 runOptions.Logger = log.New(os.Stderr, "", log.LstdFlags)
 137         }
 138         logger := runOptions.Logger
 139
 140         ticker := time.NewTicker(time.Duration(config.RunPeriod))
 141
 142         // The unbuffered channel here means we only hear SIGUSR1 if
 143         // it arrives while we're waiting in select{}.
 144         sigUSR1 := make(chan os.Signal)
 145         signal.Notify(sigUSR1, syscall.SIGUSR1)
 146
 147         logger.Printf("starting up: will scan every %v and on SIGUSR1", config.RunPeriod)
 148
 149         for {
 150                 if !runOptions.CommitPulls && !runOptions.CommitTrash {
 151                         logger.Print("WARNING: Will scan periodically, but no changes will be committed.")
 152                         logger.Print("=======  Consider using -commit-pulls and -commit-trash flags.")
 153                 }
 154
 155                 bal := &Balancer{}
 156                 var err error
 157                 runOptions, err = bal.Run(config, runOptions)
 158                 if err != nil {
 159                         logger.Print("run failed: ", err)
 160                 } else {
 161                         logger.Print("run succeeded")
 162                 }
 163
 164                 select {
 165                 case <-stop:
 166                         signal.Stop(sigUSR1)
 167                         return nil
 168                 case <-ticker.C:
 169                         logger.Print("timer went off")
 170                 case <-sigUSR1:
 171                         logger.Print("received SIGUSR1, resetting timer")
 172                         // Reset the timer so we don't start the N+1st
 173                         // run too soon after the Nth run is triggered
 174                         // by SIGUSR1.
 175                         ticker.Stop()
 176                         ticker = time.NewTicker(time.Duration(config.RunPeriod))
 177                 }
 178                 logger.Print("starting next run")
 179         }
 180 }