X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/da79c28c17eb5de3196ab49718b86c3f73db2380..e16866d0f398f6c61f11e2ecdf473d47100329c0:/services/keep-balance/main.go diff --git a/services/keep-balance/main.go b/services/keep-balance/main.go index 606fde4984..8a95d389c8 100644 --- a/services/keep-balance/main.go +++ b/services/keep-balance/main.go @@ -9,49 +9,22 @@ import ( "flag" "fmt" "io" - "log" + "net/http" + _ "net/http/pprof" "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/cmd" + "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/jmoiron/sqlx" + _ "github.com/lib/pq" + "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, _ string) 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) - 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 - } - - srv := &Server{ - Cluster: cluster, - ArvClient: ac, - RunOptions: options, - Metrics: newMetrics(), - } - - go srv.Start(ctx) - return srv -} - func main() { os.Exit(runCommand(os.Args[0], os.Args[1:], os.Stdin, os.Stdout, os.Stderr)) } @@ -59,21 +32,36 @@ func main() { func runCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int { logger := ctxlog.FromContext(context.Background()) - flags := flag.NewFlagSet(prog, flag.ExitOnError) + var options RunOptions + flags := flag.NewFlagSet(prog, flag.ContinueOnError) flags.BoolVar(&options.Once, "once", false, "balance once and then exit") 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)") flags.BoolVar(&options.CommitTrash, "commit-trash", false, "send trash requests (delete unreferenced old blocks, and excess replicas of overreplicated blocks)") - flags.Bool("version", false, "Write version information to stdout and exit 0") + flags.BoolVar(&options.CommitConfirmedFields, "commit-confirmed-fields", true, + "update collection fields (replicas_confirmed, storage_classes_confirmed, etc.)") dumpFlag := flags.Bool("dump", false, "dump details for each block to stdout") + pprofAddr := flags.String("pprof", "", "serve Go profile data at `[addr]:port`") + // "show version" is implemented by service.Command, so we + // don't need the var here -- we just need the -version flag + // to pass flags.Parse(). + flags.Bool("version", false, "Write version information to stdout and exit 0") + + if *pprofAddr != "" { + go func() { + logrus.Println(http.ListenAndServe(*pprofAddr, nil)) + }() + } loader := config.NewLoader(os.Stdin, logger) loader.SetupFlags(flags) munged := loader.MungeLegacyConfigArgs(logger, args, "-legacy-keepbalance-config") - flags.Parse(munged) + if ok, code := cmd.ParseFlags(flags, prog, munged, "", stderr); !ok { + return code + } if *dumpFlag { dumper := logrus.New() @@ -82,13 +70,66 @@ 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, + "commit-confirmed-fields": true, + "dump": true, + } flags.Visit(func(f *flag.Flag) { - if f.Name == "version" { - args = append(args, "-"+f.Name, f.Value.String()) + 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)) + } + + db, err := sqlx.Open("postgres", cluster.PostgreSQL.Connection.String()) + if err != nil { + return service.ErrorHandler(ctx, cluster, fmt.Errorf("postgresql connection failed: %s", err)) + } + if p := cluster.PostgreSQL.ConnectionPool; p > 0 { + db.SetMaxOpenConns(p) + } + err = db.Ping() + if err != nil { + return service.ErrorHandler(ctx, cluster, fmt.Errorf("postgresql connection succeeded but ping failed: %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, + DB: db, + } + 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) }