Accounts storage: do not store more than ROLLBACK_MAX blocks

pkg/importer/importer.go

@@ -27,8 +27,7 @@ func ApplyFromFile(st State, blockchainPath string, nBlocks, startHeight uint64,
 	sb := make([]byte, 4)
 	buf := make([]byte, 2*1024*1024)
 	r := bufio.NewReader(blockchain)
-	height := uint64(0)
-	for i := uint64(0); i < nBlocks; i++ {
+	for height := uint64(0); height < nBlocks; height++ {
 		if _, err := io.ReadFull(r, sb); err != nil {
 			return err
 		}
@@ -42,7 +41,7 @@ func ApplyFromFile(st State, blockchainPath string, nBlocks, startHeight uint64,
 				return err
 			}
 			if checkBlocks {
-				savedBlock, err := st.GetBlockByHeight(uint64(i))
+				savedBlock, err := st.GetBlockByHeight(height)
 				if err != nil {
 					return err
 				}
@@ -55,7 +54,6 @@ func ApplyFromFile(st State, blockchainPath string, nBlocks, startHeight uint64,
 				}
 			}
 		}
-		height++
 	}
 	if err := blockchain.Close(); err != nil {
 		return errors.Errorf("failed to close blockchain file: %v\n", err)

pkg/state/state.go

@@ -13,8 +13,9 @@ import (
 )
 
 const (
-	BLOCKS_STOR_DIR = "blocks_storage"
-	KEYVALUE_DIR    = "keyvalue"
+	ROLLBACK_MAX_BLOCKS = 4000
+	BLOCKS_STOR_DIR     = "blocks_storage"
+	KEYVALUE_DIR        = "keyvalue"
 )
 
 type WavesBalanceKey [1 + proto.AddressSize]byte
@@ -139,10 +140,11 @@ func NewStateManager(dataDir string, params BlockStorageParams) (*StateManager,
 	if err != nil {
 		return nil, errors.Errorf("failed to create block storage: %v\n", err)
 	}
-	accountsStor, err := storage.NewAccountsStorage(db)
+	accountsStor, err := storage.NewAccountsStorage(genesis, db)
 	if err != nil {
 		return nil, errors.Errorf("failed to create accounts storage: %v\n", err)
 	}
+	accountsStor.SetRollbackMax(ROLLBACK_MAX_BLOCKS, rw)
 	if err := sync(db, accountsStor, rw); err != nil {
 		return nil, errors.Errorf("failed to sync block storage and DB: %v\n", err)
 	}

pkg/storage/accountsstorage.go

@@ -10,12 +10,18 @@ import (
 )
 
 const (
-	ROLLBACK_MAX_BLOCKS = 2000
-	RECORD_SIZE         = crypto.SignatureSize + 8
+	RECORD_SIZE = crypto.SignatureSize + 8
 )
 
+type ID2Height interface {
+	HeightByBlockID(blockID crypto.Signature) (uint64, error)
+}
+
 type AccountsStorage struct {
-	Db keyvalue.IterableKeyVal
+	genesis     crypto.Signature
+	Db          keyvalue.IterableKeyVal
+	id2Height   ID2Height
+	rollbackMax int
 }
 
 var Empty = []byte{}
@@ -29,7 +35,7 @@ func toBlockID(bytes []byte) (crypto.Signature, error) {
 	return res, nil
 }
 
-func NewAccountsStorage(db keyvalue.IterableKeyVal) (*AccountsStorage, error) {
+func NewAccountsStorage(genesis crypto.Signature, db keyvalue.IterableKeyVal) (*AccountsStorage, error) {
 	has, err := db.Has([]byte{proto.DbHeightKeyPrefix})
 	if err != nil {
 		return nil, err
@@ -41,7 +47,12 @@ func NewAccountsStorage(db keyvalue.IterableKeyVal) (*AccountsStorage, error) {
 			return nil, err
 		}
 	}
-	return &AccountsStorage{Db: db}, nil
+	return &AccountsStorage{genesis: genesis, Db: db}, nil
+}
+
+func (s *AccountsStorage) SetRollbackMax(rollbackMax int, id2Height ID2Height) {
+	s.rollbackMax = rollbackMax
+	s.id2Height = id2Height
 }
 
 func (s *AccountsStorage) SetHeight(height uint64, directly bool) error {
@@ -67,6 +78,42 @@ func (s *AccountsStorage) GetHeight() (uint64, error) {
 	return binary.LittleEndian.Uint64(dbHeightBytes), nil
 }
 
+func (s *AccountsStorage) cutHistory(historyKey []byte, history []byte) ([]byte, error) {
+	historySize := len(history)
+	// Always leave at least 1 record.
+	last := historySize - RECORD_SIZE
+	for i := 0; i < last; i += RECORD_SIZE {
+		record := history[i : i+RECORD_SIZE]
+		idBytes := record[len(record)-crypto.SignatureSize:]
+		blockID, err := toBlockID(idBytes)
+		if err != nil {
+			return nil, err
+		}
+		if blockID != s.genesis {
+			blockHeight, err := s.id2Height.HeightByBlockID(blockID)
+			if err != nil {
+				return nil, err
+			}
+			currentHeight, err := s.GetHeight()
+			if err != nil {
+				return nil, err
+			}
+			if currentHeight-blockHeight > uint64(s.rollbackMax) {
+				history = history[i+RECORD_SIZE:]
+			} else {
+				break
+			}
+		}
+	}
+	if len(history) != historySize {
+		// Some records were removed, so we need to update the DB.
+		if err := s.Db.PutDirectly(historyKey, history); err != nil {
+			return nil, err
+		}
+	}
+	return history, nil
+}
+
 func (s *AccountsStorage) filterHistory(historyKey []byte, history []byte) ([]byte, error) {
 	historySize := len(history)
 	for i := historySize; i >= RECORD_SIZE; i -= RECORD_SIZE {
@@ -160,11 +207,18 @@ func (s *AccountsStorage) newHistory(newRecord []byte, key []byte, blockID crypt
 	if err != nil {
 		return nil, err
 	}
-	// Delete invalid records.
+	// Delete invalid (because of rollback) records.
 	history, err = s.filterHistory(key, history)
 	if err != nil {
 		return nil, err
 	}
+	if s.rollbackMax != 0 {
+		// Remove records which are too far in the past.
+		history, err = s.cutHistory(key, history)
+		if err != nil {
+			return nil, err
+		}
+	}
 	if len(history) < RECORD_SIZE {
 		// History is empty after filtering, new record is the first one.
 		return newRecord, nil

pkg/storage/accountsstorage_test.go

@@ -12,6 +12,7 @@ import (
 )
 
 const (
+	GENESIS_SIGNATURE   = "FSH8eAAzZNqnG8xgTZtz5xuLqXySsXgAjmFEC25hXMbEufiGjqWPnGCZFt6gLiVLJny16ipxRNAkkzjjhqTjBE2"
 	TOTAL_BLOCKS_NUMBER = 200
 )
 
@@ -25,7 +26,11 @@ func createAccountsStorage() (*AccountsStorage, []string, error) {
 	if err != nil {
 		return nil, res, err
 	}
-	stor, err := NewAccountsStorage(globalStor)
+	genesis, err := crypto.NewSignatureFromBase58(GENESIS_SIGNATURE)
+	if err != nil {
+		return nil, res, err
+	}
+	stor, err := NewAccountsStorage(genesis, globalStor)
 	if err != nil {
 		return nil, res, err
 	}