lastSnapshot marker fix

common/storage/storage.go

@@ -0,0 +1,67 @@
+package storage
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/multiversx/mx-chain-core-go/core"
+	"github.com/multiversx/mx-chain-go/common"
+	logger "github.com/multiversx/mx-chain-logger-go"
+)
+
+var log = logger.GetOrCreate("state")
+
+const (
+	// WaitTimeForSnapshotEpochCheck is the time to wait before checking the storage epoch
+	WaitTimeForSnapshotEpochCheck = time.Millisecond * 100
+
+	// SnapshotWaitTimeout is the timeout for waiting for the storage epoch to change
+	SnapshotWaitTimeout = time.Minute * 3
+)
+
+// StorageEpochChangeWaitArgs are the args needed for calling the WaitForStorageEpochChange function
+type StorageEpochChangeWaitArgs struct {
+	TrieStorageManager            common.StorageManager
+	Epoch                         uint32
+	WaitTimeForSnapshotEpochCheck time.Duration
+	SnapshotWaitTimeout           time.Duration
+}
+
+// WaitForStorageEpochChange waits for the storage epoch to change to the given epoch
+func WaitForStorageEpochChange(args StorageEpochChangeWaitArgs) error {
+	log.Debug("waiting for storage epoch change", "epoch", args.Epoch, "wait timeout", args.SnapshotWaitTimeout)
+
+	if args.SnapshotWaitTimeout < args.WaitTimeForSnapshotEpochCheck {
+		return fmt.Errorf("timeout (%s) must be greater than wait time between snapshot epoch check (%s)", args.SnapshotWaitTimeout, args.WaitTimeForSnapshotEpochCheck)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), args.SnapshotWaitTimeout)
+	defer cancel()
+
+	timer := time.NewTimer(args.WaitTimeForSnapshotEpochCheck)
+	defer timer.Stop()
+
+	for {
+		timer.Reset(args.WaitTimeForSnapshotEpochCheck)
+
+		if args.TrieStorageManager.IsClosed() {
+			return core.ErrContextClosing
+		}
+
+		latestStorageEpoch, err := args.TrieStorageManager.GetLatestStorageEpoch()
+		if err != nil {
+			return err
+		}
+
+		if latestStorageEpoch == args.Epoch {
+			return nil
+		}
+
+		select {
+		case <-timer.C:
+		case <-ctx.Done():
+			return fmt.Errorf("timeout waiting for storage epoch change, snapshot epoch %d", args.Epoch)
+		}
+	}
+}

common/storage/storage_test.go

@@ -0,0 +1,93 @@
+package storage
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/multiversx/mx-chain-core-go/core"
+	"github.com/multiversx/mx-chain-go/testscommon/storageManager"
+	"github.com/stretchr/testify/assert"
+)
+
+func getDefaultArgs() StorageEpochChangeWaitArgs {
+	return StorageEpochChangeWaitArgs{
+		Epoch:                         1,
+		WaitTimeForSnapshotEpochCheck: time.Millisecond * 100,
+		SnapshotWaitTimeout:           time.Second,
+		TrieStorageManager:            &storageManager.StorageManagerStub{},
+	}
+}
+
+func TestSnapshotsManager_WaitForStorageEpochChange(t *testing.T) {
+	t.Parallel()
+
+	t.Run("invalid args", func(t *testing.T) {
+		t.Parallel()
+
+		args := getDefaultArgs()
+		args.SnapshotWaitTimeout = time.Millisecond
+
+		err := WaitForStorageEpochChange(args)
+		assert.Error(t, err)
+	})
+	t.Run("getLatestStorageEpoch error", func(t *testing.T) {
+		t.Parallel()
+
+		expectedError := errors.New("getLatestStorageEpoch error")
+
+		args := getDefaultArgs()
+		args.TrieStorageManager = &storageManager.StorageManagerStub{
+			GetLatestStorageEpochCalled: func() (uint32, error) {
+				return 0, expectedError
+			},
+		}
+
+		err := WaitForStorageEpochChange(args)
+		assert.Equal(t, expectedError, err)
+	})
+	t.Run("storage manager closed error", func(t *testing.T) {
+		t.Parallel()
+
+		args := getDefaultArgs()
+		args.TrieStorageManager = &storageManager.StorageManagerStub{
+			GetLatestStorageEpochCalled: func() (uint32, error) {
+				return 0, nil
+			},
+			IsClosedCalled: func() bool {
+				return true
+			},
+		}
+
+		err := WaitForStorageEpochChange(args)
+		assert.Equal(t, core.ErrContextClosing, err)
+	})
+	t.Run("storage epoch change timeout", func(t *testing.T) {
+		t.Parallel()
+
+		args := getDefaultArgs()
+		args.WaitTimeForSnapshotEpochCheck = time.Millisecond
+		args.SnapshotWaitTimeout = time.Millisecond * 5
+		args.TrieStorageManager = &storageManager.StorageManagerStub{
+			GetLatestStorageEpochCalled: func() (uint32, error) {
+				return 0, nil
+			},
+		}
+
+		err := WaitForStorageEpochChange(args)
+		assert.Error(t, err)
+	})
+	t.Run("returns when latestStorageEpoch == snapshotEpoch", func(t *testing.T) {
+		t.Parallel()
+
+		args := getDefaultArgs()
+		args.TrieStorageManager = &storageManager.StorageManagerStub{
+			GetLatestStorageEpochCalled: func() (uint32, error) {
+				return 1, nil
+			},
+		}
+
+		err := WaitForStorageEpochChange(args)
+		assert.Nil(t, err)
+	})
+}

