Fix wait option for physical restore

When using --wait option during the physical restore, pbm cli waits
status done and cleanup activity to be finished.

Author: boris-ilijic

cmd/pbm/restore.go

@@ -234,11 +234,7 @@ func waitRestore(
 		return errors.Wrap(err, "get storage")
 	}
 
-	tk := time.NewTicker(time.Second * 1)
-	defer tk.Stop()
-
 	var rmeta *restore.RestoreMeta
-
 	getMeta := restore.GetRestoreMeta
 	if m.Type == defs.PhysicalBackup || m.Type == defs.IncrementalBackup {
 		getMeta = func(_ context.Context, _ connect.Client, name string) (*restore.RestoreMeta, error) {
@@ -251,6 +247,10 @@ func waitRestore(
 	if m.Type != defs.LogicalBackup {
 		frameSec = 60 * 3
 	}
+
+	tk := time.NewTicker(time.Second * 1)
+	defer tk.Stop()
+
 	for range tk.C {
 		fmt.Print(".")
 		rmeta, err = getMeta(ctx, conn, m.Name)
@@ -263,6 +263,19 @@ func waitRestore(
 
 		switch rmeta.Status {
 		case status, defs.StatusDone, defs.StatusPartlyDone:
+			if m.Type == defs.PhysicalBackup || m.Type == defs.IncrementalBackup {
+				alive, err := restore.IsCleanupHbAlive(m.Name, stg, tskew)
+				if err != nil {
+					if errors.Is(err, storage.ErrNotExist) ||
+						errors.Is(err, storage.ErrEmpty) {
+						continue
+					}
+					return errors.Wrap(err, "checking cleanup hb")
+				}
+				if alive {
+					continue
+				}
+			}
 			return nil
 		case defs.StatusError:
 			return restoreFailedError{fmt.Sprintf("operation failed with: %s", rmeta.Error)}

pbm/restore/storage.go

@@ -2,10 +2,13 @@ package restore
 
 import (
 	"encoding/json"
+	"fmt"
 	"io"
+	"path"
 	"path/filepath"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/mongodb/mongo-tools/common/db"
 
@@ -276,3 +279,36 @@ func parsePhysRestoreCond(stg storage.Storage, fname, restoreName string) (*Cond
 	}
 	return &cond, nil
 }
+
+// IsCleanupHbAlive returns true if cleanup hb is active on any agent during physical restore.
+// It reads timestamp from file and takes into account potential time skew.
+// When hb excides timeout, means that cleanup is done, false is returned.
+func IsCleanupHbAlive(restoreName string, stg storage.Storage, tskew int64) (bool, error) {
+	file := path.Join(defs.PhysRestoresDir, restoreName, fmt.Sprintf("cluster.%s", syncHbCleanupSuffix))
+	_, err := stg.FileStat(file)
+	if err != nil {
+		return false, errors.Wrap(err, "get cleanup hb file stat")
+	}
+
+	f, err := stg.SourceReader(file)
+	if err != nil {
+		return false, errors.Wrap(err, "get cleanup hb file")
+	}
+
+	b, err := io.ReadAll(f)
+	if err != nil {
+		return false, errors.Wrap(err, "read cleanup hb content")
+	}
+
+	t, err := strconv.ParseInt(strings.TrimSpace(string(b)), 10, 0)
+	if err != nil {
+		return false, errors.Wrap(err, "cleanup hb decode")
+	}
+
+	ctime := time.Now().Unix() + tskew
+	if t+int64(hbCleanupFrame.Seconds())+hbCleanupTimeout < ctime {
+		return false, nil
+	}
+
+	return true, nil
+}