Improve snapshots reader.

Author: nickeskov

cmd/snapshotsreader/snapshots_reader.go

@@ -1,64 +1,85 @@
 package main
 
 import (
+	"bufio"
 	"encoding/binary"
 	"errors"
 	"flag"
 	"fmt"
+	"io"
 	"os"
 
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+
 	"github.com/wavesplatform/gowaves/pkg/logging"
 	"github.com/wavesplatform/gowaves/pkg/proto"
 	"github.com/wavesplatform/gowaves/pkg/settings"
-	"go.uber.org/zap"
-	"go.uber.org/zap/zapcore"
 )
 
 const (
 	snapshotsByteSize = 4
 )
 
-func parseSnapshots(nBlocks int, snapshotsBody *os.File, scheme proto.Scheme) []proto.BlockSnapshot {
+type SnapshotAtHeight struct {
+	Height        proto.Height
+	BlockSnapshot proto.BlockSnapshot
+}
+
+func parseSnapshots(start, end uint64, snapshotsBody io.Reader, scheme proto.Scheme) []SnapshotAtHeight {
+	var buf []byte
 	snapshotsSizeBytes := make([]byte, snapshotsByteSize)
-	readPos := int64(0)
-	var blocksSnapshots []proto.BlockSnapshot
-	for height := uint64(1); height <= uint64(nBlocks); height++ {
-		if _, readBerr := snapshotsBody.ReadAt(snapshotsSizeBytes, readPos); readBerr != nil {
-			zap.S().Fatalf("failed to read the snapshots size in block %v", readBerr)
+	var blocksSnapshots []SnapshotAtHeight
+	for height := uint64(2); height < end; height++ {
+		if _, readBerr := io.ReadFull(snapshotsBody, snapshotsSizeBytes); readBerr != nil {
+			zap.S().Fatalf("failed to read the snapshots size in block: %v", readBerr)
 		}
 		snapshotsSize := binary.BigEndian.Uint32(snapshotsSizeBytes)
-		if snapshotsSize == 0 {
-			readPos += snapshotsByteSize
+		if snapshotsSize == 0 { // add empty block snapshot
+			if height >= start {
+				blocksSnapshots = append(blocksSnapshots, SnapshotAtHeight{
+					Height:        height,
+					BlockSnapshot: proto.BlockSnapshot{},
+				})
+			}
 			continue
 		}
-		if snapshotsSize != 0 {
-			snapshotsInBlock := proto.BlockSnapshot{}
-			snapshots := make([]byte, snapshotsSize+snapshotsByteSize) // []{snapshot, size} + 4 bytes = size of all snapshots
-			if _, readRrr := snapshotsBody.ReadAt(snapshots, readPos); readRrr != nil {
-				zap.S().Fatalf("failed to read the snapshots in block %v", readRrr)
-			}
-			unmrshlErr := snapshotsInBlock.UnmarshalBinaryImport(snapshots, scheme)
-			if unmrshlErr != nil {
-				zap.S().Fatalf("failed to unmarshal snapshots in block %v", unmrshlErr)
-			}
-			blocksSnapshots = append(blocksSnapshots, snapshotsInBlock)
-			readPos += int64(snapshotsSize) + snapshotsByteSize
+
+		if cap(buf) < int(snapshotsSize) {
+			buf = make([]byte, snapshotsSize)
 		}
+		buf = buf[:snapshotsSize]
+
+		if _, readRrr := io.ReadFull(snapshotsBody, buf); readRrr != nil {
+			zap.S().Fatalf("failed to read the snapshots in block: %v", readRrr)
+		}
+		if height < start {
+			continue
+		}
+
+		snapshotsInBlock := proto.BlockSnapshot{}
+		unmrshlErr := snapshotsInBlock.UnmarshalBinaryImport(buf, scheme)
+		if unmrshlErr != nil {
+			zap.S().Fatalf("failed to unmarshal snapshots in block: %v", unmrshlErr)
+		}
+		blocksSnapshots = append(blocksSnapshots, SnapshotAtHeight{
+			Height:        height,
+			BlockSnapshot: snapshotsInBlock,
+		})
 	}
 	return blocksSnapshots
 }
 
 func main() {
-	const (
-		defaultBlocksNumber = 1000
-	)
 	var (
 		logLevel = zap.LevelFlag("log-level", zapcore.InfoLevel,
 			"Logging level. Supported levels: DEBUG, INFO, WARN, ERROR, FATAL. Default logging level INFO.")
 		blockchainType = flag.String("blockchain-type", "mainnet",
-			"Blockchain type. Allowed values: mainnet/testnet/stagenet/custom. Default is 'mainnet'.")
+			"Blockchain type. Allowed values: mainnet/testnet/stagenet. Default is 'mainnet'.")
 		snapshotsPath = flag.String("snapshots-path", "", "Path to binary blockchain file.")
-		nBlocks       = flag.Int("blocks-number", defaultBlocksNumber, "Number of blocks to import.")
+		blocksStart   = flag.Uint64("blocks-start", 0,
+			"Start block number. Should be greater than 1, because the snapshots file doesn't include genesis.")
+		nBlocks = flag.Uint64("blocks-number", 1, "Number of blocks to read since 'blocks-start'.")
 	)
 	flag.Parse()
 
@@ -72,6 +93,9 @@ func main() {
 	if *snapshotsPath == "" {
 		zap.S().Fatalf("You must specify snapshots-path option.")
 	}
+	if *blocksStart < 2 {
+		zap.S().Fatalf("'blocks-start' must be greater than 1.")
+	}
 
 	ss, err := settings.BlockchainSettingsByTypeName(*blockchainType)
 	if err != nil {
@@ -82,7 +106,17 @@ func main() {
 	if err != nil {
 		zap.S().Fatalf("failed to open snapshots file, %v", err)
 	}
-	blocksSnapshots := parseSnapshots(*nBlocks, snapshotsBody, ss.AddressSchemeCharacter)
+	defer func(snapshotsBody *os.File) {
+		if clErr := snapshotsBody.Close(); clErr != nil {
+			zap.S().Fatalf("failed to close snapshots file, %v", clErr)
+		}
+	}(snapshotsBody)
+	const MB = 1 << 20
+	var (
+		start = *blocksStart
+		end   = start + *nBlocks
+	)
+	blocksSnapshots := parseSnapshots(start, end, bufio.NewReaderSize(snapshotsBody, MB), ss.AddressSchemeCharacter)
 
-	zap.S().Info(blocksSnapshots[0])
+	zap.S().Info(blocksSnapshots)
 }

pkg/proto/block_snapshot.go

@@ -68,11 +68,7 @@ func (bs *BlockSnapshot) UnmarshalBinary(data []byte, scheme Scheme) error {
 }
 
 func (bs *BlockSnapshot) UnmarshalBinaryImport(data []byte, scheme Scheme) error {
-	if len(data) < uint32Size {
-		return errors.Errorf("BlockSnapshot UnmarshallBinary: invalid data size")
-	}
-	snapshotsBytesSize := binary.BigEndian.Uint32(data[0:uint32Size])
-	data = data[uint32Size:] // skip size
+	snapshotsBytesSize := len(data)
 	var txSnapshots [][]AtomicSnapshot
 	for i := uint32(0); snapshotsBytesSize > 0; i++ {
 		if len(data) < uint32Size {
@@ -94,7 +90,10 @@ func (bs *BlockSnapshot) UnmarshalBinaryImport(data []byte, scheme Scheme) error
 		}
 		txSnapshots = append(txSnapshots, atomicTS)
 		data = data[oneSnapshotSize:]
-		snapshotsBytesSize = snapshotsBytesSize - oneSnapshotSize - uint32Size
+		snapshotsBytesSize = snapshotsBytesSize - int(oneSnapshotSize) - uint32Size
+	}
+	if snapshotsBytesSize != 0 { // check that all bytes were read
+		return errors.Errorf("BlockSnapshot UnmarshallBinary: invalid snapshots size")
 	}
 	bs.TxSnapshots = txSnapshots
 	return nil