state/accountsDB.go

@@ -19,6 +19,7 @@ import (
 	"github.com/multiversx/mx-chain-go/common/errChan"
 	"github.com/multiversx/mx-chain-go/common/holders"
 	"github.com/multiversx/mx-chain-go/state/iteratorChannelsProvider"
+	"github.com/multiversx/mx-chain-go/state/lastSnapshotMarker"
 	"github.com/multiversx/mx-chain-go/state/parsers"
 	"github.com/multiversx/mx-chain-go/state/stateMetrics"
 	"github.com/multiversx/mx-chain-go/trie/keyBuilder"
@@ -28,11 +29,8 @@ import (
 )
 
 const (
-	leavesChannelSize             = 100
-	missingNodesChannelSize       = 100
-	lastSnapshot                  = "lastSnapshot"
-	waitTimeForSnapshotEpochCheck = time.Millisecond * 100
-	snapshotWaitTimeout           = time.Minute
+	leavesChannelSize       = 100
+	missingNodesChannelSize = 100
 )
 
 type loadingMeasurements struct {
@@ -138,6 +136,7 @@ func NewAccountsDB(args ArgsAccountsDB) (*AccountsDB, error) {
 		StateMetrics:             sm,
 		ChannelsProvider:         iteratorChannelsProvider.NewUserStateIteratorChannelsProvider(),
 		AccountFactory:           args.AccountFactory,
+		LastSnapshotMarker:       lastSnapshotMarker.NewLastSnapshotMarker(),
 	}
 	snapshotManager, err := NewSnapshotsManager(argsSnapshotsManager)
 	if err != nil {

state/accountsDB_test.go

@@ -1026,7 +1026,7 @@ func TestAccountsDB_SnapshotStateOnAClosedStorageManagerShouldNotMarkActiveDB(t
 
 	mut.RLock()
 	defer mut.RUnlock()
-	assert.True(t, lastSnapshotStartedWasPut)
+	assert.False(t, lastSnapshotStartedWasPut)
 	assert.False(t, activeDBWasPut)
 }
 
@@ -1065,11 +1065,14 @@ func TestAccountsDB_SnapshotStateWithErrorsShouldNotMarkActiveDB(t *testing.T) {
 
 					return nil
 				},
+				GetLatestStorageEpochCalled: func() (uint32, error) {
+					return 1, nil
+				},
 			}
 		},
 	}
 	adb := generateAccountDBFromTrie(trieStub)
-	adb.SnapshotState([]byte("roothash"), 0)
+	adb.SnapshotState([]byte("roothash"), 1)
 	time.Sleep(time.Second)
 
 	mut.RLock()
