10666: Replaced 'version' package with 'version' var

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

diff --git a/services/keep-balance/main.go b/services/keep-balance/main.go
index da4fb62a6ac8f0772e550c5eb2a103de7904ad7e..ecc3f2696f6ab2ef29b4ece9d3071dd75c91389b 100644 (file)
--- a/services/keep-balance/main.go
+++ b/services/keep-balance/main.go
@@ -1,9 +1,13 @@
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
 package main
 
 import (
        "encoding/json"
        "flag"
-       "io/ioutil"
+       "fmt"
        "log"
        "os"
        "os/signal"
@@ -11,8 +15,13 @@ import (
        "time"
 
        "git.curoverse.com/arvados.git/sdk/go/arvados"
+       "git.curoverse.com/arvados.git/sdk/go/config"
 )
 
+var version = "dev"
+
+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.
 //
@@ -62,13 +71,13 @@ type RunOptions struct {
 var debugf = func(string, ...interface{}) {}
 
 func main() {
-       var config Config
+       var cfg 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")
@@ -76,22 +85,33 @@ func main() {
                "send pull requests (make more replicas of blocks that are underreplicated or are not in optimal rendezvous probe order)")
        flag.BoolVar(&runOptions.CommitTrash, "commit-trash", false,
                "send trash requests (delete unreferenced old blocks, and excess replicas of overreplicated blocks)")
+       dumpConfig := flag.Bool("dump-config", false, "write current configuration to stdout and exit")
        dumpFlag := flag.Bool("dump", false, "dump details for each block to stdout")
        debugFlag := flag.Bool("debug", false, "enable debug messages")
+       getVersion := flag.Bool("version", false, "Print version information and exit.")
        flag.Usage = usage
        flag.Parse()
 
-       if *configPath == "" {
-               log.Fatal("You must specify a config file (see `keep-balance -help`)")
+       // Print version information if requested
+       if *getVersion {
+               fmt.Printf("Version: %s\n", version)
+               os.Exit(0)
        }
-       mustReadJSON(&config, *configPath)
+
+       mustReadConfig(&cfg, *configPath)
        if *serviceListPath != "" {
-               mustReadJSON(&config.KeepServiceList, *serviceListPath)
+               mustReadConfig(&cfg.KeepServiceList, *serviceListPath)
        }
 
+       if *dumpConfig {
+               log.Fatal(config.DumpAndExit(cfg))
+       }
+
+       log.Printf("keep-balance %q started", version)
+
        if *debugFlag {
                debugf = log.Printf
-               if j, err := json.Marshal(config); err != nil {
+               if j, err := json.Marshal(cfg); err != nil {
                        log.Fatal(err)
                } else {
                        log.Printf("config is %s", j)
@@ -100,24 +120,22 @@ func main() {
        if *dumpFlag {
                runOptions.Dumper = log.New(os.Stdout, "", log.LstdFlags)
        }
-       err := CheckConfig(config, runOptions)
+       err := CheckConfig(cfg, runOptions)
        if err != nil {
                // (don't run)
        } else if runOptions.Once {
-               _, err = (&Balancer{}).Run(config, runOptions)
+               _, err = (&Balancer{}).Run(cfg, runOptions)
        } else {
-               err = RunForever(config, runOptions, nil)
+               err = RunForever(cfg, runOptions, nil)
        }
        if err != nil {
                log.Fatal(err)
        }
 }
 
-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)
        }
 }