X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/139200027a3192260b5ea7c2d0c93a8eb5f8eb7e..94f2b439783a8e63d6d7b9ba2760f54fc642a8fb:/services/keep-balance/main.go

diff --git a/services/keep-balance/main.go b/services/keep-balance/main.go
index 42a8d635b1..90235cbf31 100644
--- 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"
-	"io/ioutil"
+	"fmt"
 	"log"
+	"net/http"
 	"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.
 //
@@ -28,6 +38,17 @@ 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
+
+	// Timeout for outgoing http request/response cycle.
+	RequestTimeout arvados.Duration
 }
 
 // RunOptions controls runtime behavior. The flags/options that belong
@@ -43,18 +64,24 @@ 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{}) {}
 
 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")
@@ -62,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)")
+	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("keep-balance %s\n", version)
+		return
 	}
-	mustReadJSON(&config, *configPath)
+
+	mustReadConfig(&cfg, *configPath)
 	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 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)
@@ -86,24 +132,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)
 	}
 }
 
@@ -130,7 +174,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 {