@@ -1106,14 +1109,14 @@ func TestAccountsDB_SnapshotStateSnapshotSameRootHash(t *testing.T) {
 
 	rootHash1 := []byte("rootHash1")
 	rootHash2 := []byte("rootHash2")
-	latestEpoch := uint32(0)
+	latestEpoch := atomic.Uint32{}
 	snapshotMutex := sync.RWMutex{}
 	takeSnapshotCalled := 0
 	trieStub := &trieMock.TrieStub{
 		GetStorageManagerCalled: func() common.StorageManager {
 			return &storageManager.StorageManagerStub{
 				GetLatestStorageEpochCalled: func() (uint32, error) {
-					return latestEpoch, nil
+					return latestEpoch.Get(), nil
 				},
 				TakeSnapshotCalled: func(_ string, _ []byte, _ []byte, iteratorChannels *common.TrieIteratorChannels, _ chan []byte, stats common.SnapshotStatisticsHandler, _ uint32) {
 					snapshotMutex.Lock()
@@ -1141,7 +1144,7 @@ func TestAccountsDB_SnapshotStateSnapshotSameRootHash(t *testing.T) {
 	snapshotMutex.Unlock()
 
 	// snapshot rootHash1 and epoch 1
-	latestEpoch = 1
+	latestEpoch.Set(1)
 	adb.SnapshotState(rootHash1, 1)
 	for adb.IsSnapshotInProgress() {
 		time.Sleep(waitForOpToFinish)
@@ -1151,7 +1154,7 @@ func TestAccountsDB_SnapshotStateSnapshotSameRootHash(t *testing.T) {
 	snapshotMutex.Unlock()
 
 	// snapshot rootHash1 and epoch 0 again
-	latestEpoch = 0
+	latestEpoch.Set(0)
 	adb.SnapshotState(rootHash1, 0)
 	for adb.IsSnapshotInProgress() {
 		time.Sleep(waitForOpToFinish)
@@ -1179,7 +1182,7 @@ func TestAccountsDB_SnapshotStateSnapshotSameRootHash(t *testing.T) {
 	snapshotMutex.Unlock()
 
 	// snapshot rootHash2 and epoch 1
-	latestEpoch = 1
+	latestEpoch.Set(1)
 	adb.SnapshotState(rootHash2, 1)
 	for adb.IsSnapshotInProgress() {
 		time.Sleep(waitForOpToFinish)
@@ -1189,7 +1192,7 @@ func TestAccountsDB_SnapshotStateSnapshotSameRootHash(t *testing.T) {
 	snapshotMutex.Unlock()
 
 	// snapshot rootHash2 and epoch 1 again
-	latestEpoch = 1
+	latestEpoch.Set(1)
 	adb.SnapshotState(rootHash2, 1)
 	for adb.IsSnapshotInProgress() {
 		time.Sleep(waitForOpToFinish)
@@ -1205,34 +1208,38 @@ func TestAccountsDB_SnapshotStateSkipSnapshotIfSnapshotInProgress(t *testing.T)
 	rootHashes := [][]byte{[]byte("rootHash1"), []byte("rootHash2"), []byte("rootHash3"), []byte("rootHash4")}
 	snapshotMutex := sync.RWMutex{}
 	takeSnapshotCalled := 0
-	numPutInEpochCalled := 0
+	numPutInEpochCalled := atomic.Uint32{}
 
 	trieStub := &trieMock.TrieStub{
 		GetStorageManagerCalled: func() common.StorageManager {
 			return &storageManager.StorageManagerStub{
 				GetLatestStorageEpochCalled: func() (uint32, error) {
-					return 0, nil
+					return uint32(mathRand.Intn(5)), nil
 				},
 				TakeSnapshotCalled: func(_ string, _ []byte, _ []byte, iteratorChannels *common.TrieIteratorChannels, _ chan []byte, stats common.SnapshotStatisticsHandler, _ uint32) {
 					snapshotMutex.Lock()
 					takeSnapshotCalled++
 					close(iteratorChannels.LeavesChan)
 					stats.SnapshotFinished()
+					for numPutInEpochCalled.Get() != 4 {
+						time.Sleep(time.Millisecond * 10)
+					}
 					snapshotMutex.Unlock()
 				},
 				PutInEpochCalled: func(key []byte, val []byte, epoch uint32) error {
 					assert.Equal(t, []byte(state.LastSnapshotStarted), key)
-					assert.Equal(t, rootHashes[epoch], val)
+					assert.Equal(t, rootHashes[epoch-1], val)
 
-					numPutInEpochCalled++
+					numPutInEpochCalled.Set(numPutInEpochCalled.Get() + 1)
 					return nil
 				},
 			}
 		},
 	}
 	adb := generateAccountDBFromTrie(trieStub)
 
-	for epoch, rootHash := range rootHashes {
+	for i, rootHash := range rootHashes {
+		epoch := i + 1
 		adb.SnapshotState(rootHash, uint32(epoch))
 	}
 	for adb.IsSnapshotInProgress() {
@@ -1242,7 +1249,7 @@ func TestAccountsDB_SnapshotStateSkipSnapshotIfSnapshotInProgress(t *testing.T)
 	snapshotMutex.Lock()
 	assert.Equal(t, 1, takeSnapshotCalled)
 	snapshotMutex.Unlock()
-	assert.Equal(t, len(rootHashes), numPutInEpochCalled)
+	assert.Equal(t, len(rootHashes), int(numPutInEpochCalled.Get()))
 }
 
 func TestAccountsDB_SnapshotStateCallsRemoveFromAllActiveEpochs(t *testing.T) {

state/errors.go

@@ -144,3 +144,6 @@ var ErrNilStateMetrics = errors.New("nil sstate metrics")
 
 // ErrNilChannelsProvider signals that a nil channels provider has been given
 var ErrNilChannelsProvider = errors.New("nil channels provider")
+
+// ErrNilLastSnapshotMarker signals that a nil last snapshot marker has been given
+var ErrNilLastSnapshotMarker = errors.New("nil last snapshot marker")

state/export_test.go

@@ -1,16 +1,13 @@
 package state
 
 import (
-	"time"
-
 	"github.com/multiversx/mx-chain-core-go/marshal"
 	"github.com/multiversx/mx-chain-go/common"
-	"github.com/multiversx/mx-chain-go/testscommon/storageManager"
 	vmcommon "github.com/multiversx/mx-chain-vm-common-go"
 )
 
 // LastSnapshotStarted -
-const LastSnapshotStarted = lastSnapshot
+const LastSnapshotStarted = "lastSnapshot"
 
 // LoadCode -
 func (adb *AccountsDB) LoadCode(accountHandler baseAccountHandler) error {
@@ -92,21 +89,6 @@ func (sm *snapshotsManager) GetLastSnapshotInfo() ([]byte, uint32) {
 	return sm.lastSnapshot.rootHash, sm.lastSnapshot.epoch
 }
 
-// GetStorageEpochChangeWaitArgs -
-func GetStorageEpochChangeWaitArgs() storageEpochChangeWaitArgs {
-	return storageEpochChangeWaitArgs{
-		Epoch:                         1,
-		WaitTimeForSnapshotEpochCheck: time.Millisecond * 100,
-		SnapshotWaitTimeout:           time.Second,
-		TrieStorageManager:            &storageManager.StorageManagerStub{},
-	}
-}
-
-// WaitForStorageEpochChange
-func (sm *snapshotsManager) WaitForStorageEpochChange(args storageEpochChangeWaitArgs) error {
-	return sm.waitForStorageEpochChange(args)
-}
-
 // NewNilSnapshotsManager -
 func NewNilSnapshotsManager() *snapshotsManager {
 	return nil

state/interface.go

@@ -183,7 +183,8 @@ type DataTrie interface {
 }
 
 // PeerAccountHandler models a peer state account, which can journalize a normal account's data
-//  with some extra features like signing statistics or rating information
+//
+//	with some extra features like signing statistics or rating information
 type PeerAccountHandler interface {
 	SetBLSPublicKey([]byte) error
 	GetRewardAddress() []byte
@@ -265,3 +266,11 @@ type SignRate interface {
 	GetNumSuccess() uint32
 	GetNumFailure() uint32
 }
+
+// LastSnapshotMarker manages the lastSnapshot marker operations
+type LastSnapshotMarker interface {
+	AddMarker(trieStorageManager common.StorageManager, epoch uint32, rootHash []byte)
+	RemoveMarker(trieStorageManager common.StorageManager, epoch uint32, rootHash []byte)
+	GetMarkerInfo(trieStorageManager common.StorageManager) ([]byte, error)
+	IsInterfaceNil() bool
+}