X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/5c0d99ac1c8371910676b6375282a4fcfb5d213d..07baa0ed049746514495d1648c1aef0c40545141:/services/keep-balance/balance.go

diff --git a/services/keep-balance/balance.go b/services/keep-balance/balance.go
index 1dedb409a4..9f581751d9 100644
--- a/services/keep-balance/balance.go
+++ b/services/keep-balance/balance.go
@@ -23,7 +23,9 @@ import (
 	"syscall"
 	"time"
 
+	"git.arvados.org/arvados.git/lib/controller/dblock"
 	"git.arvados.org/arvados.git/sdk/go/arvados"
+	"git.arvados.org/arvados.git/sdk/go/ctxlog"
 	"git.arvados.org/arvados.git/sdk/go/keepclient"
 	"github.com/jmoiron/sqlx"
 	"github.com/sirupsen/logrus"
@@ -67,16 +69,19 @@ type Balancer struct {
 // subsequent balance operation.
 //
 // Run should only be called once on a given Balancer object.
-//
-// Typical usage:
-//
-//   runOptions, err = (&Balancer{}).Run(config, runOptions)
-func (bal *Balancer) Run(client *arvados.Client, cluster *arvados.Cluster, runOptions RunOptions) (nextRunOptions RunOptions, err error) {
+func (bal *Balancer) Run(ctx context.Context, client *arvados.Client, cluster *arvados.Cluster, runOptions RunOptions) (nextRunOptions RunOptions, err error) {
 	nextRunOptions = runOptions
 
+	ctxlog.FromContext(ctx).Info("acquiring active lock")
+	if !dblock.KeepBalanceActive.Lock(ctx, func(context.Context) (*sqlx.DB, error) { return bal.DB, nil }) {
+		// context canceled
+		return
+	}
+	defer dblock.KeepBalanceActive.Unlock()
+
 	defer bal.time("sweep", "wall clock time to run one full sweep")()
 
-	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(cluster.Collections.BalanceTimeout.Duration()))
+	ctx, cancel := context.WithDeadline(ctx, time.Now().Add(cluster.Collections.BalanceTimeout.Duration()))
 	defer cancel()
 
 	var lbFile *os.File