EVM-660 Use batch write inside blockchain.go (LevelDB/memory)

blockchain/blockchain_test.go

@@ -25,8 +25,7 @@ func TestGenesis(t *testing.T) {
 	genesis := &types.Header{Difficulty: 1, Number: 0}
 	genesis.ComputeHash()
 
-	_, err := b.advanceHead(genesis)
-	assert.NoError(t, err)
+	assert.NoError(t, b.writeGenesisImpl(genesis))
 
 	header := b.Header()
 	assert.Equal(t, header.Hash, genesis.Hash)
@@ -533,8 +532,12 @@ func TestForkUnknownParents(t *testing.T) {
 	h1 := AppendNewTestHeaders(h0[:5], 10)
 
 	// Write genesis
-	_, err := b.advanceHead(h0[0])
-	assert.NoError(t, err)
+	batchWriter := storage.NewBatchWriter(b.db)
+	td := new(big.Int).SetUint64(h0[0].Difficulty)
+
+	batchWriter.PutCanonicalHeader(h0[0], td)
+
+	assert.NoError(t, b.writeBatchAndUpdate(batchWriter, h0[0], td, true))
 
 	// Write 10 headers
 	assert.NoError(t, b.WriteHeaders(h0[1:]))
@@ -553,14 +556,15 @@ func TestBlockchainWriteBody(t *testing.T) {
 	newChain := func(
 		t *testing.T,
 		txFromByTxHash map[types.Hash]types.Address,
+		path string,
 	) *Blockchain {
 		t.Helper()
 
-		storage, err := memory.NewMemoryStorage(nil)
+		dbStorage, err := memory.NewMemoryStorage(nil)
 		assert.NoError(t, err)
 
 		chain := &Blockchain{
-			db: storage,
+			db: dbStorage,
 			txSigner: &mockSigner{
 				txFromByTxHash: txFromByTxHash,
 			},
@@ -590,12 +594,15 @@ func TestBlockchainWriteBody(t *testing.T) {
 
 		txFromByTxHash := map[types.Hash]types.Address{}
 
-		chain := newChain(t, txFromByTxHash)
+		chain := newChain(t, txFromByTxHash, "t1")
+		defer chain.db.Close()
+		batchWriter := storage.NewBatchWriter(chain.db)
 
 		assert.NoError(
 			t,
-			chain.writeBody(block),
+			chain.writeBody(batchWriter, block),
 		)
+		assert.NoError(t, batchWriter.WriteBatch())
 	})
 
 	t.Run("should return error if tx doesn't have from and recovering address fails", func(t *testing.T) {
@@ -618,13 +625,16 @@ func TestBlockchainWriteBody(t *testing.T) {
 
 		txFromByTxHash := map[types.Hash]types.Address{}
 
-		chain := newChain(t, txFromByTxHash)
+		chain := newChain(t, txFromByTxHash, "t2")
+		defer chain.db.Close()
+		batchWriter := storage.NewBatchWriter(chain.db)
 
 		assert.ErrorIs(
 			t,
 			errRecoveryAddressFailed,
-			chain.writeBody(block),
+			chain.writeBody(batchWriter, block),
 		)
+		assert.NoError(t, batchWriter.WriteBatch())
 	})
 
 	t.Run("should recover from address and store to storage", func(t *testing.T) {
@@ -649,9 +659,12 @@ func TestBlockchainWriteBody(t *testing.T) {
 			tx.Hash: addr,
 		}
 
-		chain := newChain(t, txFromByTxHash)
+		chain := newChain(t, txFromByTxHash, "t3")
+		defer chain.db.Close()
+		batchWriter := storage.NewBatchWriter(chain.db)
 
-		assert.NoError(t, chain.writeBody(block))
+		assert.NoError(t, chain.writeBody(batchWriter, block))
+		assert.NoError(t, batchWriter.WriteBatch())
 
 		readBody, ok := chain.readBody(block.Hash())
 		assert.True(t, ok)
@@ -854,20 +867,22 @@ func Test_recoverFromFieldsInTransactions(t *testing.T) {
 }
 
 func TestBlockchainReadBody(t *testing.T) {
-	storage, err := memory.NewMemoryStorage(nil)
+	dbStorage, err := memory.NewMemoryStorage(nil)
 	assert.NoError(t, err)
 
 	txFromByTxHash := make(map[types.Hash]types.Address)
 	addr := types.StringToAddress("1")
 
 	b := &Blockchain{
 		logger: hclog.NewNullLogger(),
-		db:     storage,
+		db:     dbStorage,
 		txSigner: &mockSigner{
 			txFromByTxHash: txFromByTxHash,
 		},
 	}
 
+	batchWriter := storage.NewBatchWriter(b.db)
+
 	tx := &types.Transaction{
 		Value: big.NewInt(10),
 		V:     big.NewInt(1),
@@ -886,10 +901,12 @@ func TestBlockchainReadBody(t *testing.T) {
 
 	txFromByTxHash[tx.Hash] = types.ZeroAddress
 
-	if err := b.writeBody(block); err != nil {
+	if err := b.writeBody(batchWriter, block); err != nil {
 		t.Fatal(err)
 	}
 
+	assert.NoError(t, batchWriter.WriteBatch())
+
 	txFromByTxHash[tx.Hash] = addr
 
 	readBody, found := b.readBody(block.Hash())

blockchain/storage/batch_writer.go

@@ -0,0 +1,96 @@
+package storage
+
+import (
+	"math/big"
+
+	"github.com/0xPolygon/polygon-edge/helper/common"
+	"github.com/0xPolygon/polygon-edge/types"
+	"github.com/umbracle/fastrlp"
+)
+
+type Batch interface {
+	Delete(key []byte)
+	Write() error
+	Put(k []byte, v []byte)
+}
+
+type BatchWriter struct {
+	batch Batch
+}
+
+func NewBatchWriter(storage Storage) *BatchWriter {
+	return &BatchWriter{batch: storage.NewBatch()}
+}
+
+func (b *BatchWriter) PutHeader(h *types.Header) {
+	b.putRlp(HEADER, h.Hash.Bytes(), h)
+}
+
+func (b *BatchWriter) PutBody(hash types.Hash, body *types.Body) {
+	b.putRlp(BODY, hash.Bytes(), body)
+}
+
+func (b *BatchWriter) PutHeadHash(h types.Hash) {
+	b.putWithPrefix(HEAD, HASH, h.Bytes())
+}
+
+func (b *BatchWriter) PutTxLookup(hash types.Hash, blockHash types.Hash) {
+	ar := &fastrlp.Arena{}
+	vr := ar.NewBytes(blockHash.Bytes()).MarshalTo(nil)
+
+	b.putWithPrefix(TX_LOOKUP_PREFIX, hash.Bytes(), vr)
+}
+
+func (b *BatchWriter) PutHeadNumber(n uint64) {
+	b.putWithPrefix(HEAD, NUMBER, common.EncodeUint64ToBytes(n))
+}
+
+func (b *BatchWriter) PutReceipts(hash types.Hash, receipts []*types.Receipt) {
+	rr := types.Receipts(receipts)
+
+	b.putRlp(RECEIPTS, hash.Bytes(), &rr)
+}
+
+func (b *BatchWriter) PutCanonicalHeader(h *types.Header, diff *big.Int) {
+	b.PutHeader(h)
+	b.PutHeadHash(h.Hash)
+	b.PutHeadNumber(h.Number)
+	b.PutCanonicalHash(h.Number, h.Hash)
+	b.PutTotalDifficulty(h.Hash, diff)
+}
+
+func (b *BatchWriter) PutCanonicalHash(n uint64, hash types.Hash) {
+	b.putWithPrefix(CANONICAL, common.EncodeUint64ToBytes(n), hash.Bytes())
+}
+
+func (b *BatchWriter) PutTotalDifficulty(hash types.Hash, diff *big.Int) {
+	b.putWithPrefix(DIFFICULTY, hash.Bytes(), diff.Bytes())
+}
+
+func (b *BatchWriter) PutForks(forks []types.Hash) {
+	ff := Forks(forks)
+
+	b.putRlp(FORK, EMPTY, &ff)
+}
+
+func (b *BatchWriter) putRlp(p, k []byte, raw types.RLPMarshaler) {
+	var data []byte
+
+	if obj, ok := raw.(types.RLPStoreMarshaler); ok {
+		data = obj.MarshalStoreRLPTo(nil)
+	} else {
+		data = raw.MarshalRLPTo(nil)
+	}
+
+	b.putWithPrefix(p, k, data)
+}
+
+func (b *BatchWriter) putWithPrefix(p, k, data []byte) {
+	fullKey := append(append(make([]byte, 0, len(p)+len(k)), p...), k...)
+
+	b.batch.Put(fullKey, data)
+}
+
+func (b *BatchWriter) WriteBatch() error {
+	return b.batch.Write()
+}