X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/8a879fd6ffb38927150be85c63d926cd6a4c0d42..e8d1a643cdbc3a5f4c0e5c745da58d9f7e1248d8:/services/keep-balance/main.go

diff --git a/services/keep-balance/main.go b/services/keep-balance/main.go
index 82e02cce81..8b4ee84c71 100644
--- a/services/keep-balance/main.go
+++ b/services/keep-balance/main.go
@@ -9,57 +9,17 @@ import (
 	"flag"
 	"fmt"
 	"io"
-	"log"
 	"os"
 
-	"git.curoverse.com/arvados.git/lib/config"
-	"git.curoverse.com/arvados.git/lib/service"
-	"git.curoverse.com/arvados.git/sdk/go/arvados"
-	"git.curoverse.com/arvados.git/sdk/go/ctxlog"
+	"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"
+	"git.arvados.org/arvados.git/sdk/go/health"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/sirupsen/logrus"
 )
 
-var (
-	version = "dev"
-	debugf  = func(string, ...interface{}) {}
-	command = service.Command(arvados.ServiceNameKeepbalance, newHandler)
-	options RunOptions
-)
-
-func newHandler(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("You must either run keep-balance with the -once flag, or set Collections.BalancePeriod in the config. "+
-			"If using the legacy keep-balance.yml config, RunPeriod is the equivalant of Collections.BalancePeriod."))
-	}
-
-	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))
-	}
-
-	if cluster.SystemLogs.LogLevel == "debug" {
-		debugf = log.Printf
-	}
-
-	if options.Logger == nil {
-		options.Logger = ctxlog.FromContext(ctx)
-	}
-
-	srv := &Server{
-		Cluster:    cluster,
-		ArvClient:  ac,
-		RunOptions: options,
-		Metrics:    newMetrics(registry),
-		Logger:     options.Logger,
-		Dumper:     options.Dumper,
-	}
-
-	srv.Start()
-	return srv
-}
-
 func main() {
 	os.Exit(runCommand(os.Args[0], os.Args[1:], os.Stdin, os.Stdout, os.Stderr))
 }
@@ -67,6 +27,7 @@ func main() {
 func runCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
 	logger := ctxlog.FromContext(context.Background())
 
+	var options RunOptions
 	flags := flag.NewFlagSet(prog, flag.ExitOnError)
 	flags.BoolVar(&options.Once, "once", false,
 		"balance once and then exit")
@@ -90,13 +51,52 @@ func runCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.W
 		options.Dumper = dumper
 	}
 
-	// Only pass along the version flag, which gets handled in RunCommand
+	// 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,
+	}
 	flags.Visit(func(f *flag.Flag) {
-		if f.Name == "version" {
+		if !dropFlag[f.Name] {
 			args = append(args, "-"+f.Name, f.Value.String())
 		}
 	})
 
-	return command.RunCommand(prog, args, stdin, stdout, stderr)
+	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)"))
+			}
+
+			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))
+			}
+
+			if options.Logger == nil {
+				options.Logger = ctxlog.FromContext(ctx)
+			}
+
+			srv := &Server{
+				Cluster:    cluster,
+				ArvClient:  ac,
+				RunOptions: options,
+				Metrics:    newMetrics(registry),
+				Logger:     options.Logger,
+				Dumper:     options.Dumper,
+			}
+			srv.Handler = &health.Handler{
+				Token:  cluster.ManagementToken,
+				Prefix: "/_health/",
+				Routes: health.Routes{"ping": srv.CheckHealth},
+			}
+
+			go srv.run()
+			return srv
+		}).RunCommand(prog, args, stdin, stdout, stderr)
 }