PBM-1452: Expand procedure for disabling balancer (#1094)

* Expand procedure for disabling balancer

* Fix line length

Author: boris-ilijic

cmd/pbm-agent/restore.go

@@ -102,8 +102,13 @@ func (a *Agent) Restore(ctx context.Context, r *ctrl.RestoreCmd, opid ctrl.OPID,
 			}
 
 			l.Debug("waiting for balancer off")
-			bs := topo.WaitForBalancerOff(ctx, a.leadConn, time.Second*30, l)
-			l.Debug("balancer status: %s", bs)
+			bs := topo.WaitForBalancerDisabled(ctx, a.leadConn, time.Second*30, l)
+			if bs.IsDisabled() {
+				l.Debug("balancer is disabled")
+			} else {
+				l.Warning("balancer is not disabled: balancer mode: %s, in balancer round: %t",
+					bs.Mode, bs.InBalancerRound)
+			}
 		}
 	}
 

pbm/backup/backup.go

@@ -293,8 +293,14 @@ func (b *Backup) Run(ctx context.Context, bcp *ctrl.BackupCmd, opid ctrl.OPID, l
 			}
 
 			l.Debug("waiting for balancer off")
-			bs := topo.WaitForBalancerOff(ctx, b.leadConn, time.Second*30, l)
-			l.Debug("balancer status: %s", bs)
+			bs := topo.WaitForBalancerDisabled(ctx, b.leadConn, time.Second*30, l)
+			if bs.IsDisabled() {
+				l.Debug("balancer is disabled")
+			} else {
+				l.Warning("balancer is not disabled: balancer mode: %s, in balancer round: %t",
+					bs.Mode, bs.InBalancerRound)
+			}
+
 		}
 	}
 

pbm/topo/cluster.go

@@ -215,6 +215,10 @@ func (b *BalancerStatus) IsOn() bool {
 	return b.Mode == BalancerModeOn
 }
 
+func (b *BalancerStatus) IsDisabled() bool {
+	return b.Mode == BalancerModeOff && !b.InBalancerRound
+}
+
 // SetBalancerStatus sets balancer status
 func SetBalancerStatus(ctx context.Context, m connect.Client, mode BalancerMode) error {
 	var cmd string

pbm/topo/topo.go

@@ -230,7 +230,16 @@ func ReplicationLag(ctx context.Context, m *mongo.Client, self string) (int, err
 	return primaryOptime - nodeOptime, nil
 }
 
-func WaitForBalancerOff(ctx context.Context, conn connect.Client, t time.Duration, l log.LogEvent) BalancerMode {
+// WaitForBalancerDisabled waits balancer to be disabled for specified time duration t.
+// Disabling balancer means that balancer is in OFF mode and that it doesn't have active balancing round.
+// If balancer is not desibled during the specified duration t, function will stop waiting
+// and it'll return current statuses.
+func WaitForBalancerDisabled(
+	ctx context.Context,
+	conn connect.Client,
+	t time.Duration,
+	l log.LogEvent,
+) *BalancerStatus {
 	dn := time.NewTimer(t)
 	defer dn.Stop()
 
@@ -249,17 +258,17 @@ Loop:
 				l.Error("get balancer status: %v", err)
 				continue
 			}
-			if bs.Mode == BalancerModeOff {
-				return BalancerModeOff
+			if bs.IsDisabled() {
+				return bs
 			}
 		case <-dn.C:
 			break Loop
 		}
 	}
 
 	if bs == nil {
-		return BalancerMode("")
+		return &BalancerStatus{}
 	}
 
-	return bs.Mode
+	return bs
 }