9998: Change count=true/false to exact/none. Add tests.

[arvados.git] / services / keep-balance / main.go

diff --git a/services/keep-balance/main.go b/services/keep-balance/main.go
index 42a8d635b131f14a78f2ac7ced915a32d0f5a7f7..310c77a21c228c12779de5ac39c686cc08d7b624 100644 (file)
--- a/services/keep-balance/main.go
+++ b/services/keep-balance/main.go
@@ -3,7 +3,6 @@ package main
 import (
        "encoding/json"
        "flag"
-       "io/ioutil"
        "log"
        "os"
        "os/signal"
@@ -11,8 +10,11 @@ import (
        "time"
 
        "git.curoverse.com/arvados.git/sdk/go/arvados"
+       "git.curoverse.com/arvados.git/sdk/go/config"
 )
 
+const defaultConfigPath = "/etc/arvados/keep-balance/keep-balance.yml"
+
 // Config specifies site configuration, like API credentials and the
 // choice of which servers are to be balanced.
 //
@@ -28,6 +30,14 @@ type Config struct {
 
        // How often to check
        RunPeriod arvados.Duration
+
+       // Number of collections to request in each API call
+       CollectionBatchSize int
+
+       // Max collections to buffer in memory (bigger values consume
+       // more memory, but can reduce store-and-forward latency when
+       // fetching pages)
+       CollectionBuffers int
 }
 
 // RunOptions controls runtime behavior. The flags/options that belong
@@ -43,6 +53,12 @@ type RunOptions struct {
        CommitTrash bool
        Logger      *log.Logger
        Dumper      *log.Logger
+
+       // SafeRendezvousState from the most recent balance operation,
+       // or "" if unknown. If this changes from one run to the next,
+       // we need to watch out for races. See
+       // (*Balancer)ClearTrashLists.
+       SafeRendezvousState string
 }
 
 var debugf = func(string, ...interface{}) {}
@@ -51,10 +67,10 @@ func main() {
        var config Config
        var runOptions RunOptions
 
-       configPath := flag.String("config", "",
-               "`path` of json configuration file")
+       configPath := flag.String("config", defaultConfigPath,
+               "`path` of JSON or YAML configuration file")
        serviceListPath := flag.String("config.KeepServiceList", "",
-               "`path` of json file with list of keep services to balance, as given by \"arv keep_service list\" "+
+               "`path` of JSON or YAML file with list of keep services to balance, as given by \"arv keep_service list\" "+
                        "(default: config[\"KeepServiceList\"], or if none given, get all available services and filter by config[\"KeepServiceTypes\"])")
        flag.BoolVar(&runOptions.Once, "once", false,
                "balance once and then exit")
@@ -67,12 +83,9 @@ func main() {
        flag.Usage = usage
        flag.Parse()
 
-       if *configPath == "" {
-               log.Fatal("You must specify a config file (see `keep-balance -help`)")
-       }
-       mustReadJSON(&config, *configPath)
+       mustReadConfig(&config, *configPath)
        if *serviceListPath != "" {
-               mustReadJSON(&config.KeepServiceList, *serviceListPath)
+               mustReadConfig(&config.KeepServiceList, *serviceListPath)
        }
 
        if *debugFlag {
@@ -90,7 +103,7 @@ func main() {
        if err != nil {
                // (don't run)
        } else if runOptions.Once {
-               err = (&Balancer{}).Run(config, runOptions)
+               _, err = (&Balancer{}).Run(config, runOptions)
        } else {
                err = RunForever(config, runOptions, nil)
        }
@@ -99,11 +112,9 @@ func main() {
        }
 }
 
-func mustReadJSON(dst interface{}, path string) {
-       if buf, err := ioutil.ReadFile(path); err != nil {
-               log.Fatalf("Reading %q: %v", path, err)
-       } else if err = json.Unmarshal(buf, dst); err != nil {
-               log.Fatalf("Decoding %q: %v", path, err)
+func mustReadConfig(dst interface{}, path string) {
+       if err := config.LoadFile(dst, path); err != nil {
+               log.Fatal(err)
        }
 }
 
@@ -130,7 +141,9 @@ func RunForever(config Config, runOptions RunOptions, stop <-chan interface{}) e
                        logger.Print("=======  Consider using -commit-pulls and -commit-trash flags.")
                }
 
-               err := (&Balancer{}).Run(config, runOptions)
+               bal := &Balancer{}
+               var err error
+               runOptions, err = bal.Run(config, runOptions)
                if err != nil {
                        logger.Print("run failed: ", err)
                } else {