X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/4620478d694697eff07e501187d784c6c98ccfa9..166f27869b2f530aa3dd05b02cdc46f787dd8216:/services/keep-balance/main.go

diff --git a/services/keep-balance/main.go b/services/keep-balance/main.go
index 3316a17240..65bd8d4cf0 100644
--- a/services/keep-balance/main.go
+++ b/services/keep-balance/main.go
@@ -5,96 +5,94 @@
 package main
 
 import (
-	"encoding/json"
+	"context"
 	"flag"
 	"fmt"
-	"log"
+	"io"
 	"net/http"
 	"os"
-	"time"
 
-	"git.curoverse.com/arvados.git/sdk/go/arvados"
-	"git.curoverse.com/arvados.git/sdk/go/config"
+	"git.arvados.org/arvados.git/lib/config"
+	"git.arvados.org/arvados.git/lib/service"
+	"git.arvados.org/arvados.git/sdk/go/arvados"
+	"git.arvados.org/arvados.git/sdk/go/ctxlog"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/sirupsen/logrus"
 )
 
-var debugf = func(string, ...interface{}) {}
-
 func main() {
-	var cfg Config
-	var runOptions RunOptions
+	os.Exit(runCommand(os.Args[0], os.Args[1:], os.Stdin, os.Stdout, os.Stderr))
+}
+
+func runCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
+	logger := ctxlog.FromContext(context.Background())
 
-	configPath := flag.String("config", defaultConfigPath,
-		"`path` of JSON or YAML configuration file")
-	serviceListPath := flag.String("config.KeepServiceList", "",
-		"`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,
+	var options RunOptions
+	flags := flag.NewFlagSet(prog, flag.ExitOnError)
+	flags.BoolVar(&options.Once, "once", false,
 		"balance once and then exit")
-	flag.BoolVar(&runOptions.CommitPulls, "commit-pulls", false,
+	flags.BoolVar(&options.CommitPulls, "commit-pulls", false,
 		"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,
+	flags.BoolVar(&options.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()
+	flags.Bool("version", false, "Write version information to stdout and exit 0")
+	dumpFlag := flags.Bool("dump", false, "dump details for each block to stdout")
 
-	// Print version information if requested
-	if *getVersion {
-		fmt.Printf("keep-balance %s\n", version)
-		return
-	}
+	loader := config.NewLoader(os.Stdin, logger)
+	loader.SetupFlags(flags)
 
-	mustReadConfig(&cfg, *configPath)
-	if *serviceListPath != "" {
-		mustReadConfig(&cfg.KeepServiceList, *serviceListPath)
-	}
+	munged := loader.MungeLegacyConfigArgs(logger, args, "-legacy-keepbalance-config")
+	flags.Parse(munged)
 
-	if *dumpConfig {
-		log.Fatal(config.DumpAndExit(cfg))
+	if *dumpFlag {
+		dumper := logrus.New()
+		dumper.Out = os.Stdout
+		dumper.Formatter = &logrus.TextFormatter{}
+		options.Dumper = dumper
 	}
 
-	to := time.Duration(cfg.RequestTimeout)
-	if to == 0 {
-		to = 30 * time.Minute
+	// Drop our custom args that would be rejected by the generic
+	// service.Command
+	args = nil
+	dropFlag := map[string]bool{
+		"once":         true,
+		"commit-pulls": true,
+		"commit-trash": true,
+		"dump":         true,
 	}
-	arvados.DefaultSecureClient.Timeout = to
-	arvados.InsecureHTTPClient.Timeout = to
-	http.DefaultClient.Timeout = to
+	flags.Visit(func(f *flag.Flag) {
+		if !dropFlag[f.Name] {
+			args = append(args, "-"+f.Name, f.Value.String())
+		}
+	})
 
-	log.Printf("keep-balance %s started", version)
+	return service.Command(arvados.ServiceNameKeepbalance,
+		func(ctx context.Context, cluster *arvados.Cluster, token string, registry *prometheus.Registry) service.Handler {
+			if !options.Once && cluster.Collections.BalancePeriod == arvados.Duration(0) {
+				return service.ErrorHandler(ctx, cluster, fmt.Errorf("cannot start service: Collections.BalancePeriod is zero (if you want to run once and then exit, use the -once flag)"))
+			}
 
-	if *debugFlag {
-		debugf = log.Printf
-		if j, err := json.Marshal(cfg); err != nil {
-			log.Fatal(err)
-		} else {
-			log.Printf("config is %s", j)
-		}
-	}
-	if *dumpFlag {
-		runOptions.Dumper = logrus.New()
-		runOptions.Dumper.Out = os.Stdout
-		runOptions.Dumper.Formatter = &logrus.TextFormatter{}
-	}
-	srv, err := NewServer(cfg, runOptions)
-	if err != nil {
-		// (don't run)
-	} else if runOptions.Once {
-		_, err = srv.Run()
-	} else {
-		err = srv.RunForever(nil)
-	}
-	if err != nil {
-		log.Fatal(err)
-	}
-}
+			ac, err := arvados.NewClientFromConfig(cluster)
+			ac.AuthToken = token
+			if err != nil {
+				return service.ErrorHandler(ctx, cluster, fmt.Errorf("error initializing client from cluster config: %s", err))
+			}
 
-func mustReadConfig(dst interface{}, path string) {
-	if err := config.LoadFile(dst, path); err != nil {
-		log.Fatal(err)
-	}
+			if options.Logger == nil {
+				options.Logger = ctxlog.FromContext(ctx)
+			}
+
+			srv := &Server{
+				Handler:    http.NotFoundHandler(),
+				Cluster:    cluster,
+				ArvClient:  ac,
+				RunOptions: options,
+				Metrics:    newMetrics(registry),
+				Logger:     options.Logger,
+				Dumper:     options.Dumper,
+			}
+
+			go srv.run()
+			return srv
+		}).RunCommand(prog, args, stdin, stdout, stderr)
 }