Use separate storage for physical restore heartbeats

Author: jcechace

pbm/restore/physical.go

@@ -92,15 +92,16 @@ type PhysRestore struct {
 	startTS int64
 	secOpts *topo.MongodOptsSec
 
-	name      string
-	opid      string
-	nodeInfo  *topo.NodeInfo
-	stg       storage.Storage
-	bcpStg    storage.Storage
-	bcp       *backup.BackupMeta
-	files     []files
-	bcpSizeRS int64 // total uncompressed size of the backup for RS (including all increments)
-	restoreTS bson.Timestamp
+	name       string
+	opid       string
+	nodeInfo   *topo.NodeInfo
+	stg        storage.Storage
+	bcpStg     storage.Storage
+	bcp        *backup.BackupMeta
+	files      []files
+	bcpSizeRS  int64 // total uncompressed size of the backup for RS (including all increments)
+	restoreTS  bson.Timestamp
+	newStorage func() (storage.Storage, error)
 
 	confOpts *config.RestoreConf
 
@@ -120,7 +121,6 @@ type PhysRestore struct {
 
 	stopHB        chan struct{}
 	stopCleanupHB chan struct{}
-	hbWG          sync.WaitGroup
 
 	log     log.LogEvent
 	logBuff *logBuff
@@ -213,9 +213,7 @@ func NewPhysical(
 	}, nil
 }
 
-func (r *PhysRestore) closeStoragesAfterHeartbeats() {
-	r.hbWG.Wait()
-
+func (r *PhysRestore) closeStorages() {
 	storage.Close(r.bcpStg, r.log)
 	r.bcpStg = nil
 
@@ -289,7 +287,7 @@ func (r *PhysRestore) close(noerr bool, progress nodeStatus) (err error) {
 		if r.stopCleanupHB != nil {
 			close(r.stopCleanupHB)
 		}
-		r.closeStoragesAfterHeartbeats()
+		r.closeStorages()
 	}()
 
 	// resolve and exec cleanup
@@ -1209,6 +1207,7 @@ func (r *PhysRestore) Snapshot(
 	defer func() {
 		if cmd.Exit && err == nil {
 			// nothing to cleanup in case of successful ext restore with exit
+			r.closeStorages()
 			return
 		}
 		if err != nil && !errors.Is(err, ErrNoDataForShard) {
@@ -2414,7 +2413,12 @@ func (r *PhysRestore) init(ctx context.Context, name string, opid ctrl.OPID, l l
 		return errors.Wrap(err, "get pbm config")
 	}
 
-	r.stg, err = util.StorageFromConfig(&cfg.Storage, r.nodeInfo.Me, l)
+	storageConf := cfg.Storage
+	r.newStorage = func() (storage.Storage, error) {
+		return util.StorageFromConfig(&storageConf, r.nodeInfo.Me, l)
+	}
+
+	r.stg, err = r.newStorage()
 	if err != nil {
 		return errors.Wrap(err, "get storage")
 	}
@@ -2468,26 +2472,31 @@ func (r *PhysRestore) init(ctx context.Context, name string, opid ctrl.OPID, l l
 
 // startHB starts heartbeats in separate go routine.
 func (r *PhysRestore) startHB(l log.LogEvent) {
-	err := r.hb()
+	stg, err := r.newStorage()
+	if err != nil {
+		l.Error("get heartbeat storage: %v", err)
+		return
+	}
+
+	err = r.hb(stg)
 	if err != nil {
 		l.Error("send init heartbeat: %v", err)
 	}
 
 	r.stopHB = make(chan struct{})
-	r.hbWG.Add(1)
 	go func() {
 		tk := time.NewTicker(time.Second * hbFrameSec)
 		defer func() {
-			r.hbWG.Done()
 			tk.Stop()
+			storage.Close(stg, l)
 			r.stopHB = nil
 			l.Debug("heartbeats stopped")
 		}()
 
 		for {
 			select {
 			case <-tk.C:
-				err := r.hb()
+				err := r.hb(stg)
 				if err != nil {
 					l.Warning("send heartbeat: %v", err)
 				}
@@ -2498,31 +2507,31 @@ func (r *PhysRestore) startHB(l log.LogEvent) {
 	}()
 }
 
-func (r *PhysRestore) hb() error {
+func (r *PhysRestore) hb(stg storage.Storage) error {
 	now := []byte(strconv.FormatInt(time.Now().Unix(), 10))
 
-	err := storage.RetryableWrite(r.stg, r.syncPathNode+"."+syncHbSuffix, now)
+	err := storage.RetryableWrite(stg, r.syncPathNode+"."+syncHbSuffix, now)
 	if err != nil {
 		return errors.Wrap(err, "write node hb")
 	}
 
-	err = storage.RetryableWrite(r.stg, r.syncPathRS+"."+syncHbSuffix, now)
+	err = storage.RetryableWrite(stg, r.syncPathRS+"."+syncHbSuffix, now)
 	if err != nil {
 		return errors.Wrap(err, "write rs hb")
 	}
 
-	err = storage.RetryableWrite(r.stg, r.syncPathCluster+"."+syncHbSuffix, now)
+	err = storage.RetryableWrite(stg, r.syncPathCluster+"."+syncHbSuffix, now)
 	if err != nil {
 		return errors.Wrap(err, "write cluster hb")
 	}
 
 	return nil
 }
 
-func (r *PhysRestore) hbCleanup() error {
+func (r *PhysRestore) hbCleanup(stg storage.Storage) error {
 	now := []byte(strconv.FormatInt(time.Now().Unix(), 10))
 
-	err := storage.RetryableWrite(r.stg, r.syncPathCluster+"."+syncHbCleanupSuffix, now)
+	err := storage.RetryableWrite(stg, r.syncPathCluster+"."+syncHbCleanupSuffix, now)
 	if err != nil {
 		return errors.Wrap(err, "write cleanup hb")
 	}
@@ -2533,26 +2542,31 @@ func (r *PhysRestore) hbCleanup() error {
 // startCleanupHb generates cleanup hb for the purpose of detecting physical restore
 // cleanup activity.
 func (r *PhysRestore) startCleanupHb() {
-	err := r.hbCleanup()
+	stg, err := r.newStorage()
+	if err != nil {
+		r.log.Warning("get cleanup heartbeat storage: %v", err)
+		return
+	}
+
+	err = r.hbCleanup(stg)
 	if err != nil {
 		r.log.Warning("send init cleanup heartbeat: %v", err)
 	}
 
 	r.stopCleanupHB = make(chan struct{})
-	r.hbWG.Add(1)
 	go func() {
 		tk := time.NewTicker(hbCleanupFrame)
 		defer func() {
-			r.hbWG.Done()
 			tk.Stop()
+			storage.Close(stg, r.log)
 			r.stopCleanupHB = nil
 			r.log.Debug("cleanup heartbeats stopped")
 		}()
 
 		for {
 			select {
 			case <-tk.C:
-				err := r.hbCleanup()
+				err := r.hbCleanup(stg)
 				if err != nil {
 					r.log.Warning("send cleanup heartbeat: %v", err)
 				}
@@ -3186,12 +3200,13 @@ func PhysRestoreFinish(l log.LogEvent, cmd *ExtFinishCmd) error {
 		return errors.Wrap(err, "creating restore object from storage dump")
 	}
 
+	defer r.closeStorages()
+
 	r.startHB(l)
 	defer func() {
 		if r.stopHB != nil {
 			close(r.stopHB)
 		}
-		r.closeStoragesAfterHeartbeats()
 	}()
 
 	excfg, err := r.prepareExtRestore(l)
@@ -3295,6 +3310,9 @@ func physRestoreFromExtDump(
 	physRestore := PhysRestore{
 		stg: stg,
 		log: l,
+		newStorage: func() (storage.Storage, error) {
+			return GetRestoreMetaStg(cmd.CfgPath, cmd.Node)
+		},
 
 		dbpath:             extDump.DBpath,
 		tmpPort:            extDump.TmpPort,