Merge branch '18947-keep-balance'

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

diff --git a/services/keep-balance/server.go b/services/keep-balance/server.go
index 9f192d6355cacc791e3ac2f71cf3afc66da2cadc..e485f5b2061f28134306d1d897b22cb62e4190e9 100644 (file)
--- a/services/keep-balance/server.go
+++ b/services/keep-balance/server.go
@@ -2,7 +2,7 @@
 //
 // SPDX-License-Identifier: AGPL-3.0
 
-package main
+package keepbalance
 
 import (
        "net/http"
@@ -11,10 +11,8 @@ import (
        "syscall"
        "time"
 
-       "git.curoverse.com/arvados.git/sdk/go/arvados"
-       "git.curoverse.com/arvados.git/sdk/go/auth"
-       "github.com/julienschmidt/httprouter"
-       "github.com/prometheus/client_golang/prometheus/promhttp"
+       "git.arvados.org/arvados.git/sdk/go/arvados"
+       "github.com/jmoiron/sqlx"
        "github.com/sirupsen/logrus"
 )
 
@@ -26,11 +24,12 @@ import (
 //
 // RunOptions fields are controlled by command line flags.
 type RunOptions struct {
-       Once        bool
-       CommitPulls bool
-       CommitTrash bool
-       Logger      logrus.FieldLogger
-       Dumper      logrus.FieldLogger
+       Once                  bool
+       CommitPulls           bool
+       CommitTrash           bool
+       CommitConfirmedFields bool
+       Logger                logrus.FieldLogger
+       Dumper                logrus.FieldLogger
 
        // SafeRendezvousState from the most recent balance operation,
        // or "" if unknown. If this changes from one run to the next,
@@ -40,41 +39,27 @@ type RunOptions struct {
 }
 
 type Server struct {
+       http.Handler
+
        Cluster    *arvados.Cluster
        ArvClient  *arvados.Client
        RunOptions RunOptions
        Metrics    *metrics
 
-       httpHandler http.Handler
-
        Logger logrus.FieldLogger
        Dumper logrus.FieldLogger
-}
 
-// ServeHTTP implements service.Handler.
-func (srv *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-       srv.httpHandler.ServeHTTP(w, r)
+       DB *sqlx.DB
 }
 
 // CheckHealth implements service.Handler.
 func (srv *Server) CheckHealth() error {
-       return nil
+       return srv.DB.Ping()
 }
 
-func (srv *Server) setup() {
-       if srv.Cluster.ManagementToken == "" {
-               srv.httpHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-                       http.Error(w, "Management API authentication is not configured", http.StatusForbidden)
-               })
-       } else {
-               mux := httprouter.New()
-               metricsH := promhttp.HandlerFor(srv.Metrics.reg, promhttp.HandlerOpts{
-                       ErrorLog: srv.Logger,
-               })
-               mux.Handler("GET", "/metrics", metricsH)
-               mux.Handler("GET", "/metrics.json", metricsH)
-               srv.httpHandler = auth.RequireLiteralToken(srv.Cluster.ManagementToken, mux)
-       }
+// Done implements service.Handler.
+func (srv *Server) Done() <-chan struct{} {
+       return nil
 }
 
 func (srv *Server) run() {
@@ -86,11 +71,15 @@ func (srv *Server) run() {
        }
        if err != nil {
                srv.Logger.Error(err)
+               os.Exit(1)
+       } else {
+               os.Exit(0)
        }
 }
 
 func (srv *Server) runOnce() (*Balancer, error) {
        bal := &Balancer{
+               DB:             srv.DB,
                Logger:         srv.Logger,
                Dumper:         srv.Dumper,
                Metrics:        srv.Metrics,