Merge branch '13219-jobs-time-limit'

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

diff --git a/services/keep-balance/main.go b/services/keep-balance/main.go
index da4fb62a6ac8f0772e550c5eb2a103de7904ad7e..90235cbf3188d91bc274412ddd5522dc639fa812 100644 (file)
--- a/services/keep-balance/main.go
+++ b/services/keep-balance/main.go
@@ -1,18 +1,28 @@
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+

 package main
 
 import (
        "encoding/json"
        "flag"
 package main
 
 import (
        "encoding/json"
        "flag"

-       "io/ioutil"
+       "fmt"

        "log"
        "log"

+       "net/http"

        "os"
        "os/signal"
        "syscall"
        "time"
 
        "git.curoverse.com/arvados.git/sdk/go/arvados"
        "os"
        "os/signal"
        "syscall"
        "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.
 //
 // Config specifies site configuration, like API credentials and the
 // choice of which servers are to be balanced.
 //


@@ -36,6 +46,9 @@ type Config struct {
        // more memory, but can reduce store-and-forward latency when
        // fetching pages)
        CollectionBuffers int
        // more memory, but can reduce store-and-forward latency when
        // fetching pages)
        CollectionBuffers int

+
+       // Timeout for outgoing http request/response cycle.
+       RequestTimeout arvados.Duration

 }
 
 // RunOptions controls runtime behavior. The flags/options that belong
 }
 
 // RunOptions controls runtime behavior. The flags/options that belong


@@ -62,13 +75,13 @@ type RunOptions struct {
 var debugf = func(string, ...interface{}) {}
 
 func main() {
 var debugf = func(string, ...interface{}) {}
 
 func main() {

-       var config Config
+       var cfg Config

        var runOptions RunOptions
 
        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", "",
        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")
                        "(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 +89,41 @@ 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)")
                "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")
        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()
 
        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("keep-balance %s\n", version)
+               return

        }
        }

-       mustReadJSON(&config, *configPath)
+
+       mustReadConfig(&cfg, *configPath)

        if *serviceListPath != "" {
        if *serviceListPath != "" {

-               mustReadJSON(&config.KeepServiceList, *serviceListPath)
+               mustReadConfig(&cfg.KeepServiceList, *serviceListPath)

        }
 
        }
 

+       if *dumpConfig {
+               log.Fatal(config.DumpAndExit(cfg))
+       }
+
+       to := time.Duration(cfg.RequestTimeout)
+       if to == 0 {
+               to = 30 * time.Minute
+       }
+       arvados.DefaultSecureClient.Timeout = to
+       arvados.InsecureHTTPClient.Timeout = to
+       http.DefaultClient.Timeout = to
+
+       log.Printf("keep-balance %s started", version)
+

        if *debugFlag {
                debugf = log.Printf
        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)
                        log.Fatal(err)
                } else {
                        log.Printf("config is %s", j)


@@ -100,24 +132,22 @@ func main() {
        if *dumpFlag {
                runOptions.Dumper = log.New(os.Stdout, "", log.LstdFlags)
        }
        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 {
        if err != nil {
                // (don't run)
        } else if runOptions.Once {

-               _, err = (&Balancer{}).Run(config, runOptions)
+               _, err = (&Balancer{}).Run(cfg, runOptions)

        } else {
        } else {

-               err = RunForever(config, runOptions, nil)
+               err = RunForever(cfg, runOptions, nil)

        }
        if err != nil {
                log.Fatal(err)
        }
 }
 
        }
        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)

        }
 }
 
        }
 }