snapshot/revert views

Author: CoderZhi

action/protocol/protocol.go

@@ -120,16 +120,49 @@ type (
 
 	// Views stores the view for all protocols
 	Views struct {
-		vm map[string]View
+		snapshotID int
+		snapshots  map[int]map[string]View
+		vm         map[string]View
 	}
 )
 
 func NewViews() *Views {
 	return &Views{
-		vm: make(map[string]View),
+		snapshotID: 0,
+		snapshots:  make(map[int]map[string]View),
+		vm:         make(map[string]View),
 	}
 }
 
+func (views *Views) Snapshot() int {
+	views.snapshotID++
+	views.snapshots[views.snapshotID] = make(map[string]View)
+	keys := make([]string, 0, len(views.vm))
+	for key := range views.vm {
+		keys = append(keys, key)
+	}
+	for _, key := range keys {
+		views.snapshots[views.snapshotID][key] = views.vm[key]
+		views.vm[key] = views.vm[key].Clone()
+	}
+	return views.snapshotID
+}
+
+func (views *Views) Revert(id int) error {
+	if id > views.snapshotID || id < 0 {
+		return errors.Errorf("invalid snapshot id %d, max id is %d", id, views.snapshotID)
+	}
+	for k, v := range views.snapshots[id] {
+		views.vm[k] = v
+	}
+	views.snapshotID = id
+	// clean up snapshots that are not needed anymore
+	for i := id + 1; i <= views.snapshotID; i++ {
+		delete(views.snapshots, i)
+	}
+	return nil
+}
+
 func (views *Views) Clone() *Views {
 	clone := NewViews()
 	for key, view := range views.vm {

state/factory/workingset.go

@@ -7,6 +7,7 @@ package factory
 
 import (
 	"context"
+	"math"
 	"math/big"
 	"sort"
 	"time"
@@ -70,6 +71,7 @@ type (
 		workingSetStoreFactory WorkingSetStoreFactory
 		height                 uint64
 		views                  *protocol.Views
+		viewsSnapshots         map[int]int
 		store                  workingSetStore
 		finalized              bool
 		txValidator            *protocol.GenericValidator
@@ -81,6 +83,7 @@ func newWorkingSet(height uint64, views *protocol.Views, store workingSetStore,
 	ws := &workingSet{
 		height:                 height,
 		views:                  views,
+		viewsSnapshots:         make(map[int]int),
 		store:                  store,
 		workingSetStoreFactory: storeFactory,
 	}
@@ -280,14 +283,36 @@ func (ws *workingSet) finalizeTx(ctx context.Context) {
 }
 
 func (ws *workingSet) Snapshot() int {
-	return ws.store.Snapshot()
+	id := ws.store.Snapshot()
+	vid := ws.views.Snapshot()
+	ws.viewsSnapshots[id] = vid
+
+	return id
 }
 
 func (ws *workingSet) Revert(snapshot int) error {
+	vid, ok := ws.viewsSnapshots[snapshot]
+	if !ok {
+		return errors.Errorf("snapshot %d not found", snapshot)
+	}
+	if err := ws.views.Revert(vid); err != nil {
+		return errors.Wrapf(err, "failed to revert views to snapshot %d", vid)
+	}
 	return ws.store.RevertSnapshot(snapshot)
 }
 
 func (ws *workingSet) ResetSnapshots() {
+	if len(ws.viewsSnapshots) > 0 {
+		minVID := math.MaxInt
+		for _, vid := range ws.viewsSnapshots {
+			if vid < minVID {
+				minVID = vid
+			}
+		}
+		if err := ws.views.Revert(minVID); err != nil {
+			log.L().Panic("failed to revert views to minimum snapshot", zap.Error(err))
+		}
+	}
 	ws.store.ResetSnapshots()
 }