Improve snapshots reader.

cmd/snapshotsreader/snapshots_reader.go

@@ -1,64 +1,71 @@
 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 {
+func parseSnapshots(start, end uint64, snapshotsBody io.Reader, scheme proto.Scheme) []proto.BlockSnapshot {
+	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)
+	for height := uint64(1); 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 { // skip empty block snapshot
 			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
+
+		readSize := int(snapshotsByteSize + snapshotsSize) // []{snapshot, size} + 4 bytes = size of all snapshots
+		if cap(buf) < readSize {
+			buf = make([]byte, readSize)
+		}
+		buf = buf[:readSize]
+
+		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, 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.")
+		nBlocks       = flag.Uint64("blocks-number", 1, "Number of blocks to read since 'blocks-start'.")
 	)
 	flag.Parse()
 
@@ -72,6 +79,9 @@ func main() {
 	if *snapshotsPath == "" {
 		zap.S().Fatalf("You must specify snapshots-path option.")
 	}
+	if *blocksStart == 0 {
+		zap.S().Fatalf("You must specify 'blocks-start' option.")
+	}
 
 	ss, err := settings.BlockchainSettingsByTypeName(*blockchainType)
 	if err != nil {
@@ -82,7 +92,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)
 }