Set state node diff field to false for snapshots (#18)

Needed for cerc-io/ipld-eth-state-snapshot#1

Reviewed-on: #18

.gitea/workflows/test.yml

@@ -89,3 +89,47 @@ jobs:
           pip install pytest
           pip install -r requirements.txt
           pytest -v -k test_basic_db
+
+  compliance-test:
+    name: Run compliance tests
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          path: ./plugeth-statediff
+      - uses: actions/checkout@v3
+        with:
+          repository: cerc-io/eth-statediff-compliance
+          ref: v0.1.0
+          path: ./eth-statediff-compliance
+          token: ${{ secrets.CICD_REPO_TOKEN }}
+      - uses: actions/setup-go@v4
+        with:
+          go-version-file: './eth-statediff-compliance/go.mod'
+          check-latest: true
+      - name: Install jq
+        run: apt-get update && apt-get install -yq jq
+      - name: Set up Gitea access token
+        env:
+          TOKEN: ${{ secrets.CICD_REPO_TOKEN }}
+        run: |
+          git config --global url."https://$TOKEN:@git.vdb.to/".insteadOf https://git.vdb.to/
+
+      - name: Update go.mod for dumpdiff-geth
+        working-directory: ./eth-statediff-compliance/
+        run: ./scripts/update-mod.sh ../plugeth-statediff dumpdiff-geth/
+      - name: Update go.mod for dumpdiff-plugeth
+        working-directory: ./eth-statediff-compliance/
+        run: ./scripts/update-mod.sh ../plugeth-statediff dumpdiff-plugeth/
+      - name: Update go.mod for dumpdiff-plugeth-parallel
+        working-directory: ./eth-statediff-compliance/
+        run: ./scripts/update-mod.sh ../plugeth-statediff dumpdiff-plugeth-parallel/
+      - name: Build tools
+        working-directory: ./eth-statediff-compliance/
+        run: make all
+      - name: Compare output of geth and plugeth
+        working-directory: ./eth-statediff-compliance/
+        run: ./scripts/compare-diffs.sh geth plugeth
+      - name: Compare output of geth and plugeth-parallel
+        working-directory: ./eth-statediff-compliance/
+        run: ./scripts/compare-diffs.sh geth plugeth-parallel

builder.go

@@ -21,14 +21,19 @@ package statediff
 
 import (
 	"bytes"
+	"context"
 	"fmt"
+	"sync"
 	"time"
 
+	iterutils "github.com/cerc-io/eth-iterator-utils"
+	"github.com/cerc-io/eth-iterator-utils/tracker"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/cerc-io/plugeth-statediff/adapt"
 	"github.com/cerc-io/plugeth-statediff/indexer/database/metrics"
@@ -44,6 +49,8 @@ var (
 	emptyContractRoot = crypto.Keccak256Hash(emptyNode)
 	nullCodeHash      = crypto.Keccak256([]byte{})
 	zeroHash          common.Hash
+
+	defaultSubtrieWorkers uint = 1
 )
 
 // Builder interface exposes the method for building a state diff between two blocks
@@ -52,13 +59,10 @@ type Builder interface {
 	WriteStateDiff(Args, Params, sdtypes.StateNodeSink, sdtypes.IPLDSink) error
 }
 
-type StateDiffBuilder struct {
+type builder struct {
 	// state cache is safe for concurrent reads
-	stateCache adapt.StateView
+	stateCache     adapt.StateView
-}
+	subtrieWorkers uint
-
-type iterPair struct {
-	Older, Newer trie.NodeIterator
 }
 
 type accountUpdate struct {
@@ -74,19 +78,39 @@ func appender[T any](to *[]T) func(T) error {
 	}
 }
 
-// NewBuilder is used to create a statediff builder
+func syncedAppender[T any](to *[]T) func(T) error {
-func NewBuilder(stateCache adapt.StateView) Builder {
+	var mtx sync.Mutex
-	return &StateDiffBuilder{
+	return func(a T) error {
-		stateCache: stateCache,
+		mtx.Lock()
+		*to = append(*to, a)
+		mtx.Unlock()
+		return nil
 	}
 }
 
+// NewBuilder is used to create a statediff builder
+func NewBuilder(stateCache adapt.StateView) *builder {
+	return &builder{
+		stateCache:     stateCache,
+		subtrieWorkers: defaultSubtrieWorkers,
+	}
+}
+
+// SetSubtrieWorkers sets the number of disjoint subtries to divide among parallel workers.
+// Passing 0 will reset this to the default value.
+func (sdb *builder) SetSubtrieWorkers(n uint) {
+	if n == 0 {
+		n = defaultSubtrieWorkers
+	}
+	sdb.subtrieWorkers = n
+}
+
 // BuildStateDiffObject builds a statediff object from two blocks and the provided parameters
-func (sdb *StateDiffBuilder) BuildStateDiffObject(args Args, params Params) (sdtypes.StateObject, error) {
+func (sdb *builder) BuildStateDiffObject(args Args, params Params) (sdtypes.StateObject, error) {
 	defer metrics.UpdateDuration(time.Now(), metrics.IndexerMetrics.BuildStateDiffObjectTimer)
 	var stateNodes []sdtypes.StateLeafNode
 	var iplds []sdtypes.IPLD
-	err := sdb.WriteStateDiff(args, params, appender(&stateNodes), appender(&iplds))
+	err := sdb.WriteStateDiff(args, params, syncedAppender(&stateNodes), syncedAppender(&iplds))
 	if err != nil {
 		return sdtypes.StateObject{}, err
 	}
@@ -99,41 +123,96 @@ func (sdb *StateDiffBuilder) BuildStateDiffObject(args Args, params Params) (sdt
 }
 
 // WriteStateDiff writes a statediff object to output sinks
-func (sdb *StateDiffBuilder) WriteStateDiff(
+func (sdb *builder) WriteStateDiff(
 	args Args, params Params,
 	nodeSink sdtypes.StateNodeSink,
 	ipldSink sdtypes.IPLDSink,
 ) error {
 	defer metrics.UpdateDuration(time.Now(), metrics.IndexerMetrics.WriteStateDiffTimer)
 	// Load tries for old and new states
-	oldTrie, err := sdb.stateCache.OpenTrie(args.OldStateRoot)
+	triea, err := sdb.stateCache.OpenTrie(args.OldStateRoot)
 	if err != nil {
 		return fmt.Errorf("error opening old state trie: %w", err)
 	}
-	newTrie, err := sdb.stateCache.OpenTrie(args.NewStateRoot)
+	trieb, err := sdb.stateCache.OpenTrie(args.NewStateRoot)
+	if err != nil {
+		return fmt.Errorf("error opening new state trie: %w", err)
+	}
+	subitersA := iterutils.SubtrieIterators(triea.NodeIterator, uint(sdb.subtrieWorkers))
+	subitersB := iterutils.SubtrieIterators(trieb.NodeIterator, uint(sdb.subtrieWorkers))
+
+	logger := log.New("hash", args.BlockHash, "number", args.BlockNumber)
+	// errgroup will cancel if any group fails
+	g, ctx := errgroup.WithContext(context.Background())
+	for i := uint(0); i < sdb.subtrieWorkers; i++ {
+		func(subdiv uint) {
+			g.Go(func() error {
+				a, b := subitersA[subdiv], subitersB[subdiv]
+				it := utils.NewSymmetricDifferenceIterator(a, b)
+				return sdb.processAccounts(ctx,
+					it, &it.SymmDiffState,
+					params.watchedAddressesLeafPaths,
+					nodeSink, ipldSink, logger,
+				)
+			})
+		}(i)
+	}
+	return g.Wait()
+}
+
+// WriteStateDiff writes a statediff object to output sinks
+func (sdb *builder) WriteStateSnapshot(
+	stateRoot common.Hash, params Params,
+	nodeSink sdtypes.StateNodeSink,
+	ipldSink sdtypes.IPLDSink,
+	tracker tracker.IteratorTracker,
+) error {
+	defer metrics.UpdateDuration(time.Now(), metrics.IndexerMetrics.WriteStateDiffTimer)
+	// Load tries for old and new states
+	tree, err := sdb.stateCache.OpenTrie(stateRoot)
 	if err != nil {
 		return fmt.Errorf("error opening new state trie: %w", err)
 	}
 
-	iters := iterPair{
+	subiters, _, err := tracker.Restore(tree.NodeIterator)
-		Older: oldTrie.NodeIterator(nil),
-		Newer: newTrie.NodeIterator(nil),
-	}
-	logger := log.New("hash", args.BlockHash, "number", args.BlockNumber)
-
-	err = sdb.processAccounts(
-		iters.Older, iters.Newer,
-		params.watchedAddressesLeafPaths,
-		nodeSink, ipldSink, logger)
 	if err != nil {
-		return fmt.Errorf("error collecting createdAndUpdatedNodes: %w", err)
+		return fmt.Errorf("error restoring iterators: %w", err)
 	}
-	return nil
+	if len(subiters) != 0 {
+		// Completed iterators are not saved by the tracker, so restoring fewer than configured is ok,
+		// but having too many is a problem.
+		if len(subiters) > int(sdb.subtrieWorkers) {
+			return fmt.Errorf("restored too many iterators: expected %d, got %d",
+				sdb.subtrieWorkers, len(subiters))
+		}
+	} else {
+		subiters = iterutils.SubtrieIterators(tree.NodeIterator, uint(sdb.subtrieWorkers))
+		for i := range subiters {
+			subiters[i] = tracker.Tracked(subiters[i])
+		}
+	}
+	// errgroup will cancel if any group fails
+	g, ctx := errgroup.WithContext(context.Background())
+	for i := range subiters {
+		func(subdiv uint) {
+			g.Go(func() error {
+				symdiff := utils.AlwaysBState()
+				return sdb.processAccounts(ctx,
+					subiters[subdiv], &symdiff,
+					params.watchedAddressesLeafPaths,
+					nodeSink, ipldSink, log.DefaultLogger,
+				)
+			})
+		}(uint(i))
+	}
+	return g.Wait()
 }
 
-// processAccounts processes account creations and deletions, and returns a set of updated
+// processAccounts processes account creations, deletions, and updates
-// existing accounts, indexed by leaf key.
+// the NodeIterator and SymmDiffIterator instances should refer to the same object, will only be used
-func (sdb *StateDiffBuilder) processAccounts(a, b trie.NodeIterator,
+func (sdb *builder) processAccounts(
+	ctx context.Context,
+	it trie.NodeIterator, symdiff *utils.SymmDiffState,
 	watchedAddressesLeafPaths [][]byte,
 	nodeSink sdtypes.StateNodeSink, ipldSink sdtypes.IPLDSink,
 	logger log.Logger,
@@ -144,14 +223,19 @@ func (sdb *StateDiffBuilder) processAccounts(a, b trie.NodeIterator,
 
 	updates := make(accountUpdateMap)
 	// Cache the RLP of the previous node. When we hit a value node this will be the parent blob.
-	var prevBlob []byte
+	var prevBlob = it.NodeBlob()
-	it, itCount := utils.NewSymmetricDifferenceIterator(a, b)
 	for it.Next(true) {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+
 		// ignore node if it is not along paths of interest
 		if !isWatchedPathPrefix(watchedAddressesLeafPaths, it.Path()) {
 			continue
 		}
-		if it.FromA() { // Node exists in the old trie
+		if symdiff.FromA() { // Node exists in the old trie
 			if it.Leaf() {
 				var account types.StateAccount
 				if err := rlp.DecodeBytes(it.LeafBlob(), &account); err != nil {
@@ -160,7 +244,7 @@ func (sdb *StateDiffBuilder) processAccounts(a, b trie.NodeIterator,
 				leafKey := make([]byte, len(it.LeafKey()))
 				copy(leafKey, it.LeafKey())
 
-				if it.CommonPath() {
+				if symdiff.CommonPath() {
 					// If B also contains this leaf node, this is the old state of an updated account.
 					if update, ok := updates[string(leafKey)]; ok {
 						update.oldRoot = account.Root
@@ -178,14 +262,14 @@ func (sdb *StateDiffBuilder) processAccounts(a, b trie.NodeIterator,
 			}
 			continue
 		}
-		// Node exists in the new trie
+		// Node exists in the new trie (B)
 		if it.Leaf() {
 			accountW, err := sdb.decodeStateLeaf(it, prevBlob)
 			if err != nil {
 				return err
 			}
 
-			if it.CommonPath() {
+			if symdiff.CommonPath() {
 				// If A also contains this leaf node, this is the new state of an updated account.
 				if update, ok := updates[string(accountW.LeafKey)]; ok {
 					update.new = *accountW
@@ -198,41 +282,41 @@ func (sdb *StateDiffBuilder) processAccounts(a, b trie.NodeIterator,
 					return err
 				}
 			}
-		} else {
+			continue
-			// New trie nodes will be written to blockstore only.
+		}
-			// Reminder: this includes leaf nodes, since the geth iterator.Leaf() actually
+		// New inner trie nodes will be written to blockstore only.
-			// signifies a "value" node.
+		// Reminder: this includes leaf nodes, since the geth iterator.Leaf() actually
-			if it.Hash() == zeroHash {
+		// signifies a "value" node.
-				continue
+		if it.Hash() == zeroHash {
-			}
+			continue
-			nodeVal := make([]byte, len(it.NodeBlob()))
+		}
-			copy(nodeVal, it.NodeBlob())
+		nodeVal := make([]byte, len(it.NodeBlob()))
-			// if doing a selective diff, we need to ensure this is a watched path
+		copy(nodeVal, it.NodeBlob())
-			if len(watchedAddressesLeafPaths) > 0 {
+		// if doing a selective diff, we need to ensure this is a watched path
-				var elements []interface{}
+		if len(watchedAddressesLeafPaths) > 0 {
-				if err := rlp.DecodeBytes(nodeVal, &elements); err != nil {
+			var elements []interface{}
-					return err
+			if err := rlp.DecodeBytes(nodeVal, &elements); err != nil {
-				}
-				ok, err := isLeaf(elements)
-				if err != nil {
-					return err
-				}
-				if ok {
-					partialPath := utils.CompactToHex(elements[0].([]byte))
-					valueNodePath := append(it.Path(), partialPath...)
-					if !isWatchedPath(watchedAddressesLeafPaths, valueNodePath) {
-						continue
-					}
-				}
-			}
-			if err := ipldSink(sdtypes.IPLD{
-				CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, it.Hash().Bytes()).String(),
-				Content: nodeVal,
-			}); err != nil {
 				return err
 			}
-			prevBlob = nodeVal
+			ok, err := isLeaf(elements)
+			if err != nil {
+				return err
+			}
+			if ok {
+				partialPath := utils.CompactToHex(elements[0].([]byte))
+				valueNodePath := append(it.Path(), partialPath...)
+				if !isWatchedPath(watchedAddressesLeafPaths, valueNodePath) {
+					continue
+				}
+			}
 		}
+		if err := ipldSink(sdtypes.IPLD{
+			CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, it.Hash().Bytes()).String(),
+			Content: nodeVal,
+		}); err != nil {
+			return err
+		}
+		prevBlob = nodeVal
 	}
 
 	for key, update := range updates {
@@ -252,12 +336,10 @@ func (sdb *StateDiffBuilder) processAccounts(a, b trie.NodeIterator,
 			return err
 		}
 	}
-
-	metrics.IndexerMetrics.DifferenceIteratorCounter.Inc(int64(*itCount))
 	return it.Error()
 }
 
-func (sdb *StateDiffBuilder) processAccountDeletion(
+func (sdb *builder) processAccountDeletion(
 	leafKey []byte, account types.StateAccount, nodeSink sdtypes.StateNodeSink,
 ) error {
 	diff := sdtypes.StateLeafNode{
@@ -274,7 +356,7 @@ func (sdb *StateDiffBuilder) processAccountDeletion(
 	return nodeSink(diff)
 }
 
-func (sdb *StateDiffBuilder) processAccountCreation(
+func (sdb *builder) processAccountCreation(
 	accountW *sdtypes.AccountWrapper, ipldSink sdtypes.IPLDSink, nodeSink sdtypes.StateNodeSink,
 ) error {
 	diff := sdtypes.StateLeafNode{
@@ -305,7 +387,7 @@ func (sdb *StateDiffBuilder) processAccountCreation(
 // decodes account at leaf and encodes RLP data to CID
 // reminder: it.Leaf() == true when the iterator is positioned at a "value node" (which is not something
 // that actually exists in an MMPT), therefore we pass the parent node blob as the leaf RLP.
-func (sdb *StateDiffBuilder) decodeStateLeaf(it trie.NodeIterator, parentBlob []byte) (*sdtypes.AccountWrapper, error) {
+func (sdb *builder) decodeStateLeaf(it trie.NodeIterator, parentBlob []byte) (*sdtypes.AccountWrapper, error) {
 	var account types.StateAccount
 	if err := rlp.DecodeBytes(it.LeafBlob(), &account); err != nil {
 		return nil, fmt.Errorf("error decoding account at leaf key %x: %w", it.LeafKey(), err)
@@ -322,7 +404,7 @@ func (sdb *StateDiffBuilder) decodeStateLeaf(it trie.NodeIterator, parentBlob []
 
 // processStorageCreations processes the storage node records for a newly created account
 // i.e. it returns all the storage nodes at this state, since there is no previous state.
-func (sdb *StateDiffBuilder) processStorageCreations(
+func (sdb *builder) processStorageCreations(
 	sr common.Hash, storageSink sdtypes.StorageNodeSink, ipldSink sdtypes.IPLDSink,
 ) error {
 	defer metrics.UpdateDuration(time.Now(), metrics.IndexerMetrics.ProcessStorageCreationsTimer)
@@ -359,8 +441,9 @@ func (sdb *StateDiffBuilder) processStorageCreations(
 	return it.Error()
 }
 
-// processStorageUpdates builds the storage diff node objects for all nodes that exist in a different state at B than A
+// processStorageUpdates builds the storage diff node objects for all nodes that exist in a
-func (sdb *StateDiffBuilder) processStorageUpdates(
+// different state at B than A
+func (sdb *builder) processStorageUpdates(
 	oldroot common.Hash, newroot common.Hash,
 	storageSink sdtypes.StorageNodeSink,
 	ipldSink sdtypes.IPLDSink,
@@ -381,7 +464,7 @@ func (sdb *StateDiffBuilder) processStorageUpdates(
 
 	var prevBlob []byte
 	a, b := oldTrie.NodeIterator(nil), newTrie.NodeIterator(nil)
-	it, _ := utils.NewSymmetricDifferenceIterator(a, b)
+	it := utils.NewSymmetricDifferenceIterator(a, b)
 	for it.Next(true) {
 		if it.FromA() {
 			if it.Leaf() && !it.CommonPath() {
@@ -422,7 +505,7 @@ func (sdb *StateDiffBuilder) processStorageUpdates(
 }
 
 // processRemovedAccountStorage builds the "removed" diffs for all the storage nodes for a destroyed account
-func (sdb *StateDiffBuilder) processRemovedAccountStorage(
+func (sdb *builder) processRemovedAccountStorage(
 	sr common.Hash, storageSink sdtypes.StorageNodeSink,
 ) error {
 	defer metrics.UpdateDuration(time.Now(), metrics.IndexerMetrics.ProcessRemovedAccountStorageTimer)
@@ -456,7 +539,7 @@ func (sdb *StateDiffBuilder) processRemovedAccountStorage(
 // decodes slot at leaf and encodes RLP data to CID
 // reminder: it.Leaf() == true when the iterator is positioned at a "value node" (which is not something
 // that actually exists in an MMPT), therefore we pass the parent node blob as the leaf RLP.
-func (sdb *StateDiffBuilder) decodeStorageLeaf(it trie.NodeIterator, parentBlob []byte) sdtypes.StorageLeafNode {
+func (sdb *builder) decodeStorageLeaf(it trie.NodeIterator, parentBlob []byte) sdtypes.StorageLeafNode {
 	leafKey := make([]byte, len(it.LeafKey()))
 	copy(leafKey, it.LeafKey())
 	value := make([]byte, len(it.LeafBlob()))

builder_snapshot_test.go

@@ -0,0 +1,538 @@
+package statediff_test
+
+import (
+	"testing"
+
+	statediff "github.com/cerc-io/plugeth-statediff"
+	"github.com/cerc-io/plugeth-statediff/indexer/ipld"
+	"github.com/cerc-io/plugeth-statediff/test_helpers"
+	sdtypes "github.com/cerc-io/plugeth-statediff/types"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/crypto"
+)
+
+func TestBuilderSnapshot(t *testing.T) {
+	blocks, chain := test_helpers.MakeChain(3, test_helpers.Genesis, test_helpers.TestChainGen)
+	contractLeafKey = test_helpers.AddressToLeafKey(test_helpers.ContractAddr)
+	defer chain.Stop()
+	block0 = test_helpers.Genesis
+	block1 = blocks[0]
+	block2 = blocks[1]
+	block3 = blocks[2]
+	params := statediff.Params{}
+
+	tests := []test_helpers.SnapshotTestCase{
+		{
+			"testEmptyDiff",
+			common.Hash{},
+			&sdtypes.StateObject{
+				Nodes: emptyDiffs,
+			},
+		},
+		{
+			"testBlock0",
+			//10000 transferred from testBankAddress to account1Addr
+			block0.Root(),
+			&sdtypes.StateObject{
+				Nodes: []sdtypes.StateLeafNode{
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: bankAccountAtBlock0,
+							LeafKey: test_helpers.BankLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock0LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+				},
+				IPLDs: []sdtypes.IPLD{
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock0LeafNode)).String(),
+						Content: bankAccountAtBlock0LeafNode,
+					},
+				},
+			},
+		},
+		{
+			"testBlock1",
+			//10000 transferred from testBankAddress to account1Addr
+			block1.Root(),
+			&sdtypes.StateObject{
+				Nodes: []sdtypes.StateLeafNode{
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: bankAccountAtBlock1,
+							LeafKey: test_helpers.BankLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock1LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: minerAccountAtBlock1,
+							LeafKey: minerLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(minerAccountAtBlock1LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account1AtBlock1,
+							LeafKey: test_helpers.Account1LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock1LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+				},
+				IPLDs: []sdtypes.IPLD{
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(block1BranchRootNode)).String(),
+						Content: block1BranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock1LeafNode)).String(),
+						Content: bankAccountAtBlock1LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(minerAccountAtBlock1LeafNode)).String(),
+						Content: minerAccountAtBlock1LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock1LeafNode)).String(),
+						Content: account1AtBlock1LeafNode,
+					},
+				},
+			},
+		},
+		{
+			"testBlock2",
+			//1000 transferred from testBankAddress to account1Addr
+			//1000 transferred from account1Addr to account2Addr
+			block2.Root(),
+			&sdtypes.StateObject{
+				Nodes: []sdtypes.StateLeafNode{
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: bankAccountAtBlock2,
+							LeafKey: test_helpers.BankLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: minerAccountAtBlock2,
+							LeafKey: minerLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(minerAccountAtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account1AtBlock2,
+							LeafKey: test_helpers.Account1LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: contractAccountAtBlock2,
+							LeafKey: contractLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock2LeafNode)).String()},
+						StorageDiff: []sdtypes.StorageLeafNode{
+							{
+								Removed: false,
+								Value:   slot0StorageValue,
+								LeafKey: slot0StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+							},
+							{
+								Removed: false,
+								Value:   slot1StorageValue,
+								LeafKey: slot1StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+							},
+						},
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account2AtBlock2,
+							LeafKey: test_helpers.Account2LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account2AtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+				},
+				IPLDs: []sdtypes.IPLD{
+					{
+						CID:     ipld.Keccak256ToCid(ipld.RawBinary, test_helpers.CodeHash.Bytes()).String(),
+						Content: test_helpers.ByteCodeAfterDeployment,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(block2BranchRootNode)).String(),
+						Content: block2BranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock2LeafNode)).String(),
+						Content: bankAccountAtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(minerAccountAtBlock2LeafNode)).String(),
+						Content: minerAccountAtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String(),
+						Content: account1AtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock2LeafNode)).String(),
+						Content: contractAccountAtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(block2StorageBranchRootNode)).String(),
+						Content: block2StorageBranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+						Content: slot0StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+						Content: slot1StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account2AtBlock2LeafNode)).String(),
+						Content: account2AtBlock2LeafNode,
+					},
+				},
+			},
+		},
+		{
+			"testBlock3",
+			//the contract's storage is changed
+			//and the block is mined by account 2
+			block3.Root(),
+			&sdtypes.StateObject{
+				Nodes: []sdtypes.StateLeafNode{
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: minerAccountAtBlock2,
+							LeafKey: minerLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(minerAccountAtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account1AtBlock2,
+							LeafKey: test_helpers.Account1LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: bankAccountAtBlock3,
+							LeafKey: test_helpers.BankLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock3LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: contractAccountAtBlock3,
+							LeafKey: contractLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock3LeafNode)).String()},
+						StorageDiff: []sdtypes.StorageLeafNode{
+
+							{
+								Removed: false,
+								Value:   slot0StorageValue,
+								LeafKey: slot0StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+							},
+							{
+								Removed: false,
+								Value:   slot1StorageValue,
+								LeafKey: slot1StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+							},
+
+							{
+								Removed: false,
+								Value:   slot3StorageValue,
+								LeafKey: slot3StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot3StorageLeafNode)).String(),
+							},
+						},
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account2AtBlock3,
+							LeafKey: test_helpers.Account2LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account2AtBlock3LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+				},
+				IPLDs: []sdtypes.IPLD{
+					{
+						CID:     ipld.Keccak256ToCid(ipld.RawBinary, test_helpers.CodeHash.Bytes()).String(),
+						Content: test_helpers.ByteCodeAfterDeployment,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(minerAccountAtBlock2LeafNode)).String(),
+						Content: minerAccountAtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String(),
+						Content: account1AtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+						Content: slot0StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+						Content: slot1StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(block3BranchRootNode)).String(),
+						Content: block3BranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock3LeafNode)).String(),
+						Content: bankAccountAtBlock3LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock3LeafNode)).String(),
+						Content: contractAccountAtBlock3LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(block3StorageBranchRootNode)).String(),
+						Content: block3StorageBranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot3StorageLeafNode)).String(),
+						Content: slot3StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account2AtBlock3LeafNode)).String(),
+						Content: account2AtBlock3LeafNode,
+					},
+				},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		test_helpers.RunStateSnapshot(t, chain.StateCache(), test, params)
+	}
+}
+
+func TestBuilderSnapshotWithWatchedAddressList(t *testing.T) {
+	blocks, chain := test_helpers.MakeChain(3, test_helpers.Genesis, test_helpers.TestChainGen)
+	contractLeafKey = test_helpers.AddressToLeafKey(test_helpers.ContractAddr)
+	defer chain.Stop()
+	block0 = test_helpers.Genesis
+	block1 = blocks[0]
+	block2 = blocks[1]
+	block3 = blocks[2]
+	params := statediff.Params{
+		WatchedAddresses: []common.Address{test_helpers.Account1Addr, test_helpers.ContractAddr},
+	}
+	params.ComputeWatchedAddressesLeafPaths()
+
+	var tests = []test_helpers.SnapshotTestCase{
+		{
+			"testBlock0",
+			//10000 transferred from testBankAddress to account1Addr
+			block0.Root(),
+			&sdtypes.StateObject{
+				Nodes: emptyDiffs,
+			},
+		},
+		{
+			"testBlock1",
+			//10000 transferred from testBankAddress to account1Addr
+			block1.Root(),
+			&sdtypes.StateObject{
+				Nodes: []sdtypes.StateLeafNode{
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account1AtBlock1,
+							LeafKey: test_helpers.Account1LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock1LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+				},
+				IPLDs: []sdtypes.IPLD{
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(block1BranchRootNode)).String(),
+						Content: block1BranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock1LeafNode)).String(),
+						Content: account1AtBlock1LeafNode,
+					},
+				},
+			},
+		},
+		{
+			"testBlock2",
+			//1000 transferred from testBankAddress to account1Addr
+			//1000 transferred from account1Addr to account2Addr
+			block2.Root(),
+			&sdtypes.StateObject{
+				Nodes: []sdtypes.StateLeafNode{
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: contractAccountAtBlock2,
+							LeafKey: contractLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock2LeafNode)).String(),
+						},
+						StorageDiff: []sdtypes.StorageLeafNode{
+							{
+								Removed: false,
+								Value:   slot0StorageValue,
+								LeafKey: slot0StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+							},
+							{
+								Removed: false,
+								Value:   slot1StorageValue,
+								LeafKey: slot1StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+							},
+						},
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account1AtBlock2,
+							LeafKey: test_helpers.Account1LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+				},
+				IPLDs: []sdtypes.IPLD{
+					{
+						CID:     ipld.Keccak256ToCid(ipld.RawBinary, test_helpers.CodeHash.Bytes()).String(),
+						Content: test_helpers.ByteCodeAfterDeployment,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(block2BranchRootNode)).String(),
+						Content: block2BranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock2LeafNode)).String(),
+						Content: contractAccountAtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(block2StorageBranchRootNode)).String(),
+						Content: block2StorageBranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+						Content: slot0StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+						Content: slot1StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String(),
+						Content: account1AtBlock2LeafNode,
+					},
+				},
+			},
+		},
+		{
+			"testBlock3",
+			//the contract's storage is changed
+			//and the block is mined by account 2
+			block3.Root(),
+			&sdtypes.StateObject{
+				Nodes: []sdtypes.StateLeafNode{
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account1AtBlock2,
+							LeafKey: test_helpers.Account1LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: contractAccountAtBlock3,
+							LeafKey: contractLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock3LeafNode)).String()},
+						StorageDiff: []sdtypes.StorageLeafNode{
+							{
+								Removed: false,
+								Value:   slot0StorageValue,
+								LeafKey: slot0StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+							},
+							{
+								Removed: false,
+								Value:   slot1StorageValue,
+								LeafKey: slot1StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+							},
+							{
+								Removed: false,
+								Value:   slot3StorageValue,
+								LeafKey: slot3StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot3StorageLeafNode)).String(),
+							},
+						},
+					},
+				},
+				IPLDs: []sdtypes.IPLD{
+					{
+						CID:     ipld.Keccak256ToCid(ipld.RawBinary, test_helpers.CodeHash.Bytes()).String(),
+						Content: test_helpers.ByteCodeAfterDeployment,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String(),
+						Content: account1AtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+						Content: slot0StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+						Content: slot1StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(block3BranchRootNode)).String(),
+						Content: block3BranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock3LeafNode)).String(),
+						Content: contractAccountAtBlock3LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(block3StorageBranchRootNode)).String(),
+						Content: block3StorageBranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot3StorageLeafNode)).String(),
+						Content: slot3StorageLeafNode,
+					},
+				},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		test_helpers.RunStateSnapshot(t, chain.StateCache(), test, params)
+	}
+}

builder_test.go

@@ -503,7 +503,7 @@ func TestBuilder(t *testing.T) {
 	block3 = blocks[2]
 	params := statediff.Params{}
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		{
 			"testEmptyDiff",
 			statediff.Args{
@@ -795,13 +795,13 @@ func TestBuilder(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(),
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
-		tests, params, test_helpers.CheckedRoots{
+	test_helpers.CheckedRoots{
-			block0: bankAccountAtBlock0LeafNode,
+		block0: bankAccountAtBlock0LeafNode,
-			block1: block1BranchRootNode,
+		block1: block1BranchRootNode,
-			block2: block2BranchRootNode,
+		block2: block2BranchRootNode,
-			block3: block3BranchRootNode,
+		block3: block3BranchRootNode,
-		})
+	}.Check(t)
 }
 
 func TestBuilderWithWatchedAddressList(t *testing.T) {
@@ -817,7 +817,7 @@ func TestBuilderWithWatchedAddressList(t *testing.T) {
 	}
 	params.ComputeWatchedAddressesLeafPaths()
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		{
 			"testEmptyDiff",
 			statediff.Args{
@@ -1009,12 +1009,13 @@ func TestBuilderWithWatchedAddressList(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, test_helpers.CheckedRoots{
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
+	test_helpers.CheckedRoots{
 		block0: bankAccountAtBlock0LeafNode,
 		block1: block1BranchRootNode,
 		block2: block2BranchRootNode,
 		block3: block3BranchRootNode,
-	})
+	}.Check(t)
 }
 
 func TestBuilderWithRemovedAccountAndStorage(t *testing.T) {
@@ -1027,7 +1028,7 @@ func TestBuilderWithRemovedAccountAndStorage(t *testing.T) {
 	block6 = blocks[5]
 	params := statediff.Params{}
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		// blocks 0-3 are the same as in TestBuilderWithIntermediateNodes
 		{
 			"testBlock4",
@@ -1259,11 +1260,12 @@ func TestBuilderWithRemovedAccountAndStorage(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, test_helpers.CheckedRoots{
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
+	test_helpers.CheckedRoots{
 		block4: block4BranchRootNode,
 		block5: block5BranchRootNode,
 		block6: block6BranchRootNode,
-	})
+	}.Check(t)
 }
 
 func TestBuilderWithRemovedNonWatchedAccount(t *testing.T) {
@@ -1279,7 +1281,7 @@ func TestBuilderWithRemovedNonWatchedAccount(t *testing.T) {
 	}
 	params.ComputeWatchedAddressesLeafPaths()
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		{
 			"testBlock4",
 			statediff.Args{
@@ -1393,11 +1395,12 @@ func TestBuilderWithRemovedNonWatchedAccount(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, test_helpers.CheckedRoots{
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
+	test_helpers.CheckedRoots{
 		block4: block4BranchRootNode,
 		block5: block5BranchRootNode,
 		block6: block6BranchRootNode,
-	})
+	}.Check(t)
 }
 
 func TestBuilderWithRemovedWatchedAccount(t *testing.T) {
@@ -1413,7 +1416,7 @@ func TestBuilderWithRemovedWatchedAccount(t *testing.T) {
 	}
 	params.ComputeWatchedAddressesLeafPaths()
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		{
 			"testBlock4",
 			statediff.Args{
@@ -1596,11 +1599,12 @@ func TestBuilderWithRemovedWatchedAccount(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, test_helpers.CheckedRoots{
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
+	test_helpers.CheckedRoots{
 		block4: block4BranchRootNode,
 		block5: block5BranchRootNode,
 		block6: block6BranchRootNode,
-	})
+	}.Check(t)
 }
 
 var (
@@ -1696,7 +1700,7 @@ func TestBuilderWithMovedAccount(t *testing.T) {
 	block2 = blocks[1]
 	params := statediff.Params{}
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		{
 			"testBlock1",
 			statediff.Args{
@@ -1823,10 +1827,11 @@ func TestBuilderWithMovedAccount(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, test_helpers.CheckedRoots{
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
+	test_helpers.CheckedRoots{
 		block1: block01BranchRootNode,
 		block2: bankAccountAtBlock02LeafNode,
-	})
+	}.Check(t)
 }
 
 /*
@@ -2083,7 +2088,7 @@ func TestBuilderWithInternalizedLeafNode(t *testing.T) {
 	block3 = blocks[2]
 	params := statediff.Params{}
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		{
 			"testEmptyDiff",
 			statediff.Args{
@@ -2349,11 +2354,12 @@ func TestBuilderWithInternalizedLeafNode(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, test_helpers.CheckedRoots{
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
+	test_helpers.CheckedRoots{
 		block1: block1bBranchRootNode,
 		block2: block2bBranchRootNode,
 		block3: block3bBranchRootNode,
-	})
+	}.Check(t)
 }
 
 func TestBuilderWithInternalizedLeafNodeAndWatchedAddress(t *testing.T) {
@@ -2371,7 +2377,7 @@ func TestBuilderWithInternalizedLeafNodeAndWatchedAddress(t *testing.T) {
 	}
 	params.ComputeWatchedAddressesLeafPaths()
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		{
 			"testEmptyDiff",
 			statediff.Args{
@@ -2550,11 +2556,12 @@ func TestBuilderWithInternalizedLeafNodeAndWatchedAddress(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, test_helpers.CheckedRoots{
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
+	test_helpers.CheckedRoots{
 		block1: block1bBranchRootNode,
 		block2: block2bBranchRootNode,
 		block3: block3bBranchRootNode,
-	})
+	}.Check(t)
 }
 
 /*

config.go

@@ -44,18 +44,20 @@ type Config struct {
 	BackfillCheckPastBlocks uint64
 	// Size of the worker pool
 	NumWorkers uint
+	// Number of subtries to iterate in parallel
+	SubtrieWorkers uint
 	// Should the statediff service wait until geth has synced to the head of the blockchain?
 	WaitForSync bool
-	// Context used during DB initialization
+	// Context passed to all DB method calls
 	Context context.Context
 }
 
 // Params contains config parameters for the state diff builder
 type Params struct {
-	IncludeBlock              bool
+	IncludeBlock              bool // TODO: not used in write-requests
 	IncludeReceipts           bool
 	IncludeTD                 bool
-	IncludeCode               bool
+	IncludeCode               bool // TODO: not used by anything?
 	WatchedAddresses          []common.Address
 	watchedAddressesLeafPaths [][]byte
 }

go.mod

@@ -3,6 +3,7 @@ module github.com/cerc-io/plugeth-statediff
 go 1.19
 
 require (
+	github.com/cerc-io/eth-iterator-utils v0.1.1
 	github.com/cerc-io/eth-testing v0.2.1
 	github.com/ethereum/go-ethereum v1.11.6
 	github.com/georgysavva/scany v0.2.9
@@ -18,8 +19,9 @@ require (
 	github.com/openrelayxyz/plugeth-utils v1.2.0
 	github.com/pganalyze/pg_query_go/v4 v4.2.1
 	github.com/shopspring/decimal v1.2.0
-	github.com/stretchr/testify v1.8.1
+	github.com/stretchr/testify v1.8.2
-	github.com/thoas/go-funk v0.9.2
+	github.com/thoas/go-funk v0.9.3
+	golang.org/x/sync v0.1.0
 )
 
 require (
@@ -109,7 +111,6 @@ require (
 	golang.org/x/crypto v0.1.0 // indirect
 	golang.org/x/exp v0.0.0-20230206171751-46f607a40771 // indirect
 	golang.org/x/net v0.8.0 // indirect
-	golang.org/x/sync v0.1.0 // indirect
 	golang.org/x/sys v0.6.0 // indirect
 	golang.org/x/term v0.6.0 // indirect
 	golang.org/x/text v0.8.0 // indirect
@@ -123,7 +124,8 @@ require (
 )
 
 replace (
-	github.com/cerc-io/eth-testing => git.vdb.to/cerc-io/eth-testing v0.2.1
+	github.com/cerc-io/eth-iterator-utils => git.vdb.to/cerc-io/eth-iterator-utils v0.1.2
+	github.com/cerc-io/eth-testing => git.vdb.to/cerc-io/eth-testing v0.3.1
 	github.com/ethereum/go-ethereum => git.vdb.to/cerc-io/plugeth v0.0.0-20230808125822-691dc334fab1
 	github.com/openrelayxyz/plugeth-utils => git.vdb.to/cerc-io/plugeth-utils v0.0.0-20230706160122-cd41de354c46
 )

go.sum

@@ -1,6 +1,8 @@
 cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
-git.vdb.to/cerc-io/eth-testing v0.2.1 h1:IZAX7DVgzPkSmu1xdKZ5aOemdEYbvtgae7GUl/TUNtQ=
+git.vdb.to/cerc-io/eth-iterator-utils v0.1.2 h1:PdMR5B9wrQSYuYpFhN+9Kc8AEZ0pTt5eKCmu8oCtFcY=
-git.vdb.to/cerc-io/eth-testing v0.2.1/go.mod h1:qdvpc/W1xvf2MKx3rMOqvFvYaYIHG77Z1g0lwsmw0Uk=
+git.vdb.to/cerc-io/eth-iterator-utils v0.1.2/go.mod h1:OvXbdWbZ5viBXC/Ui1EkhsSmGB+AUX+TjGa3UDAfjfg=
+git.vdb.to/cerc-io/eth-testing v0.3.1 h1:sPnlMev6oEgTjsW7GtUkSsjKNG/+X6P9q0izSejLGpM=
+git.vdb.to/cerc-io/eth-testing v0.3.1/go.mod h1:qdvpc/W1xvf2MKx3rMOqvFvYaYIHG77Z1g0lwsmw0Uk=
 git.vdb.to/cerc-io/plugeth v0.0.0-20230808125822-691dc334fab1 h1:KLjxHwp9Zp7xhECccmJS00RiL+VwTuUGLU7qeIctg8g=
 git.vdb.to/cerc-io/plugeth v0.0.0-20230808125822-691dc334fab1/go.mod h1:cYXZu70+6xmDgIgrTD81GPasv16piiAFJnKyAbwVPMU=
 git.vdb.to/cerc-io/plugeth-utils v0.0.0-20230706160122-cd41de354c46 h1:KYcbbne/RXd7AuxbUd/3hgk1jPN+33k2CKiNsUsMCC0=
@@ -493,12 +495,12 @@ github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
-github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
+github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8=
-github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/syndtr/goleveldb v1.0.1-0.20210819022825-2ae1ddf74ef7 h1:epCh84lMvA70Z7CTTCmYQn2CKbY8j86K7/FAIr141uY=
 github.com/syndtr/goleveldb v1.0.1-0.20210819022825-2ae1ddf74ef7/go.mod h1:q4W45IWZaF22tdD+VEXcAWRA037jwmWEB5VWYORlTpc=
-github.com/thoas/go-funk v0.9.2 h1:oKlNYv0AY5nyf9g+/GhMgS/UO2ces0QRdPKwkhY3VCk=
+github.com/thoas/go-funk v0.9.3 h1:7+nAEx3kn5ZJcnDm2Bh23N2yOtweO14bi//dvRtgLpw=
-github.com/thoas/go-funk v0.9.2/go.mod h1:+IWnUfUmFO1+WVYQWQtIJHeRRdaIyyYglZN7xzUPe4Q=
+github.com/thoas/go-funk v0.9.3/go.mod h1:+IWnUfUmFO1+WVYQWQtIJHeRRdaIyyYglZN7xzUPe4Q=
 github.com/tklauser/go-sysconf v0.3.5 h1:uu3Xl4nkLzQfXNsWn15rPc/HQCJKObbt1dKJeWp3vU4=
 github.com/tklauser/go-sysconf v0.3.5/go.mod h1:MkWzOF4RMCshBAMXuhXJs64Rte09mITnppBXY/rYEFI=
 github.com/tklauser/numcpus v0.2.2 h1:oyhllyrScuYI6g+h/zUvNXNp1wy7x8qQy3t/piefldA=

indexer/constructor.go

@@ -33,7 +33,24 @@ import (
 )
 
 // NewStateDiffIndexer creates and returns an implementation of the StateDiffIndexer interface.
-func NewStateDiffIndexer(ctx context.Context, chainConfig *params.ChainConfig, nodeInfo node.Info, config interfaces.Config) (sql.Database, interfaces.StateDiffIndexer, error) {
+// The returned indexer is used to process state diffs and write them to a database.
+// The returned SQL database, if non-nil, is the indexer's backing database.
+// `ctx` is used to cancel an underlying DB connection.
+// `chainConfig` is used when processing chain state.
+// `nodeInfo` contains metadata on the Ethereum node, which is inserted with the indexed state.
+// `config` contains configuration specific to the indexer.
+// `isDiff` means to mark state nodes as belonging to an incremental diff, as opposed to a full snapshot.
+func NewStateDiffIndexer(
+	ctx context.Context,
+	chainConfig *params.ChainConfig,
+	nodeInfo node.Info,
+	config interfaces.Config,
+	isDiff bool,
+) (
+	sql.Database,
+	interfaces.StateDiffIndexer,
+	error,
+) {
 	switch config.Type() {
 	case shared.FILE:
 		log.Info("Starting statediff service in SQL file writing mode")
@@ -41,8 +58,7 @@ func NewStateDiffIndexer(ctx context.Context, chainConfig *params.ChainConfig, n
 		if !ok {
 			return nil, nil, fmt.Errorf("file config is not the correct type: got %T, expected %T", config, file.Config{})
 		}
-		fc.NodeInfo = nodeInfo
+		ind, err := file.NewStateDiffIndexer(chainConfig, fc, nodeInfo, isDiff)
-		ind, err := file.NewStateDiffIndexer(chainConfig, fc)
 		return nil, ind, err
 	case shared.POSTGRES:
 		log.Info("Starting statediff service in Postgres writing mode")
@@ -67,7 +83,7 @@ func NewStateDiffIndexer(ctx context.Context, chainConfig *params.ChainConfig, n
 			return nil, nil, fmt.Errorf("unrecognized Postgres driver type: %s", pgc.Driver)
 		}
 		db := postgres.NewPostgresDB(driver, pgc.Upsert)
-		ind, err := sql.NewStateDiffIndexer(ctx, chainConfig, db)
+		ind, err := sql.NewStateDiffIndexer(ctx, chainConfig, db, isDiff)
 		return db, ind, err
 	case shared.DUMP:
 		log.Info("Starting statediff service in data dump mode")

indexer/database/dump/batch_tx.go

@@ -19,26 +19,55 @@ package dump
 import (
 	"fmt"
 	"io"
+	"math/big"
 
 	"github.com/cerc-io/plugeth-statediff/indexer/ipld"
-
 	"github.com/cerc-io/plugeth-statediff/indexer/models"
+	"github.com/cerc-io/plugeth-statediff/utils/log"
 )
 
 // BatchTx wraps a void with the state necessary for building the tx concurrently during trie difference iteration
 type BatchTx struct {
-	BlockNumber string
+	blockNum  string
-	dump        io.Writer
+	dump      io.Writer
-	quit        chan struct{}
+	quit      chan struct{}
-	iplds       chan models.IPLDModel
+	iplds     chan models.IPLDModel
-	ipldCache   models.IPLDBatch
+	ipldCache models.IPLDBatch
-
-	submit func(blockTx *BatchTx, err error) error
 }
 
-// Submit satisfies indexer.AtomicTx
+func NewBatch(number *big.Int, dest io.Writer) *BatchTx {
-func (tx *BatchTx) Submit(err error) error {
+	batch := &BatchTx{
-	return tx.submit(tx, err)
+		blockNum:  number.String(),
+		dump:      dest,
+		iplds:     make(chan models.IPLDModel),
+		quit:      make(chan struct{}),
+		ipldCache: models.IPLDBatch{},
+	}
+	go batch.cache()
+	return batch
+}
+
+func (self *BatchTx) Submit() error {
+	close(self.quit)
+	close(self.iplds)
+
+	if err := self.flush(); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (tx *BatchTx) BlockNumber() string {
+	return tx.blockNum
+}
+
+func (tx *BatchTx) RollbackOnFailure(err error) {
+	if p := recover(); p != nil {
+		log.Info("panic detected before tx submission, but rollback not supported", "panic", p)
+		panic(p)
+	} else if err != nil {
+		log.Info("error detected before tx submission, but rollback not supported", "error", err)
+	}
 }
 
 func (tx *BatchTx) flush() error {
@@ -65,7 +94,7 @@ func (tx *BatchTx) cache() {
 
 func (tx *BatchTx) cacheDirect(key string, value []byte) {
 	tx.iplds <- models.IPLDModel{
-		BlockNumber: tx.BlockNumber,
+		BlockNumber: tx.BlockNumber(),
 		Key:         key,
 		Data:        value,
 	}
@@ -73,7 +102,7 @@ func (tx *BatchTx) cacheDirect(key string, value []byte) {
 
 func (tx *BatchTx) cacheIPLD(i ipld.IPLD) {
 	tx.iplds <- models.IPLDModel{
-		BlockNumber: tx.BlockNumber,
+		BlockNumber: tx.BlockNumber(),
 		Key:         i.Cid().String(),
 		Data:        i.RawData(),
 	}

indexer/database/dump/indexer.go

@@ -17,6 +17,7 @@
 package dump
 
 import (
+	"context"
 	"encoding/hex"
 	"fmt"
 	"io"
@@ -37,7 +38,6 @@ import (
 	"github.com/cerc-io/plugeth-statediff/indexer/models"
 	"github.com/cerc-io/plugeth-statediff/indexer/shared"
 	sdtypes "github.com/cerc-io/plugeth-statediff/types"
-	"github.com/cerc-io/plugeth-statediff/utils/log"
 )
 
 var _ interfaces.StateDiffIndexer = &StateDiffIndexer{}
@@ -62,7 +62,7 @@ func (sdi *StateDiffIndexer) ReportDBMetrics(time.Duration, <-chan bool) {}
 // PushBlock pushes and indexes block data in sql, except state & storage nodes (includes header, uncles, transactions & receipts)
 // Returns an initiated DB transaction which must be Closed via defer to commit or rollback
 func (sdi *StateDiffIndexer) PushBlock(block *types.Block, receipts types.Receipts, totalDifficulty *big.Int) (interfaces.Batch, error) {
-	start, t := time.Now(), time.Now()
+	t := time.Now()
 	blockHash := block.Hash()
 	blockHashStr := blockHash.String()
 	height := block.NumberU64()
@@ -74,7 +74,7 @@ func (sdi *StateDiffIndexer) PushBlock(block *types.Block, receipts types.Receip
 	}
 
 	// Generate the block iplds
-	headerNode, txNodes, rctNodes, logNodes, err := ipld.FromBlockAndReceipts(block, receipts)
+	txNodes, rctNodes, logNodes, err := ipld.FromBlockAndReceipts(block, receipts)
 	if err != nil {
 		return nil, fmt.Errorf("error creating IPLD nodes from block and receipts: %v", err)
 	}
@@ -91,49 +91,17 @@ func (sdi *StateDiffIndexer) PushBlock(block *types.Block, receipts types.Receip
 	} else {
 		reward = shared.CalcEthBlockReward(block.Header(), block.Uncles(), block.Transactions(), receipts)
 	}
-	t = time.Now()
 
-	blockTx := &BatchTx{
+	blockTx := NewBatch(block.Number(), sdi.dump)
-		BlockNumber: block.Number().String(),
-		dump:        sdi.dump,
-		iplds:       make(chan models.IPLDModel),
-		quit:        make(chan struct{}),
-		ipldCache:   models.IPLDBatch{},
-		submit: func(self *BatchTx, err error) error {
-			close(self.quit)
-			close(self.iplds)
-			tDiff := time.Since(t)
-			metrics.IndexerMetrics.StateStoreCodeProcessingTimer.Update(tDiff)
-			traceMsg += fmt.Sprintf("state, storage, and code storage processing time: %s\r\n", tDiff.String())
-			t = time.Now()
-			if err := self.flush(); err != nil {
-				traceMsg += fmt.Sprintf(" TOTAL PROCESSING DURATION: %s\r\n", time.Since(start).String())
-				log.Debug(traceMsg)
-				return err
-			}
-			tDiff = time.Since(t)
-			metrics.IndexerMetrics.PostgresCommitTimer.Update(tDiff)
-			traceMsg += fmt.Sprintf("postgres transaction commit duration: %s\r\n", tDiff.String())
-			traceMsg += fmt.Sprintf(" TOTAL PROCESSING DURATION: %s\r\n", time.Since(start).String())
-			log.Debug(traceMsg)
-			return err
-		},
-	}
-	go blockTx.cache()
-
-	tDiff := time.Since(t)
-	metrics.IndexerMetrics.FreePostgresTimer.Update(tDiff)
-
-	traceMsg += fmt.Sprintf("time spent waiting for free postgres tx: %s:\r\n", tDiff.String())
 	t = time.Now()
 
 	// Publish and index header, collect headerID
 	var headerID string
-	headerID, err = sdi.processHeader(blockTx, block.Header(), headerNode, reward, totalDifficulty)
+	headerID, err = sdi.PushHeader(blockTx, block.Header(), reward, totalDifficulty)
 	if err != nil {
 		return nil, err
 	}
-	tDiff = time.Since(t)
+	tDiff := time.Since(t)
 	metrics.IndexerMetrics.HeaderProcessingTimer.Update(tDiff)
 	traceMsg += fmt.Sprintf("header processing time: %s\r\n", tDiff.String())
 	t = time.Now()
@@ -167,9 +135,17 @@ func (sdi *StateDiffIndexer) PushBlock(block *types.Block, receipts types.Receip
 	return blockTx, err
 }
 
-// processHeader publishes and indexes a header IPLD in Postgres
+// PushHeader publishes and indexes a header IPLD in Postgres
 // it returns the headerID
-func (sdi *StateDiffIndexer) processHeader(tx *BatchTx, header *types.Header, headerNode ipld.IPLD, reward, td *big.Int) (string, error) {
+func (sdi *StateDiffIndexer) PushHeader(batch interfaces.Batch, header *types.Header, reward, td *big.Int) (string, error) {
+	tx, ok := batch.(*BatchTx)
+	if !ok {
+		return "", fmt.Errorf("sql: batch is expected to be of type %T, got %T", &BatchTx{}, batch)
+	}
+	headerNode, err := ipld.NewEthHeader(header)
+	if err != nil {
+		return "", err
+	}
 	tx.cacheIPLD(headerNode)
 
 	headerID := header.Hash().String()
@@ -189,7 +165,7 @@ func (sdi *StateDiffIndexer) processHeader(tx *BatchTx, header *types.Header, he
 		Coinbase:        header.Coinbase.String(),
 		Canonical:       true,
 	}
-	_, err := fmt.Fprintf(sdi.dump, "%+v\r\n", mod)
+	_, err = fmt.Fprintf(sdi.dump, "%+v\r\n", mod)
 	return headerID, err
 }
 
@@ -344,7 +320,7 @@ func (sdi *StateDiffIndexer) PushStateNode(batch interfaces.Batch, stateNode sdt
 		// short circuit if it is a Removed node
 		// this assumes the db has been initialized and a ipld.blocks entry for the Removed node is present
 		stateModel = models.StateNodeModel{
-			BlockNumber: tx.BlockNumber,
+			BlockNumber: tx.BlockNumber(),
 			HeaderID:    headerID,
 			StateKey:    common.BytesToHash(stateNode.AccountWrapper.LeafKey).String(),
 			CID:         shared.RemovedNodeStateCID,
@@ -352,7 +328,7 @@ func (sdi *StateDiffIndexer) PushStateNode(batch interfaces.Batch, stateNode sdt
 		}
 	} else {
 		stateModel = models.StateNodeModel{
-			BlockNumber: tx.BlockNumber,
+			BlockNumber: tx.BlockNumber(),
 			HeaderID:    headerID,
 			StateKey:    common.BytesToHash(stateNode.AccountWrapper.LeafKey).String(),
 			CID:         stateNode.AccountWrapper.CID,
@@ -375,7 +351,7 @@ func (sdi *StateDiffIndexer) PushStateNode(batch interfaces.Batch, stateNode sdt
 			// short circuit if it is a Removed node
 			// this assumes the db has been initialized and a ipld.blocks entry for the Removed node is present
 			storageModel := models.StorageNodeModel{
-				BlockNumber: tx.BlockNumber,
+				BlockNumber: tx.BlockNumber(),
 				HeaderID:    headerID,
 				StateKey:    common.BytesToHash(stateNode.AccountWrapper.LeafKey).String(),
 				StorageKey:  common.BytesToHash(storageNode.LeafKey).String(),
@@ -388,7 +364,7 @@ func (sdi *StateDiffIndexer) PushStateNode(batch interfaces.Batch, stateNode sdt
 			continue
 		}
 		storageModel := models.StorageNodeModel{
-			BlockNumber: tx.BlockNumber,
+			BlockNumber: tx.BlockNumber(),
 			HeaderID:    headerID,
 			StateKey:    common.BytesToHash(stateNode.AccountWrapper.LeafKey).String(),
 			StorageKey:  common.BytesToHash(storageNode.LeafKey).String(),
@@ -432,6 +408,10 @@ func (sdi *StateDiffIndexer) DetectGaps(beginBlockNumber uint64, endBlockNumber
 	return nil, nil
 }
 
+func (sdi *StateDiffIndexer) BeginTx(number *big.Int, _ context.Context) interfaces.Batch {
+	return NewBatch(number, sdi.dump)
+}
+
 // Close satisfies io.Closer
 func (sdi *StateDiffIndexer) Close() error {
 	return sdi.dump.Close()

indexer/database/file/batch_tx.go

@@ -16,14 +16,29 @@
 
 package file
 
+import "github.com/cerc-io/plugeth-statediff/utils/log"
+
 // BatchTx wraps a void with the state necessary for building the tx concurrently during trie difference iteration
 type BatchTx struct {
-	BlockNumber string
+	blockNum   string
-
+	fileWriter FileWriter
-	submit func(blockTx *BatchTx, err error) error
 }
 
 // Submit satisfies indexer.AtomicTx
-func (tx *BatchTx) Submit(err error) error {
+func (tx *BatchTx) Submit() error {
-	return tx.submit(tx, err)
+	tx.fileWriter.Flush()
+	return nil
+}
+
+func (tx *BatchTx) BlockNumber() string {
+	return tx.blockNum
+}
+
+func (tx *BatchTx) RollbackOnFailure(err error) {
+	if p := recover(); p != nil {
+		log.Info("panic detected before tx submission, but rollback not supported", "panic", p)
+		panic(p)
+	} else if err != nil {
+		log.Info("error detected before tx submission, but rollback not supported", "error", err)
+	}
 }

indexer/database/file/config.go

@@ -20,7 +20,6 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/cerc-io/plugeth-statediff/indexer/node"
 	"github.com/cerc-io/plugeth-statediff/indexer/shared"
 )
 
@@ -30,7 +29,6 @@ type Config struct {
 	OutputDir                string
 	FilePath                 string
 	WatchedAddressesFilePath string
-	NodeInfo                 node.Info
 }
 
 // FileMode to explicitly type the mode of file writer we are using
@@ -70,20 +68,11 @@ func (c Config) Type() shared.DBType {
 	return shared.FILE
 }
 
-var nodeInfo = node.Info{
-	GenesisBlock: "0xd4e56740f876aef8c010b86a40d5f56745a118d0906a34e69aec8c0db1cb8fa3",
-	NetworkID:    "1",
-	ChainID:      1,
-	ID:           "mockNodeID",
-	ClientName:   "go-ethereum",
-}
-
 // CSVTestConfig config for unit tests
 var CSVTestConfig = Config{
 	Mode:                     CSV,
 	OutputDir:                "./statediffing_test",
 	WatchedAddressesFilePath: "./statediffing_watched_addresses_test_file.csv",
-	NodeInfo:                 nodeInfo,
 }
 
 // SQLTestConfig config for unit tests
@@ -91,5 +80,4 @@ var SQLTestConfig = Config{
 	Mode:                     SQL,
 	FilePath:                 "./statediffing_test_file.sql",
 	WatchedAddressesFilePath: "./statediffing_watched_addresses_test_file.sql",
-	NodeInfo:                 nodeInfo,
 }

indexer/database/file/csv_indexer_legacy_test.go

@@ -43,7 +43,7 @@ func setupLegacyCSVIndexer(t *testing.T) {
 		require.NoError(t, err)
 	}
 
-	ind, err = file.NewStateDiffIndexer(test.LegacyConfig, file.CSVTestConfig)
+	ind, err = file.NewStateDiffIndexer(test.LegacyConfig, file.CSVTestConfig, test.LegacyNodeInfo, true)
 	require.NoError(t, err)
 
 	db, err = postgres.SetupSQLXDB()

indexer/database/file/csv_indexer_test.go

@@ -41,7 +41,7 @@ func setupCSVIndexer(t *testing.T) {
 		require.NoError(t, err)
 	}
 
-	ind, err = file.NewStateDiffIndexer(mocks.TestChainConfig, file.CSVTestConfig)
+	ind, err = file.NewStateDiffIndexer(mocks.TestChainConfig, file.CSVTestConfig, test.LegacyNodeInfo, true)
 	require.NoError(t, err)
 
 	db, err = postgres.SetupSQLXDB()

indexer/database/file/csv_writer.go

@@ -57,7 +57,8 @@ type tableRow struct {
 
 type CSVWriter struct {
 	// dir containing output files
-	dir string
+	dir    string
+	isDiff bool
 
 	writers                fileWriters
 	watchedAddressesWriter fileWriter
@@ -128,7 +129,7 @@ func (tx fileWriters) flush() error {
 	return nil
 }
 
-func NewCSVWriter(path string, watchedAddressesFilePath string) (*CSVWriter, error) {
+func NewCSVWriter(path string, watchedAddressesFilePath string, diff bool) (*CSVWriter, error) {
 	if err := os.MkdirAll(path, 0777); err != nil {
 		return nil, fmt.Errorf("unable to create directory '%s': %w", path, err)
 	}
@@ -147,6 +148,7 @@ func NewCSVWriter(path string, watchedAddressesFilePath string) (*CSVWriter, err
 		writers:                writers,
 		watchedAddressesWriter: watchedAddressesWriter,
 		dir:                    path,
+		isDiff:                 diff,
 		rows:                   make(chan tableRow),
 		flushChan:              make(chan struct{}),
 		flushFinished:          make(chan struct{}),
@@ -278,14 +280,14 @@ func (csw *CSVWriter) upsertStateCID(stateNode models.StateNodeModel) {
 
 	var values []interface{}
 	values = append(values, stateNode.BlockNumber, stateNode.HeaderID, stateNode.StateKey, stateNode.CID,
-		true, balance, strconv.FormatUint(stateNode.Nonce, 10), stateNode.CodeHash, stateNode.StorageRoot, stateNode.Removed)
+		csw.isDiff, balance, strconv.FormatUint(stateNode.Nonce, 10), stateNode.CodeHash, stateNode.StorageRoot, stateNode.Removed)
 	csw.rows <- tableRow{schema.TableStateNode, values}
 }
 
 func (csw *CSVWriter) upsertStorageCID(storageCID models.StorageNodeModel) {
 	var values []interface{}
 	values = append(values, storageCID.BlockNumber, storageCID.HeaderID, storageCID.StateKey, storageCID.StorageKey, storageCID.CID,
-		true, storageCID.Value, storageCID.Removed)
+		csw.isDiff, storageCID.Value, storageCID.Removed)
 	csw.rows <- tableRow{schema.TableStorageNode, values}
 }
 

indexer/database/file/indexer.go

@@ -17,6 +17,7 @@
 package file
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"math/big"
@@ -37,6 +38,7 @@ import (
 	"github.com/cerc-io/plugeth-statediff/indexer/interfaces"
 	"github.com/cerc-io/plugeth-statediff/indexer/ipld"
 	"github.com/cerc-io/plugeth-statediff/indexer/models"
+	"github.com/cerc-io/plugeth-statediff/indexer/node"
 	"github.com/cerc-io/plugeth-statediff/indexer/shared"
 	sdtypes "github.com/cerc-io/plugeth-statediff/types"
 	"github.com/cerc-io/plugeth-statediff/utils/log"
@@ -57,11 +59,15 @@ type StateDiffIndexer struct {
 	chainConfig      *params.ChainConfig
 	nodeID           string
 	wg               *sync.WaitGroup
-	removedCacheFlag *uint32
+	removedCacheFlag uint32
 }
 
 // NewStateDiffIndexer creates a void implementation of interfaces.StateDiffIndexer
-func NewStateDiffIndexer(chainConfig *params.ChainConfig, config Config) (*StateDiffIndexer, error) {
+// `chainConfig` is used when processing chain state.
+// `config` contains configuration specific to the indexer.
+// `nodeInfo` contains metadata on the Ethereum node, which is inserted with the indexed state.
+// `isDiff` means to mark state nodes as belonging to an incremental diff, as opposed to a full snapshot.
+func NewStateDiffIndexer(chainConfig *params.ChainConfig, config Config, nodeInfo node.Info, isDiff bool) (*StateDiffIndexer, error) {
 	var err error
 	var writer FileWriter
 
@@ -84,7 +90,7 @@ func NewStateDiffIndexer(chainConfig *params.ChainConfig, config Config) (*State
 		}
 		log.Info("Writing watched addresses to file", "file", watchedAddressesFilePath)
 
-		writer, err = NewCSVWriter(outputDir, watchedAddressesFilePath)
+		writer, err = NewCSVWriter(outputDir, watchedAddressesFilePath, isDiff)
 		if err != nil {
 			return nil, err
 		}
@@ -107,19 +113,19 @@ func NewStateDiffIndexer(chainConfig *params.ChainConfig, config Config) (*State
 		}
 		log.Info("Writing watched addresses to file", "file", watchedAddressesFilePath)
 
-		writer = NewSQLWriter(file, watchedAddressesFilePath)
+		writer = NewSQLWriter(file, watchedAddressesFilePath, isDiff)
 	default:
 		return nil, fmt.Errorf("unrecognized file mode: %s", config.Mode)
 	}
 
 	wg := new(sync.WaitGroup)
 	writer.Loop()
-	writer.upsertNode(config.NodeInfo)
+	writer.upsertNode(nodeInfo)
 
 	return &StateDiffIndexer{
 		fileWriter:  writer,
 		chainConfig: chainConfig,
-		nodeID:      config.NodeInfo.ID,
+		nodeID:      nodeInfo.ID,
 		wg:          wg,
 	}, nil
 }
@@ -130,8 +136,7 @@ func (sdi *StateDiffIndexer) ReportDBMetrics(time.Duration, <-chan bool) {}
 // PushBlock pushes and indexes block data in sql, except state & storage nodes (includes header, uncles, transactions & receipts)
 // Returns an initiated DB transaction which must be Closed via defer to commit or rollback
 func (sdi *StateDiffIndexer) PushBlock(block *types.Block, receipts types.Receipts, totalDifficulty *big.Int) (interfaces.Batch, error) {
-	sdi.removedCacheFlag = new(uint32)
+	t := time.Now()
-	start, t := time.Now(), time.Now()
 	blockHash := block.Hash()
 	blockHashStr := blockHash.String()
 	height := block.NumberU64()
@@ -143,7 +148,7 @@ func (sdi *StateDiffIndexer) PushBlock(block *types.Block, receipts types.Receip
 	}
 
 	// Generate the block iplds
-	headerNode, txNodes, rctNodes, logNodes, err := ipld.FromBlockAndReceipts(block, receipts)
+	txNodes, rctNodes, logNodes, err := ipld.FromBlockAndReceipts(block, receipts)
 	if err != nil {
 		return nil, fmt.Errorf("error creating IPLD nodes from block and receipts: %v", err)
 	}
@@ -160,32 +165,19 @@ func (sdi *StateDiffIndexer) PushBlock(block *types.Block, receipts types.Receip
 	} else {
 		reward = shared.CalcEthBlockReward(block.Header(), block.Uncles(), block.Transactions(), receipts)
 	}
-	t = time.Now()
 
 	blockTx := &BatchTx{
-		BlockNumber: block.Number().String(),
+		blockNum:   block.Number().String(),
-		submit: func(self *BatchTx, err error) error {
+		fileWriter: sdi.fileWriter,
-			tDiff := time.Since(t)
-			metrics.IndexerMetrics.StateStoreCodeProcessingTimer.Update(tDiff)
-			traceMsg += fmt.Sprintf("state, storage, and code storage processing time: %s\r\n", tDiff.String())
-			t = time.Now()
-			sdi.fileWriter.Flush()
-			tDiff = time.Since(t)
-			metrics.IndexerMetrics.PostgresCommitTimer.Update(tDiff)
-			traceMsg += fmt.Sprintf("postgres transaction commit duration: %s\r\n", tDiff.String())
-			traceMsg += fmt.Sprintf(" TOTAL PROCESSING DURATION: %s\r\n", time.Since(start).String())
-			log.Trace(traceMsg)
-			return err
-		},
 	}
-	tDiff := time.Since(t)
-	metrics.IndexerMetrics.FreePostgresTimer.Update(tDiff)
-	traceMsg += fmt.Sprintf("time spent waiting for free postgres tx: %s:\r\n", tDiff.String())
 	t = time.Now()
 
 	// write header, collect headerID
-	headerID := sdi.processHeader(block.Header(), headerNode, reward, totalDifficulty)
+	headerID, err := sdi.PushHeader(blockTx, block.Header(), reward, totalDifficulty)
-	tDiff = time.Since(t)
+	if err != nil {
+		return nil, err
+	}
+	tDiff := time.Since(t)
 	metrics.IndexerMetrics.HeaderProcessingTimer.Update(tDiff)
 	traceMsg += fmt.Sprintf("header processing time: %s\r\n", tDiff.String())
 	t = time.Now()
@@ -218,16 +210,16 @@ func (sdi *StateDiffIndexer) PushBlock(block *types.Block, receipts types.Receip
 	return blockTx, err
 }
 
-// processHeader write a header IPLD insert SQL stmt to a file
+// PushHeader write a header IPLD insert SQL stmt to a file
 // it returns the headerID
-func (sdi *StateDiffIndexer) processHeader(header *types.Header, headerNode ipld.IPLD, reward, td *big.Int) string {
+func (sdi *StateDiffIndexer) PushHeader(_ interfaces.Batch, header *types.Header, reward, td *big.Int) (string, error) {
+	// Process the header
+	headerNode, err := ipld.NewEthHeader(header)
+	if err != nil {
+		return "", err
+	}
 	sdi.fileWriter.upsertIPLDNode(header.Number.String(), headerNode)
 
-	var baseFee *string
-	if header.BaseFee != nil {
-		baseFee = new(string)
-		*baseFee = header.BaseFee.String()
-	}
 	headerID := header.Hash().String()
 	sdi.fileWriter.upsertHeaderCID(models.HeaderModel{
 		NodeIDs:         pq.StringArray([]string{sdi.nodeID}),
@@ -246,7 +238,7 @@ func (sdi *StateDiffIndexer) processHeader(header *types.Header, headerNode ipld
 		Coinbase:        header.Coinbase.String(),
 		Canonical:       true,
 	})
-	return headerID
+	return headerID, nil
 }
 
 // processUncles publishes and indexes uncle IPLDs in Postgres
@@ -380,20 +372,16 @@ func (sdi *StateDiffIndexer) processReceiptsAndTxs(args processArgs) error {
 }
 
 // PushStateNode writes a state diff node object (including any child storage nodes) IPLD insert SQL stmt to a file
-func (sdi *StateDiffIndexer) PushStateNode(batch interfaces.Batch, stateNode sdtypes.StateLeafNode, headerID string) error {
+func (sdi *StateDiffIndexer) PushStateNode(tx interfaces.Batch, stateNode sdtypes.StateLeafNode, headerID string) error {
-	tx, ok := batch.(*BatchTx)
-	if !ok {
-		return fmt.Errorf("file: batch is expected to be of type %T, got %T", &BatchTx{}, batch)
-	}
 	// publish the state node
 	var stateModel models.StateNodeModel
 	if stateNode.Removed {
-		if atomic.LoadUint32(sdi.removedCacheFlag) == 0 {
+		if atomic.LoadUint32(&sdi.removedCacheFlag) == 0 {
-			atomic.StoreUint32(sdi.removedCacheFlag, 1)
+			atomic.StoreUint32(&sdi.removedCacheFlag, 1)
-			sdi.fileWriter.upsertIPLDDirect(tx.BlockNumber, shared.RemovedNodeStateCID, []byte{})
+			sdi.fileWriter.upsertIPLDDirect(tx.BlockNumber(), shared.RemovedNodeStateCID, []byte{})
 		}
 		stateModel = models.StateNodeModel{
-			BlockNumber: tx.BlockNumber,
+			BlockNumber: tx.BlockNumber(),
 			HeaderID:    headerID,
 			StateKey:    common.BytesToHash(stateNode.AccountWrapper.LeafKey).String(),
 			CID:         shared.RemovedNodeStateCID,
@@ -401,7 +389,7 @@ func (sdi *StateDiffIndexer) PushStateNode(batch interfaces.Batch, stateNode sdt
 		}
 	} else {
 		stateModel = models.StateNodeModel{
-			BlockNumber: tx.BlockNumber,
+			BlockNumber: tx.BlockNumber(),
 			HeaderID:    headerID,
 			StateKey:    common.BytesToHash(stateNode.AccountWrapper.LeafKey).String(),
 			CID:         stateNode.AccountWrapper.CID,
@@ -419,12 +407,12 @@ func (sdi *StateDiffIndexer) PushStateNode(batch interfaces.Batch, stateNode sdt
 	// if there are any storage nodes associated with this node, publish and index them
 	for _, storageNode := range stateNode.StorageDiff {
 		if storageNode.Removed {
-			if atomic.LoadUint32(sdi.removedCacheFlag) == 0 {
+			if atomic.LoadUint32(&sdi.removedCacheFlag) == 0 {
-				atomic.StoreUint32(sdi.removedCacheFlag, 1)
+				atomic.StoreUint32(&sdi.removedCacheFlag, 1)
-				sdi.fileWriter.upsertIPLDDirect(tx.BlockNumber, shared.RemovedNodeStorageCID, []byte{})
+				sdi.fileWriter.upsertIPLDDirect(tx.BlockNumber(), shared.RemovedNodeStorageCID, []byte{})
 			}
 			storageModel := models.StorageNodeModel{
-				BlockNumber: tx.BlockNumber,
+				BlockNumber: tx.BlockNumber(),
 				HeaderID:    headerID,
 				StateKey:    common.BytesToHash(stateNode.AccountWrapper.LeafKey).String(),
 				StorageKey:  common.BytesToHash(storageNode.LeafKey).String(),
@@ -436,7 +424,7 @@ func (sdi *StateDiffIndexer) PushStateNode(batch interfaces.Batch, stateNode sdt
 			continue
 		}
 		storageModel := models.StorageNodeModel{
-			BlockNumber: tx.BlockNumber,
+			BlockNumber: tx.BlockNumber(),
 			HeaderID:    headerID,
 			StateKey:    common.BytesToHash(stateNode.AccountWrapper.LeafKey).String(),
 			StorageKey:  common.BytesToHash(storageNode.LeafKey).String(),
@@ -451,12 +439,8 @@ func (sdi *StateDiffIndexer) PushStateNode(batch interfaces.Batch, stateNode sdt
 }
 
 // PushIPLD writes iplds to ipld.blocks
-func (sdi *StateDiffIndexer) PushIPLD(batch interfaces.Batch, ipld sdtypes.IPLD) error {
+func (sdi *StateDiffIndexer) PushIPLD(tx interfaces.Batch, ipld sdtypes.IPLD) error {
-	tx, ok := batch.(*BatchTx)
+	sdi.fileWriter.upsertIPLDDirect(tx.BlockNumber(), ipld.CID, ipld.Content)
-	if !ok {
-		return fmt.Errorf("file: batch is expected to be of type %T, got %T", &BatchTx{}, batch)
-	}
-	sdi.fileWriter.upsertIPLDDirect(tx.BlockNumber, ipld.CID, ipld.Content)
 	return nil
 }
 
@@ -478,6 +462,13 @@ func (sdi *StateDiffIndexer) HasBlock(hash common.Hash, number uint64) (bool, er
 	return false, nil
 }
 
+func (sdi *StateDiffIndexer) BeginTx(number *big.Int, _ context.Context) interfaces.Batch {
+	return &BatchTx{
+		blockNum:   number.String(),
+		fileWriter: sdi.fileWriter,
+	}
+}
+
 // Close satisfies io.Closer
 func (sdi *StateDiffIndexer) Close() error {
 	return sdi.fileWriter.Close()

indexer/database/file/mainnet_tests/indexer_test.go

@@ -83,7 +83,7 @@ func setupMainnetIndexer(t *testing.T) {
 		require.NoError(t, err)
 	}
 
-	ind, err = file.NewStateDiffIndexer(chainConf, file.CSVTestConfig)
+	ind, err = file.NewStateDiffIndexer(chainConf, file.CSVTestConfig, test.LegacyNodeInfo, true)
 	require.NoError(t, err)
 
 	db, err = postgres.SetupSQLXDB()

indexer/database/file/sql_indexer_legacy_test.go

@@ -44,7 +44,7 @@ func setupLegacySQLIndexer(t *testing.T) {
 		require.NoError(t, err)
 	}
 
-	ind, err = file.NewStateDiffIndexer(test.LegacyConfig, file.SQLTestConfig)
+	ind, err = file.NewStateDiffIndexer(test.LegacyConfig, file.SQLTestConfig, test.LegacyNodeInfo, true)
 	require.NoError(t, err)
 
 	db, err = postgres.SetupSQLXDB()

indexer/database/file/sql_indexer_test.go

@@ -41,7 +41,7 @@ func setupIndexer(t *testing.T) {
 		require.NoError(t, err)
 	}
 
-	ind, err = file.NewStateDiffIndexer(mocks.TestChainConfig, file.SQLTestConfig)
+	ind, err = file.NewStateDiffIndexer(mocks.TestChainConfig, file.SQLTestConfig, test.LegacyNodeInfo, true)
 	require.NoError(t, err)
 
 	db, err = postgres.SetupSQLXDB()

indexer/database/file/sql_writer.go

@@ -44,6 +44,7 @@ var (
 type SQLWriter struct {
 	wc             io.WriteCloser
 	stmts          chan []byte
+	isDiff         bool
 	collatedStmt   []byte
 	collationIndex int
 
@@ -55,11 +56,15 @@ type SQLWriter struct {
 	watchedAddressesFilePath string
 }
 
-// NewSQLWriter creates a new pointer to a Writer
+// NewSQLWriter creates a new Writer.
-func NewSQLWriter(wc io.WriteCloser, watchedAddressesFilePath string) *SQLWriter {
+// `wc` is the underlying io.WriteCloser to write to.
+// `watchedAddressesFilePath` is the path to the file containing watched addresses.
+// `isDiff` means to mark state nodes as belonging to an incremental diff, as opposed to a full snapshot.
+func NewSQLWriter(wc io.WriteCloser, watchedAddressesFilePath string, isDiff bool) *SQLWriter {
 	return &SQLWriter{
 		wc:                       wc,
 		stmts:                    make(chan []byte),
+		isDiff:                   isDiff,
 		collatedStmt:             make([]byte, writeBufferSize),
 		flushChan:                make(chan struct{}),
 		flushFinished:            make(chan struct{}),
@@ -225,12 +230,12 @@ func (sqw *SQLWriter) upsertStateCID(stateNode models.StateNodeModel) {
 		balance = "0"
 	}
 	sqw.stmts <- []byte(fmt.Sprintf(stateInsert, stateNode.BlockNumber, stateNode.HeaderID, stateNode.StateKey, stateNode.CID,
-		stateNode.Removed, true, balance, stateNode.Nonce, stateNode.CodeHash, stateNode.StorageRoot))
+		stateNode.Removed, sqw.isDiff, balance, stateNode.Nonce, stateNode.CodeHash, stateNode.StorageRoot))
 }
 
 func (sqw *SQLWriter) upsertStorageCID(storageCID models.StorageNodeModel) {
 	sqw.stmts <- []byte(fmt.Sprintf(storageInsert, storageCID.BlockNumber, storageCID.HeaderID, storageCID.StateKey, storageCID.StorageKey, storageCID.CID,
-		storageCID.Removed, true, storageCID.Value))
+		storageCID.Removed, sqw.isDiff, storageCID.Value))
 }
 
 // LoadWatchedAddresses loads watched addresses from a file

indexer/database/sql/batch_tx.go

@@ -18,11 +18,14 @@ package sql
 
 import (
 	"context"
+	"math/big"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/lib/pq"
 
+	"github.com/cerc-io/plugeth-statediff/indexer/database/metrics"
 	"github.com/cerc-io/plugeth-statediff/indexer/ipld"
 	"github.com/cerc-io/plugeth-statediff/indexer/models"
 	"github.com/cerc-io/plugeth-statediff/utils/log"
@@ -32,7 +35,7 @@ const startingCacheCapacity = 1024 * 24
 
 // BatchTx wraps a sql tx with the state necessary for building the tx concurrently during trie difference iteration
 type BatchTx struct {
-	BlockNumber      string
+	blockNum         string
 	ctx              context.Context
 	dbtx             Tx
 	stm              string
@@ -42,13 +45,68 @@ type BatchTx struct {
 	removedCacheFlag *uint32
 	// Tracks expected cache size and ensures cache is caught up before flush
 	cacheWg sync.WaitGroup
-
-	submit func(blockTx *BatchTx, err error) error
 }
 
-// Submit satisfies indexer.AtomicTx
+func NewBatch(insertStm string, ctx context.Context, number *big.Int, tx Tx) *BatchTx {
-func (tx *BatchTx) Submit(err error) error {
+	blockTx := &BatchTx{
-	return tx.submit(tx, err)
+		removedCacheFlag: new(uint32),
+		ctx:              ctx,
+		blockNum:         number.String(),
+		stm:              insertStm,
+		iplds:            make(chan models.IPLDModel),
+		quit:             make(chan (chan<- struct{})),
+		ipldCache: models.IPLDBatch{
+			BlockNumbers: make([]string, 0, startingCacheCapacity),
+			Keys:         make([]string, 0, startingCacheCapacity),
+			Values:       make([][]byte, 0, startingCacheCapacity),
+		},
+		dbtx: tx,
+	}
+	go blockTx.cache()
+	return blockTx
+}
+
+// Submit satisfies indexer.Batch
+func (tx *BatchTx) Submit() error {
+	defer tx.close()
+
+	t := time.Now()
+	if err := tx.flush(); err != nil {
+		rollback(tx.ctx, tx.dbtx)
+		return err
+	}
+	err := tx.dbtx.Commit(tx.ctx)
+	metrics.IndexerMetrics.PostgresCommitTimer.Update(time.Since(t))
+	return err
+}
+
+func (tx *BatchTx) BlockNumber() string {
+	return tx.blockNum
+}
+
+func (tx *BatchTx) RollbackOnFailure(err error) {
+	if p := recover(); p != nil {
+		defer tx.close()
+		log.Info("panic detected before tx submission, rolling back the tx", "panic", p)
+		rollback(tx.ctx, tx.dbtx)
+		panic(p)
+	} else if err != nil {
+		defer tx.close()
+		log.Info("error detected before tx submission, rolling back the tx", "error", err)
+		rollback(tx.ctx, tx.dbtx)
+	}
+}
+
+func (tx *BatchTx) close() {
+	if tx.quit == nil {
+		return
+	}
+	confirm := make(chan struct{})
+	tx.quit <- confirm
+	close(tx.quit)
+	<-confirm
+	close(tx.iplds)
+	tx.quit = nil
 }
 
 func (tx *BatchTx) flush() error {
@@ -92,7 +150,7 @@ func (tx *BatchTx) cache() {
 func (tx *BatchTx) cacheDirect(key string, value []byte) {
 	tx.cacheWg.Add(1)
 	tx.iplds <- models.IPLDModel{
-		BlockNumber: tx.BlockNumber,
+		BlockNumber: tx.BlockNumber(),
 		Key:         key,
 		Data:        value,
 	}
@@ -101,7 +159,7 @@ func (tx *BatchTx) cacheDirect(key string, value []byte) {
 func (tx *BatchTx) cacheIPLD(i ipld.IPLD) {
 	tx.cacheWg.Add(1)
 	tx.iplds <- models.IPLDModel{
-		BlockNumber: tx.BlockNumber,
+		BlockNumber: tx.BlockNumber(),
 		Key:         i.Cid().String(),
 		Data:        i.RawData(),
 	}
@@ -112,7 +170,7 @@ func (tx *BatchTx) cacheRemoved(key string, value []byte) {
 		atomic.StoreUint32(tx.removedCacheFlag, 1)
 		tx.cacheWg.Add(1)
 		tx.iplds <- models.IPLDModel{
-			BlockNumber: tx.BlockNumber,
+			BlockNumber: tx.BlockNumber(),
 			Key:         key,
 			Data:        value,
 		}

indexer/database/sql/indexer.go

@@ -39,12 +39,11 @@ import (
 	"github.com/cerc-io/plugeth-statediff/indexer/models"
 	"github.com/cerc-io/plugeth-statediff/indexer/shared"
 	sdtypes "github.com/cerc-io/plugeth-statediff/types"
-	"github.com/cerc-io/plugeth-statediff/utils/log"
 )
 
 var _ interfaces.StateDiffIndexer = &StateDiffIndexer{}
 
-// StateDiffIndexer satisfies the indexer.StateDiffIndexer interface for ethereum statediff objects on top of an SQL sql
+// StateDiffIndexer satisfies the indexer.StateDiffIndexer interface for ethereum statediff objects on top of an SQL DB.
 type StateDiffIndexer struct {
 	ctx         context.Context
 	chainConfig *params.ChainConfig
@@ -52,11 +51,17 @@ type StateDiffIndexer struct {
 }
 
 // NewStateDiffIndexer creates a sql implementation of interfaces.StateDiffIndexer
-func NewStateDiffIndexer(ctx context.Context, chainConfig *params.ChainConfig, db Database) (*StateDiffIndexer, error) {
+// `ctx` is used to cancel the underlying DB connection.
+// `chainConfig` is used when processing chain state.
+// `db` is the backing database to use.
+// `isDiff` means to mark state nodes as belonging to an incremental diff, as opposed to a full snapshot.
+func NewStateDiffIndexer(
+	ctx context.Context, chainConfig *params.ChainConfig, db Database, isDiff bool,
+) (*StateDiffIndexer, error) {
 	return &StateDiffIndexer{
 		ctx:         ctx,
 		chainConfig: chainConfig,
-		dbWriter:    NewWriter(db),
+		dbWriter:    NewWriter(db, isDiff),
 	}, nil
 }
 
@@ -82,25 +87,26 @@ func (sdi *StateDiffIndexer) ReportDBMetrics(delay time.Duration, quit <-chan bo
 // PushBlock pushes and indexes block data in sql, except state & storage nodes (includes header, uncles, transactions & receipts)
 // Returns an initiated DB transaction which must be Closed via defer to commit or rollback
 func (sdi *StateDiffIndexer) PushBlock(block *types.Block, receipts types.Receipts, totalDifficulty *big.Int) (interfaces.Batch, error) {
-	start, t := time.Now(), time.Now()
+	t := time.Now()
 	blockHash := block.Hash()
-	blockHashStr := blockHash.String()
 	height := block.NumberU64()
-	traceMsg := fmt.Sprintf("indexer stats for statediff at %d with hash %s:\r\n", height, blockHashStr)
 	transactions := block.Transactions()
+	var err error
+
 	// Derive any missing fields
 	if err := receipts.DeriveFields(sdi.chainConfig, blockHash, height, block.BaseFee(), transactions); err != nil {
 		return nil, err
 	}
 
 	// Generate the block iplds
-	headerNode, txNodes, rctNodes, logNodes, err := ipld.FromBlockAndReceipts(block, receipts)
+	txNodes, rctNodes, logNodes, err := ipld.FromBlockAndReceipts(block, receipts)
 	if err != nil {
-		return nil, fmt.Errorf("error creating IPLD nodes from block and receipts: %v", err)
+		return nil, fmt.Errorf("error creating IPLD nodes from block and receipts: %w", err)
 	}
 
 	if len(txNodes) != len(rctNodes) {
-		return nil, fmt.Errorf("expected number of transactions (%d), receipts (%d)", len(txNodes), len(rctNodes))
+		return nil, fmt.Errorf("expected number of transactions (%d) does not match number of receipts (%d)",
+			len(txNodes), len(rctNodes))
 	}
 
 	// Calculate reward
@@ -109,99 +115,35 @@ func (sdi *StateDiffIndexer) PushBlock(block *types.Block, receipts types.Receip
 	if sdi.chainConfig.Clique != nil {
 		reward = big.NewInt(0)
 	} else {
-		reward = shared.CalcEthBlockReward(block.Header(), block.Uncles(), block.Transactions(), receipts)
+		reward = shared.CalcEthBlockReward(block.Header(), block.Uncles(), transactions, receipts)
 	}
 	t = time.Now()
 
 	// Begin new DB tx for everything
-	tx := NewDelayedTx(sdi.dbWriter.db)
+	batch := NewBatch(
-	defer func() {
+		sdi.dbWriter.db.InsertIPLDsStm(), sdi.ctx,
-		if p := recover(); p != nil {
+		block.Number(),
-			rollback(sdi.ctx, tx)
+		NewDelayedTx(sdi.dbWriter.db),
-			panic(p)
+	)
-		} else if err != nil {
+	// handle transaction rollback for failures in this scope
-			rollback(sdi.ctx, tx)
+	defer batch.RollbackOnFailure(err)
-		}
-	}()
-	blockTx := &BatchTx{
-		removedCacheFlag: new(uint32),
-		ctx:              sdi.ctx,
-		BlockNumber:      block.Number().String(),
-		stm:              sdi.dbWriter.db.InsertIPLDsStm(),
-		iplds:            make(chan models.IPLDModel),
-		quit:             make(chan (chan<- struct{})),
-		ipldCache: models.IPLDBatch{
-			BlockNumbers: make([]string, 0, startingCacheCapacity),
-			Keys:         make([]string, 0, startingCacheCapacity),
-			Values:       make([][]byte, 0, startingCacheCapacity),
-		},
-		dbtx: tx,
-		// handle transaction commit or rollback for any return case
-		submit: func(self *BatchTx, err error) error {
-			defer func() {
-				confirm := make(chan struct{})
-				self.quit <- confirm
-				close(self.quit)
-				<-confirm
-				close(self.iplds)
-			}()
-			if p := recover(); p != nil {
-				log.Info("panic detected before tx submission, rolling back the tx", "panic", p)
-				rollback(sdi.ctx, tx)
-				panic(p)
-			} else if err != nil {
-				log.Info("error detected before tx submission, rolling back the tx", "error", err)
-				rollback(sdi.ctx, tx)
-			} else {
-				tDiff := time.Since(t)
-				metrics2.IndexerMetrics.StateStoreCodeProcessingTimer.Update(tDiff)
-				traceMsg += fmt.Sprintf("state, storage, and code storage processing time: %s\r\n", tDiff.String())
-				t = time.Now()
-				if err := self.flush(); err != nil {
-					rollback(sdi.ctx, tx)
-					traceMsg += fmt.Sprintf(" TOTAL PROCESSING DURATION: %s\r\n", time.Since(start).String())
-					log.Debug(traceMsg)
-					return err
-				}
-				err = tx.Commit(sdi.ctx)
-				tDiff = time.Since(t)
-				metrics2.IndexerMetrics.PostgresCommitTimer.Update(tDiff)
-				traceMsg += fmt.Sprintf("postgres transaction commit duration: %s\r\n", tDiff.String())
-			}
-			traceMsg += fmt.Sprintf(" TOTAL PROCESSING DURATION: %s\r\n", time.Since(start).String())
-			log.Debug(traceMsg)
-			return err
-		},
-	}
-	go blockTx.cache()
-
-	tDiff := time.Since(t)
-	metrics2.IndexerMetrics.FreePostgresTimer.Update(tDiff)
-
-	traceMsg += fmt.Sprintf("time spent waiting for free postgres tx: %s:\r\n", tDiff.String())
-	t = time.Now()
 
 	// Publish and index header, collect headerID
-	var headerID string
+	headerID, err := sdi.PushHeader(batch, block.Header(), reward, totalDifficulty)
-	headerID, err = sdi.processHeader(blockTx, block.Header(), headerNode, reward, totalDifficulty)
 	if err != nil {
 		return nil, err
 	}
-	tDiff = time.Since(t)
+	metrics2.IndexerMetrics.HeaderProcessingTimer.Update(time.Since(t))
-	metrics2.IndexerMetrics.HeaderProcessingTimer.Update(tDiff)
-	traceMsg += fmt.Sprintf("header processing time: %s\r\n", tDiff.String())
 	t = time.Now()
 	// Publish and index uncles
-	err = sdi.processUncles(blockTx, headerID, block.Number(), block.UncleHash(), block.Uncles())
+	err = sdi.processUncles(batch, headerID, block.Number(), block.UncleHash(), block.Uncles())
 	if err != nil {
 		return nil, err
 	}
-	tDiff = time.Since(t)
+	metrics2.IndexerMetrics.UncleProcessingTimer.Update(time.Since(t))
-	metrics2.IndexerMetrics.UncleProcessingTimer.Update(tDiff)
-	traceMsg += fmt.Sprintf("uncle processing time: %s\r\n", tDiff.String())
 	t = time.Now()
 	// Publish and index receipts and txs
-	err = sdi.processReceiptsAndTxs(blockTx, processArgs{
+	err = sdi.processReceiptsAndTxs(batch, processArgs{
 		headerID:    headerID,
 		blockNumber: block.Number(),
 		receipts:    receipts,
@@ -213,12 +155,9 @@ func (sdi *StateDiffIndexer) PushBlock(block *types.Block, receipts types.Receip
 	if err != nil {
 		return nil, err
 	}
-	tDiff = time.Since(t)
+	metrics2.IndexerMetrics.TxAndRecProcessingTimer.Update(time.Since(t))
-	metrics2.IndexerMetrics.TxAndRecProcessingTimer.Update(tDiff)
-	traceMsg += fmt.Sprintf("tx and receipt processing time: %s\r\n", tDiff.String())
-	t = time.Now()
 
-	return blockTx, err
+	return batch, err
 }
 
 // CurrentBlock returns the HeaderModel of the highest existing block in the database.
@@ -231,16 +170,20 @@ func (sdi *StateDiffIndexer) DetectGaps(beginBlockNumber uint64, endBlockNumber
 	return sdi.dbWriter.detectGaps(beginBlockNumber, endBlockNumber)
 }
 
-// processHeader publishes and indexes a header IPLD in Postgres
+// PushHeader publishes and indexes a header IPLD in Postgres
 // it returns the headerID
-func (sdi *StateDiffIndexer) processHeader(tx *BatchTx, header *types.Header, headerNode ipld.IPLD, reward, td *big.Int) (string, error) {
+func (sdi *StateDiffIndexer) PushHeader(batch interfaces.Batch, header *types.Header, reward, td *big.Int) (string, error) {
+	tx, ok := batch.(*BatchTx)
+	if !ok {
+		return "", fmt.Errorf("sql: batch is expected to be of type %T, got %T", &BatchTx{}, batch)
+	}
+	// Process the header
+	headerNode, err := ipld.NewEthHeader(header)
+	if err != nil {
+		return "", err
+	}
 	tx.cacheIPLD(headerNode)
 
-	var baseFee *string
-	if header.BaseFee != nil {
-		baseFee = new(string)
-		*baseFee = header.BaseFee.String()
-	}
 	headerID := header.Hash().String()
 	// index header
 	return headerID, sdi.dbWriter.upsertHeaderCID(tx.dbtx, models.HeaderModel{
@@ -412,7 +355,7 @@ func (sdi *StateDiffIndexer) PushStateNode(batch interfaces.Batch, stateNode sdt
 	if stateNode.Removed {
 		tx.cacheRemoved(shared.RemovedNodeStateCID, []byte{})
 		stateModel = models.StateNodeModel{
-			BlockNumber: tx.BlockNumber,
+			BlockNumber: tx.BlockNumber(),
 			HeaderID:    headerID,
 			StateKey:    common.BytesToHash(stateNode.AccountWrapper.LeafKey).String(),
 			CID:         shared.RemovedNodeStateCID,
@@ -420,7 +363,7 @@ func (sdi *StateDiffIndexer) PushStateNode(batch interfaces.Batch, stateNode sdt
 		}
 	} else {
 		stateModel = models.StateNodeModel{
-			BlockNumber: tx.BlockNumber,
+			BlockNumber: tx.BlockNumber(),
 			HeaderID:    headerID,
 			StateKey:    common.BytesToHash(stateNode.AccountWrapper.LeafKey).String(),
 			CID:         stateNode.AccountWrapper.CID,
@@ -442,7 +385,7 @@ func (sdi *StateDiffIndexer) PushStateNode(batch interfaces.Batch, stateNode sdt
 		if storageNode.Removed {
 			tx.cacheRemoved(shared.RemovedNodeStorageCID, []byte{})
 			storageModel := models.StorageNodeModel{
-				BlockNumber: tx.BlockNumber,
+				BlockNumber: tx.BlockNumber(),
 				HeaderID:    headerID,
 				StateKey:    common.BytesToHash(stateNode.AccountWrapper.LeafKey).String(),
 				StorageKey:  common.BytesToHash(storageNode.LeafKey).String(),
@@ -456,7 +399,7 @@ func (sdi *StateDiffIndexer) PushStateNode(batch interfaces.Batch, stateNode sdt
 			continue
 		}
 		storageModel := models.StorageNodeModel{
-			BlockNumber: tx.BlockNumber,
+			BlockNumber: tx.BlockNumber(),
 			HeaderID:    headerID,
 			StateKey:    common.BytesToHash(stateNode.AccountWrapper.LeafKey).String(),
 			StorageKey:  common.BytesToHash(storageNode.LeafKey).String(),
@@ -487,6 +430,15 @@ func (sdi *StateDiffIndexer) HasBlock(hash common.Hash, number uint64) (bool, er
 	return sdi.dbWriter.hasHeader(hash, number)
 }
 
+func (sdi *StateDiffIndexer) BeginTx(number *big.Int, ctx context.Context) interfaces.Batch {
+	return NewBatch(
+		sdi.dbWriter.db.InsertIPLDsStm(),
+		ctx,
+		number,
+		NewDelayedTx(sdi.dbWriter.db),
+	)
+}
+
 // Close satisfies io.Closer
 func (sdi *StateDiffIndexer) Close() error {
 	return sdi.dbWriter.Close()

indexer/database/sql/interfaces.go

@@ -58,18 +58,6 @@ type Statements interface {
 	InsertStorageStm() string
 	InsertIPLDStm() string
 	InsertIPLDsStm() string
-
-	// Table/column descriptions for use with CopyFrom and similar commands.
-	LogTableName() []string
-	LogColumnNames() []string
-	RctTableName() []string
-	RctColumnNames() []string
-	StateTableName() []string
-	StateColumnNames() []string
-	StorageTableName() []string
-	StorageColumnNames() []string
-	TxTableName() []string
-	TxColumnNames() []string
 }
 
 // Tx interface to accommodate different concrete SQL transaction types

indexer/database/sql/lazy_tx.go

@@ -3,7 +3,9 @@ package sql
 import (
 	"context"
 	"reflect"
+	"time"
 
+	"github.com/cerc-io/plugeth-statediff/indexer/database/metrics"
 	"github.com/cerc-io/plugeth-statediff/utils/log"
 )
 
@@ -69,10 +71,12 @@ func (tx *DelayedTx) Exec(ctx context.Context, sql string, args ...interface{})
 }
 
 func (tx *DelayedTx) Commit(ctx context.Context) error {
+	t := time.Now()
 	base, err := tx.db.Begin(ctx)
 	if err != nil {
 		return err
 	}
+	metrics.IndexerMetrics.FreePostgresTimer.Update(time.Since(t))
 	defer func() {
 		if p := recover(); p != nil {
 			rollback(ctx, base)

indexer/database/sql/mainnet_tests/indexer_test.go

@@ -71,7 +71,7 @@ func setupMainnetIndexer(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	ind, err = sql.NewStateDiffIndexer(context.Background(), chainConf, db)
+	ind, err = sql.NewStateDiffIndexer(context.Background(), chainConf, db, true)
 }
 
 func checkTxClosure(t *testing.T, idle, inUse, open int64) {

indexer/database/sql/pgx_indexer_legacy_test.go

@@ -33,7 +33,7 @@ func setupLegacyPGXIndexer(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	ind, err = sql.NewStateDiffIndexer(context.Background(), test.LegacyConfig, db)
+	ind, err = sql.NewStateDiffIndexer(context.Background(), test.LegacyConfig, db, true)
 	require.NoError(t, err)
 }
 

indexer/database/sql/pgx_indexer_test.go

@@ -44,7 +44,7 @@ func setupPGXIndexer(t *testing.T, config postgres.Config) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	ind, err = sql.NewStateDiffIndexer(context.Background(), mocks.TestChainConfig, db)
+	ind, err = sql.NewStateDiffIndexer(context.Background(), mocks.TestChainConfig, db, true)
 	require.NoError(t, err)
 }
 

indexer/database/sql/postgres/config.go

@@ -44,10 +44,6 @@ type Config struct {
 	ConnTimeout     time.Duration
 	LogStatements   bool
 
-	// node info params
-	ID         string
-	ClientName string
-
 	// driver type
 	Driver DriverType
 

indexer/database/sql/postgres/database.go

@@ -107,43 +107,3 @@ func (db *DB) InsertIPLDStm() string {
 func (db *DB) InsertIPLDsStm() string {
 	return `INSERT INTO ipld.blocks (block_number, key, data) VALUES (unnest($1::BIGINT[]), unnest($2::TEXT[]), unnest($3::BYTEA[])) ON CONFLICT DO NOTHING`
 }
-
-func (db *DB) LogTableName() []string {
-	return []string{"eth", "log_cids"}
-}
-
-func (db *DB) LogColumnNames() []string {
-	return []string{"block_number", "header_id", "cid", "rct_id", "address", "index", "topic0", "topic1", "topic2", "topic3"}
-}
-
-func (db *DB) RctTableName() []string {
-	return []string{"eth", "receipt_cids"}
-}
-
-func (db *DB) RctColumnNames() []string {
-	return []string{"block_number", "header_id", "tx_id", "cid", "contract", "post_state", "post_status"}
-}
-
-func (db *DB) StateTableName() []string {
-	return []string{"eth", "state_cids"}
-}
-
-func (db *DB) StateColumnNames() []string {
-	return []string{"block_number", "header_id", "state_leaf_key", "cid", "diff", "balance", "nonce", "code_hash", "storage_root", "removed"}
-}
-
-func (db *DB) StorageTableName() []string {
-	return []string{"eth", "storage_cids"}
-}
-
-func (db *DB) StorageColumnNames() []string {
-	return []string{"block_number", "header_id", "state_leaf_key", "storage_leaf_key", "cid", "diff", "val", "removed"}
-}
-
-func (db *DB) TxTableName() []string {
-	return []string{"eth", "transaction_cids"}
-}
-
-func (db *DB) TxColumnNames() []string {
-	return []string{"block_number", "header_id", "tx_hash", "cid", "dst", "src", "index", "tx_type", "value"}
-}

indexer/database/sql/postgres/pgx_test.go

@@ -36,12 +36,6 @@ var (
 	ctx          = context.Background()
 )
 
-func expectContainsSubstring(t *testing.T, full string, sub string) {
-	if !strings.Contains(full, sub) {
-		t.Fatalf("Expected \"%v\" to contain substring \"%v\"\n", full, sub)
-	}
-}
-
 func TestPostgresPGX(t *testing.T) {
 	t.Run("connects to the sql", func(t *testing.T) {
 		dbPool, err := pgxpool.ConnectConfig(context.Background(), pgxConfig)
@@ -105,7 +99,7 @@ func TestPostgresPGX(t *testing.T) {
 			t.Fatal("Expected an error")
 		}
 
-		expectContainsSubstring(t, err.Error(), postgres.DbConnectionFailedMsg)
+		require.Contains(t, err.Error(), postgres.DbConnectionFailedMsg)
 	})
 
 	t.Run("throws error when can't create node", func(t *testing.T) {
@@ -117,6 +111,6 @@ func TestPostgresPGX(t *testing.T) {
 			t.Fatal("Expected an error")
 		}
 
-		expectContainsSubstring(t, err.Error(), postgres.SettingNodeFailedMsg)
+		require.Contains(t, err.Error(), postgres.SettingNodeFailedMsg)
 	})
 }

indexer/database/sql/postgres/sqlx_test.go

@@ -102,7 +102,7 @@ func TestPostgresSQLX(t *testing.T) {
 			t.Fatal("Expected an error")
 		}
 
-		expectContainsSubstring(t, err.Error(), postgres.DbConnectionFailedMsg)
+		require.Contains(t, err.Error(), postgres.DbConnectionFailedMsg)
 	})
 
 	t.Run("throws error when can't create node", func(t *testing.T) {
@@ -114,6 +114,6 @@ func TestPostgresSQLX(t *testing.T) {
 			t.Fatal("Expected an error")
 		}
 
-		expectContainsSubstring(t, err.Error(), postgres.SettingNodeFailedMsg)
+		require.Contains(t, err.Error(), postgres.SettingNodeFailedMsg)
 	})
 }

indexer/database/sql/sqlx_indexer_legacy_test.go

@@ -32,7 +32,7 @@ func setupLegacySQLXIndexer(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	ind, err = sql.NewStateDiffIndexer(context.Background(), test.LegacyConfig, db)
+	ind, err = sql.NewStateDiffIndexer(context.Background(), test.LegacyConfig, db, true)
 	require.NoError(t, err)
 }
 

indexer/database/sql/sqlx_indexer_test.go

@@ -34,7 +34,7 @@ func setupSQLXIndexer(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	ind, err = sql.NewStateDiffIndexer(context.Background(), mocks.TestChainConfig, db)
+	ind, err = sql.NewStateDiffIndexer(context.Background(), mocks.TestChainConfig, db, true)
 	require.NoError(t, err)
 }
 

indexer/database/sql/writer.go

@@ -29,17 +29,20 @@ import (
 	"github.com/cerc-io/plugeth-statediff/indexer/database/metrics"
 	"github.com/cerc-io/plugeth-statediff/indexer/interfaces"
 	"github.com/cerc-io/plugeth-statediff/indexer/models"
+	"github.com/cerc-io/plugeth-statediff/indexer/shared/schema"
 )
 
 // Writer handles processing and writing of indexed IPLD objects to Postgres
 type Writer struct {
-	db Database
+	db     Database
+	isDiff bool
 }
 
-// NewWriter creates a new pointer to a Writer
+// NewWriter creates a new pointer to a Writer. `diff` indicates whether this is part of an
-func NewWriter(db Database) *Writer {
+// incremental diff (as opposed to a snapshot).
+func NewWriter(db Database, diff bool) *Writer {
 	return &Writer{
-		db: db,
+		db: db, isDiff: diff,
 	}
 }
 
@@ -65,7 +68,8 @@ func (w *Writer) detectGaps(beginBlockNumber uint64, endBlockNumber uint64) ([]*
 	// pgx misdetects the parameter OIDs and selects int8, which can overflow.
 	// unfortunately there is no good place to override it, so it is safer to pass the uint64s as text
 	// and let PG handle the cast
-	err := w.db.Select(w.db.Context(), &gaps, w.db.DetectGapsStm(), strconv.FormatUint(beginBlockNumber, 10), strconv.FormatUint(endBlockNumber, 10))
+	err := w.db.Select(w.db.Context(), &gaps, w.db.DetectGapsStm(),
+		strconv.FormatUint(beginBlockNumber, 10), strconv.FormatUint(endBlockNumber, 10))
 	return gaps, err
 }
 
@@ -177,7 +181,8 @@ func (w *Writer) upsertTransactionCID(tx Tx, transaction models.TxModel) error {
 			return insertError{"eth.transaction_cids", err, "COPY", transaction}
 		}
 
-		_, err = tx.CopyFrom(w.db.Context(), w.db.TxTableName(), w.db.TxColumnNames(),
+		_, err = tx.CopyFrom(w.db.Context(),
+			schema.TableTransaction.TableName(), schema.TableTransaction.ColumnNames(),
 			toRows(toRow(blockNum, transaction.HeaderID, transaction.TxHash, transaction.CID, transaction.Dst,
 				transaction.Src, transaction.Index, int(transaction.Type), value)))
 		if err != nil {
@@ -213,7 +218,7 @@ func (w *Writer) upsertReceiptCID(tx Tx, rct *models.ReceiptModel) error {
 			return insertError{"eth.receipt_cids", err, "COPY", rct}
 		}
 
-		_, err = tx.CopyFrom(w.db.Context(), w.db.RctTableName(), w.db.RctColumnNames(),
+		_, err = tx.CopyFrom(w.db.Context(), schema.TableReceipt.TableName(), schema.TableReceipt.ColumnNames(),
 			toRows(toRow(blockNum, rct.HeaderID, rct.TxID, rct.CID, rct.Contract,
 				rct.PostState, int(rct.PostStatus))))
 		if err != nil {
@@ -253,7 +258,7 @@ func (w *Writer) upsertLogCID(tx Tx, logs []*models.LogsModel) error {
 				log.Address, log.Index, log.Topic0, log.Topic1, log.Topic2, log.Topic3))
 		}
 		if nil != rows && len(rows) >= 0 {
-			_, err := tx.CopyFrom(w.db.Context(), w.db.LogTableName(), w.db.LogColumnNames(), rows)
+			_, err := tx.CopyFrom(w.db.Context(), schema.TableLog.TableName(), schema.TableLog.ColumnNames(), rows)
 			if err != nil {
 				return insertError{"eth.log_cids", err, "COPY", rows}
 			}
@@ -302,9 +307,11 @@ func (w *Writer) upsertStateCID(tx Tx, stateNode models.StateNodeModel) error {
 			return insertError{"eth.state_cids", err, "COPY", stateNode}
 		}
 
-		_, err = tx.CopyFrom(w.db.Context(), w.db.StateTableName(), w.db.StateColumnNames(),
+		_, err = tx.CopyFrom(w.db.Context(),
+			schema.TableStateNode.TableName(), schema.TableStateNode.ColumnNames(),
 			toRows(toRow(blockNum, stateNode.HeaderID, stateNode.StateKey, stateNode.CID,
-				true, balance, stateNode.Nonce, stateNode.CodeHash, stateNode.StorageRoot, stateNode.Removed)))
+				w.isDiff, balance, stateNode.Nonce, stateNode.CodeHash, stateNode.StorageRoot,
+				stateNode.Removed)))
 		if err != nil {
 			return insertError{"eth.state_cids", err, "COPY", stateNode}
 		}
@@ -314,7 +321,7 @@ func (w *Writer) upsertStateCID(tx Tx, stateNode models.StateNodeModel) error {
 			stateNode.HeaderID,
 			stateNode.StateKey,
 			stateNode.CID,
-			true,
+			w.isDiff,
 			bal,
 			stateNode.Nonce,
 			stateNode.CodeHash,
@@ -339,9 +346,10 @@ func (w *Writer) upsertStorageCID(tx Tx, storageCID models.StorageNodeModel) err
 			return insertError{"eth.storage_cids", err, "COPY", storageCID}
 		}
 
-		_, err = tx.CopyFrom(w.db.Context(), w.db.StorageTableName(), w.db.StorageColumnNames(),
+		_, err = tx.CopyFrom(w.db.Context(),
+			schema.TableStorageNode.TableName(), schema.TableStorageNode.ColumnNames(),
 			toRows(toRow(blockNum, storageCID.HeaderID, storageCID.StateKey, storageCID.StorageKey, storageCID.CID,
-				true, storageCID.Value, storageCID.Removed)))
+				w.isDiff, storageCID.Value, storageCID.Removed)))
 		if err != nil {
 			return insertError{"eth.storage_cids", err, "COPY", storageCID}
 		}
@@ -352,7 +360,7 @@ func (w *Writer) upsertStorageCID(tx Tx, storageCID models.StorageNodeModel) err
 			storageCID.StateKey,
 			storageCID.StorageKey,
 			storageCID.CID,
-			true,
+			w.isDiff,
 			storageCID.Value,
 			storageCID.Removed,
 		)

indexer/interfaces/interfaces.go

@@ -17,6 +17,7 @@
 package interfaces
 
 import (
+	"context"
 	"math/big"
 	"time"
 
@@ -34,10 +35,13 @@ type StateDiffIndexer interface {
 	CurrentBlock() (*models.HeaderModel, error)
 	HasBlock(hash common.Hash, number uint64) (bool, error)
 	PushBlock(block *types.Block, receipts types.Receipts, totalDifficulty *big.Int) (Batch, error)
+	PushHeader(batch Batch, header *types.Header, reward, td *big.Int) (string, error)
 	PushStateNode(tx Batch, stateNode sdtypes.StateLeafNode, headerID string) error
 	PushIPLD(tx Batch, ipld sdtypes.IPLD) error
 	ReportDBMetrics(delay time.Duration, quit <-chan bool)
 
+	BeginTx(number *big.Int, ctx context.Context) Batch
+
 	// Methods used by WatchAddress API/functionality
 	LoadWatchedAddresses() ([]common.Address, error)
 	InsertWatchedAddresses(addresses []sdtypes.WatchAddressArg, currentBlock *big.Int) error
@@ -50,7 +54,12 @@ type StateDiffIndexer interface {
 
 // Batch required for indexing data atomically
 type Batch interface {
-	Submit(err error) error
+	// Submit commits the batch transaction
+	Submit() error
+	// BlockNumber is the block number of the header this batch contains
+	BlockNumber() string
+	// RollbackOnFailure rolls back the batch transaction if the error is not nil
+	RollbackOnFailure(error)
 }
 
 // Config used to configure different underlying implementations

indexer/ipld/eth_parser.go

@@ -22,23 +22,17 @@ import (
 
 // FromBlockAndReceipts takes a block and processes it
 // to return it a set of IPLD nodes for further processing.
-func FromBlockAndReceipts(block *types.Block, receipts []*types.Receipt) (*EthHeader, []*EthTx, []*EthReceipt, [][]*EthLog, error) {
+func FromBlockAndReceipts(block *types.Block, receipts []*types.Receipt) ([]*EthTx, []*EthReceipt, [][]*EthLog, error) {
-	// Process the header
-	headerNode, err := NewEthHeader(block.Header())
-	if err != nil {
-		return nil, nil, nil, nil, err
-	}
-
 	// Process the txs
 	txNodes, err := processTransactions(block.Transactions())
 	if err != nil {
-		return nil, nil, nil, nil, err
+		return nil, nil, nil, err
 	}
 
 	// Process the receipts and logs
 	rctNodes, logNodes, err := processReceiptsAndLogs(receipts)
 
-	return headerNode, txNodes, rctNodes, logNodes, err
+	return txNodes, rctNodes, logNodes, err
 }
 
 // processTransactions will take the found transactions in a parsed block body

indexer/ipld/eth_parser_test.go

@@ -92,7 +92,7 @@ func loadBlockData(t *testing.T) []testCase {
 func TestFromBlockAndReceipts(t *testing.T) {
 	testCases := loadBlockData(t)
 	for _, tc := range testCases {
-		_, _, _, _, err := FromBlockAndReceipts(tc.block, tc.receipts)
+		_, _, _, err := FromBlockAndReceipts(tc.block, tc.receipts)
 		if err != nil {
 			t.Fatalf("error generating IPLDs from block and receipts, err %v, kind %s, block hash %s", err, tc.kind, tc.block.Hash())
 		}

indexer/reexport.go

@@ -0,0 +1,7 @@
+package indexer
+
+import "github.com/cerc-io/plugeth-statediff/indexer/interfaces"
+
+type Indexer = interfaces.StateDiffIndexer
+type Batch = interfaces.Batch
+type Config = interfaces.Config

indexer/shared/schema/table.go

@@ -45,6 +45,7 @@ type Column struct {
 	Type  colType
 	Array bool
 }
+
 type Table struct {
 	Name    string
 	Columns []Column
@@ -117,6 +118,20 @@ func (tbl *Table) ToInsertStatement(upsert bool) string {
 	)
 }
 
+// TableName returns a pgx-compatible table name.
+func (tbl *Table) TableName() []string {
+	return strings.Split(tbl.Name, ".")
+}
+
+// ColumnNames returns the ordered list of column names.
+func (tbl *Table) ColumnNames() []string {
+	var names []string
+	for _, col := range tbl.Columns {
+		names = append(names, col.Name)
+	}
+	return names
+}
+
 func sprintf(f string) colfmt {
 	return func(x interface{}) string { return fmt.Sprintf(f, x) }
 }

indexer/shared/schema/table_test.go

@@ -43,7 +43,8 @@ var testHeaderTable = Table{
 		"mh_key",
 		"times_validated",
 		"coinbase",
-	)}
+	),
+}
 
 func TestTable(t *testing.T) {
 	headerUpsert := `INSERT INTO eth.header_cids (block_number, block_hash, parent_hash, cid, td, node_id, reward, state_root, tx_root, receipt_root, uncle_root, bloom, timestamp, mh_key, times_validated, coinbase) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16) ON CONFLICT (block_hash, block_number) DO UPDATE SET (parent_hash, cid, td, node_id, reward, state_root, tx_root, receipt_root, uncle_root, bloom, timestamp, mh_key, times_validated, coinbase) = ($3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)`

indexer/test/test.go

@@ -28,7 +28,6 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
-	"github.com/cerc-io/plugeth-statediff/indexer/database/file"
 	"github.com/cerc-io/plugeth-statediff/indexer/database/sql"
 	"github.com/cerc-io/plugeth-statediff/indexer/interfaces"
 	"github.com/cerc-io/plugeth-statediff/indexer/mocks"
@@ -48,7 +47,7 @@ func SetupTestData(t *testing.T, ind interfaces.StateDiffIndexer) {
 		t.Fatal(err)
 	}
 	defer func() {
-		if err := tx.Submit(err); err != nil {
+		if err := tx.Submit(); err != nil {
 			t.Fatal(err)
 		}
 	}()
@@ -61,11 +60,7 @@ func SetupTestData(t *testing.T, ind interfaces.StateDiffIndexer) {
 		require.NoError(t, err)
 	}
 
-	if batchTx, ok := tx.(*sql.BatchTx); ok {
+	require.Equal(t, mocks.BlockNumber.String(), tx.BlockNumber())
-		require.Equal(t, mocks.BlockNumber.String(), batchTx.BlockNumber)
-	} else if batchTx, ok := tx.(*file.BatchTx); ok {
-		require.Equal(t, mocks.BlockNumber.String(), batchTx.BlockNumber)
-	}
 }
 
 func DoTestPublishAndIndexHeaderIPLDs(t *testing.T, db sql.Database) {
@@ -547,13 +542,9 @@ func SetupTestDataNonCanonical(t *testing.T, ind interfaces.StateDiffIndexer) {
 		require.NoError(t, err)
 	}
 
-	if batchTx, ok := tx1.(*sql.BatchTx); ok {
+	require.Equal(t, mocks.BlockNumber.String(), tx1.BlockNumber())
-		require.Equal(t, mocks.BlockNumber.String(), batchTx.BlockNumber)
-	} else if batchTx, ok := tx1.(*file.BatchTx); ok {
-		require.Equal(t, mocks.BlockNumber.String(), batchTx.BlockNumber)
-	}
 
-	if err := tx1.Submit(err); err != nil {
+	if err := tx1.Submit(); err != nil {
 		t.Fatal(err)
 	}
 
@@ -572,13 +563,9 @@ func SetupTestDataNonCanonical(t *testing.T, ind interfaces.StateDiffIndexer) {
 		require.NoError(t, err)
 	}
 
-	if tx, ok := tx2.(*sql.BatchTx); ok {
+	require.Equal(t, mocks.BlockNumber.String(), tx2.BlockNumber())
-		require.Equal(t, mocks.BlockNumber.String(), tx.BlockNumber)
-	} else if tx, ok := tx2.(*sql.BatchTx); ok {
-		require.Equal(t, mocks.BlockNumber.String(), tx.BlockNumber)
-	}
 
-	if err := tx2.Submit(err); err != nil {
+	if err := tx2.Submit(); err != nil {
 		t.Fatal(err)
 	}
 
@@ -597,13 +584,9 @@ func SetupTestDataNonCanonical(t *testing.T, ind interfaces.StateDiffIndexer) {
 		require.NoError(t, err)
 	}
 
-	if batchTx, ok := tx3.(*sql.BatchTx); ok {
+	require.Equal(t, mocks.Block2Number.String(), tx3.BlockNumber())
-		require.Equal(t, mocks.Block2Number.String(), batchTx.BlockNumber)
-	} else if batchTx, ok := tx3.(*file.BatchTx); ok {
-		require.Equal(t, mocks.Block2Number.String(), batchTx.BlockNumber)
-	}
 
-	if err := tx3.Submit(err); err != nil {
+	if err := tx3.Submit(); err != nil {
 		t.Fatal(err)
 	}
 }

indexer/test/test_legacy.go

@@ -20,11 +20,11 @@ import (
 	"context"
 	"testing"
 
-	"github.com/cerc-io/plugeth-statediff/indexer/database/file"
 	"github.com/cerc-io/plugeth-statediff/indexer/database/sql"
 	"github.com/cerc-io/plugeth-statediff/indexer/interfaces"
 	"github.com/cerc-io/plugeth-statediff/indexer/ipld"
 	"github.com/cerc-io/plugeth-statediff/indexer/mocks"
+	"github.com/cerc-io/plugeth-statediff/indexer/node"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/params"
 	"github.com/ipfs/go-cid"
@@ -37,6 +37,14 @@ var (
 	legacyData      = mocks.NewLegacyData(LegacyConfig)
 	mockLegacyBlock *types.Block
 	legacyHeaderCID cid.Cid
+	// Mainnet node info
+	LegacyNodeInfo = node.Info{
+		GenesisBlock: "0xd4e56740f876aef8c010b86a40d5f56745a118d0906a34e69aec8c0db1cb8fa3",
+		NetworkID:    "1",
+		ChainID:      1,
+		ID:           "mockNodeID",
+		ClientName:   "go-ethereum",
+	}
 )
 
 func SetupLegacyTestData(t *testing.T, ind interfaces.StateDiffIndexer) {
@@ -51,7 +59,7 @@ func SetupLegacyTestData(t *testing.T, ind interfaces.StateDiffIndexer) {
 	require.NoError(t, err)
 
 	defer func() {
-		if err := tx.Submit(err); err != nil {
+		if err := tx.Submit(); err != nil {
 			t.Fatal(err)
 		}
 	}()
@@ -60,11 +68,7 @@ func SetupLegacyTestData(t *testing.T, ind interfaces.StateDiffIndexer) {
 		require.NoError(t, err)
 	}
 
-	if batchTx, ok := tx.(*sql.BatchTx); ok {
+	require.Equal(t, legacyData.BlockNumber.String(), tx.BlockNumber())
-		require.Equal(t, legacyData.BlockNumber.String(), batchTx.BlockNumber)
-	} else if batchTx, ok := tx.(*file.BatchTx); ok {
-		require.Equal(t, legacyData.BlockNumber.String(), batchTx.BlockNumber)
-	}
 }
 
 func TestLegacyIndexer(t *testing.T, db sql.Database) {

indexer/test/test_mainnet.go

@@ -19,8 +19,6 @@ package test
 import (
 	"testing"
 
-	"github.com/cerc-io/plugeth-statediff/indexer/database/file"
-	"github.com/cerc-io/plugeth-statediff/indexer/database/sql"
 	"github.com/cerc-io/plugeth-statediff/indexer/interfaces"
 	"github.com/cerc-io/plugeth-statediff/indexer/mocks"
 	"github.com/ethereum/go-ethereum/core/types"
@@ -36,7 +34,7 @@ func TestBlock(t *testing.T, ind interfaces.StateDiffIndexer, testBlock *types.B
 	require.NoError(t, err)
 
 	defer func() {
-		if err := tx.Submit(err); err != nil {
+		if err := tx.Submit(); err != nil {
 			t.Fatal(err)
 		}
 	}()
@@ -45,9 +43,5 @@ func TestBlock(t *testing.T, ind interfaces.StateDiffIndexer, testBlock *types.B
 		require.NoError(t, err)
 	}
 
-	if batchTx, ok := tx.(*sql.BatchTx); ok {
+	require.Equal(t, testBlock.Number().String(), tx.BlockNumber())
-		require.Equal(t, testBlock.Number().String(), batchTx.BlockNumber)
-	} else if batchTx, ok := tx.(*file.BatchTx); ok {
-		require.Equal(t, testBlock.Number().String(), batchTx.BlockNumber)
-	}
 }

main/flags.go

@@ -46,7 +46,11 @@ func init() {
 	)
 	Flags.UintVar(&config.NumWorkers,
 		"statediff.workers", 1,
-		"Number of concurrent workers to use during statediff processing (default 1)",
+		"Number of concurrent workers to dispatch to during statediff processing",
+	)
+	Flags.UintVar(&config.SubtrieWorkers,
+		"statediff.subtries", 1,
+		"Number of subtries to iterate in parallel",
 	)
 	Flags.BoolVar(&config.WaitForSync,
 		"statediff.waitforsync", false,
@@ -169,8 +173,6 @@ func initConfig() {
 	case shared.FILE:
 		indexerConfig = fileConfig
 	case shared.POSTGRES:
-		dbConfig.ID = config.ID
-		dbConfig.ClientName = config.ClientName
 		indexerConfig = dbConfig
 	case shared.DUMP:
 		switch dbDumpDst {

main/main.go

@@ -47,8 +47,13 @@ func InitializeNode(stack core.Node, b core.Backend) {
 			ClientName:   serviceConfig.ClientName,
 		}
 		var err error
-		_, indexer, err = ind.NewStateDiffIndexer(serviceConfig.Context,
+		_, indexer, err = ind.NewStateDiffIndexer(
-			adapt.ChainConfig(backend.ChainConfig()), info, serviceConfig.IndexerConfig)
+			serviceConfig.Context,
+			adapt.ChainConfig(backend.ChainConfig()),
+			info,
+			serviceConfig.IndexerConfig,
+			true,
+		)
 		if err != nil {
 			log.Error("failed to construct indexer", "error", err)
 		}

mainnet_tests/builder_test.go

@@ -444,7 +444,7 @@ func TestBuilderOnMainnetBlocks(t *testing.T) {
 	}
 	params := statediff.Params{}
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		// note that block0 (genesis) has over 1000 nodes due to the pre-allocation for the crowd-sale
 		// it is not feasible to write a unit test of that size at this time
 		{
@@ -566,7 +566,9 @@ func TestBuilderOnMainnetBlocks(t *testing.T) {
 						},
 						StorageDiff: emptyStorage,
 					},
-					{ // this is the new account created due to the coinbase mining a block, it's creation shouldn't affect 0x 0e 05 07
+					{
+						// this is the new account created due to the coinbase mining a block, its
+						// creation shouldn't affect 0x 0e 05 07
 						Removed: false,
 						AccountWrapper: sdtypes.AccountWrapper{
 							Account: block3CoinbaseAccount,
@@ -622,11 +624,10 @@ func TestBuilderOnMainnetBlocks(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t,
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
-		chain.StateCache(),
+	test_helpers.CheckedRoots{
-		tests, params, test_helpers.CheckedRoots{
+		block1: block1RootBranchNode,
-			block1: block1RootBranchNode,
+		block2: block2RootBranchNode,
-			block2: block2RootBranchNode,
+		block3: block3RootBranchNode,
-			block3: block3RootBranchNode,
+	}.Check(t)
-		})
 }

service.go

@@ -166,10 +166,12 @@ func NewService(cfg Config, blockChain BlockChain, backend plugeth.Backend, inde
 		workers = 1
 	}
 
+	builder := NewBuilder(blockChain.StateCache())
+	builder.SetSubtrieWorkers(cfg.SubtrieWorkers)
 	quitCh := make(chan bool)
 	sds := &Service{
 		BlockChain:              blockChain,
-		Builder:                 NewBuilder(blockChain.StateCache()),
+		Builder:                 builder,
 		QuitChan:                quitCh,
 		Subscriptions:           make(map[common.Hash]map[SubID]Subscription),
 		SubscriptionTypes:       make(map[common.Hash]Params),
@@ -785,6 +787,10 @@ func (sds *Service) claimExclusiveAccess(block *types.Block) (bool, func()) {
 
 // Writes a state diff from the current block, parent state root, and provided params
 func (sds *Service) writeStateDiff(block *types.Block, parentRoot common.Hash, params Params) error {
+	if sds.indexer == nil {
+		return fmt.Errorf("indexer is not set; cannot write indexed diffs")
+	}
+
 	log := log.New("hash", block.Hash(), "number", block.Number())
 	if granted, relinquish := sds.claimExclusiveAccess(block); granted {
 		defer relinquish()
@@ -804,9 +810,6 @@ func (sds *Service) writeStateDiff(block *types.Block, parentRoot common.Hash, p
 
 	start := countStateDiffBegin(block, log)
 	defer countStateDiffEnd(start, log, &err)
-	if sds.indexer == nil {
-		return fmt.Errorf("indexer is not set; cannot write indexed diffs")
-	}
 
 	if params.IncludeTD {
 		totalDifficulty = sds.BlockChain.GetTd(block.Hash(), block.NumberU64())
@@ -814,19 +817,26 @@ func (sds *Service) writeStateDiff(block *types.Block, parentRoot common.Hash, p
 	if params.IncludeReceipts {
 		receipts = sds.BlockChain.GetReceiptsByHash(block.Hash())
 	}
+
+	t := time.Now()
 	tx, err = sds.indexer.PushBlock(block, receipts, totalDifficulty)
 	if err != nil {
 		return err
 	}
+	defer tx.RollbackOnFailure(err)
 
+	// TODO: review/remove the need to sync here
+	var nodeMtx, ipldMtx sync.Mutex
 	nodeSink := func(node types2.StateLeafNode) error {
-		defer metrics.ReportAndUpdateDuration("statediff output", time.Now(), log,
+		defer metrics.UpdateDuration(time.Now(), metrics.IndexerMetrics.OutputTimer)
-			metrics.IndexerMetrics.OutputTimer)
+		nodeMtx.Lock()
+		defer nodeMtx.Unlock()
 		return sds.indexer.PushStateNode(tx, node, block.Hash().String())
 	}
 	ipldSink := func(c types2.IPLD) error {
-		defer metrics.ReportAndUpdateDuration("statediff ipldOutput", time.Now(), log,
+		defer metrics.UpdateDuration(time.Now(), metrics.IndexerMetrics.IPLDOutputTimer)
-			metrics.IndexerMetrics.IPLDOutputTimer)
+		ipldMtx.Lock()
+		defer ipldMtx.Unlock()
 		return sds.indexer.PushIPLD(tx, c)
 	}
 
@@ -836,9 +846,12 @@ func (sds *Service) writeStateDiff(block *types.Block, parentRoot common.Hash, p
 		BlockHash:    block.Hash(),
 		BlockNumber:  block.Number(),
 	}, params, nodeSink, ipldSink)
+	if err != nil {
+		return err
+	}
 
-	// TODO this anti-pattern needs to be sorted out eventually
+	metrics.IndexerMetrics.StateStoreCodeProcessingTimer.Update(time.Since(t))
-	if err = tx.Submit(err); err != nil {
+	if err = tx.Submit(); err != nil {
 		return fmt.Errorf("batch transaction submission failed: %w", err)
 	}
 	return nil

test_helpers/builder.go

@@ -2,46 +2,147 @@ package test_helpers
 
 import (
 	"bytes"
+	"fmt"
+	"math/big"
+	"math/rand"
+	"path/filepath"
 	"sort"
+	"sync"
 	"testing"
 
+	"github.com/cerc-io/eth-iterator-utils/tracker"
 	statediff "github.com/cerc-io/plugeth-statediff"
 	"github.com/cerc-io/plugeth-statediff/adapt"
 	sdtypes "github.com/cerc-io/plugeth-statediff/types"
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/stretchr/testify/require"
 )
 
-type TestCase struct {
+var subtrieCounts = []uint{1, 8, 32}
+
+type DiffTestCase struct {
 	Name     string
 	Args     statediff.Args
 	Expected *sdtypes.StateObject
 }
 
-type CheckedRoots = map[*types.Block][]byte
+type SnapshotTestCase struct {
+	Name      string
+	StateRoot common.Hash
+	Expected  *sdtypes.StateObject
+}
 
-func RunBuilderTests(
+type CheckedRoots map[*types.Block][]byte
+
+// Replicates the statediff object, but indexes nodes by CID
+type normalizedStateDiff struct {
+	BlockNumber *big.Int
+	BlockHash   common.Hash
+	Nodes       map[string]sdtypes.StateLeafNode
+	IPLDs       map[string]sdtypes.IPLD
+}
+
+func RunBuildStateDiff(
 	t *testing.T,
 	sdb state.Database,
-	tests []TestCase,
+	tests []DiffTestCase,
 	params statediff.Params,
-	roots CheckedRoots,
 ) {
 	builder := statediff.NewBuilder(adapt.GethStateView(sdb))
 	for _, test := range tests {
-		t.Run(test.Name, func(t *testing.T) {
+		for _, subtries := range subtrieCounts {
-			diff, err := builder.BuildStateDiffObject(test.Args, params)
+			t.Run(fmt.Sprintf("%s with %d subtries", test.Name, subtries), func(t *testing.T) {
-			if err != nil {
+				builder.SetSubtrieWorkers(subtries)
-				t.Error(err)
+				diff, err := builder.BuildStateDiffObject(test.Args, params)
-			}
+				if err != nil {
+					t.Error(err)
+				}
+				require.Equal(t,
+					normalize(test.Expected),
+					normalize(&diff),
+				)
+			})
+		}
+	}
+}
 
-			normalize(test.Expected)
+func RunStateSnapshot(
-			normalize(&diff)
+	t *testing.T,
-			require.Equal(t, *test.Expected, diff)
+	sdb state.Database,
+	test SnapshotTestCase,
+	params statediff.Params,
+) {
+	builder := statediff.NewBuilder(adapt.GethStateView(sdb))
+
+	for _, subtries := range subtrieCounts {
+		// Skip the recovery test for empty diffs
+		doRecovery := len(test.Expected.Nodes) != 0
+
+		t.Run(fmt.Sprintf("%s with %d subtries", test.Name, subtries), func(t *testing.T) {
+			builder.SetSubtrieWorkers(subtries)
+			var stateNodes []sdtypes.StateLeafNode
+			var iplds []sdtypes.IPLD
+			interrupt := randomInterrupt(len(test.Expected.IPLDs))
+			stateAppender := failingSyncedAppender(&stateNodes, -1)
+			ipldAppender := failingSyncedAppender(&iplds, interrupt)
+			recoveryFile := filepath.Join(t.TempDir(), "recovery.txt")
+			build := func() error {
+				tr := tracker.New(recoveryFile, subtries)
+				defer tr.CloseAndSave()
+				return builder.WriteStateSnapshot(
+					test.StateRoot, params, stateAppender, ipldAppender, tr,
+				)
+			}
+			if doRecovery {
+				// First attempt fails, second succeeds
+				if build() == nil {
+					t.Fatal("expected an error")
+				}
+				require.FileExists(t, recoveryFile)
+			}
+			ipldAppender = failingSyncedAppender(&iplds, -1)
+			if err := build(); err != nil {
+				t.Fatal(err)
+			}
+			diff := sdtypes.StateObject{
+				Nodes: stateNodes,
+				IPLDs: iplds,
+			}
+			require.Equal(t,
+				normalize(test.Expected),
+				normalize(&diff),
+			)
 		})
 	}
+
+}
+
+// an appender which fails on a configured trigger
+func failingSyncedAppender[T any](to *[]T, failAt int) func(T) error {
+	var mtx sync.Mutex
+	return func(item T) error {
+		mtx.Lock()
+		defer mtx.Unlock()
+		if len(*to) == failAt {
+			return fmt.Errorf("failing at %d items", failAt)
+		}
+		*to = append(*to, item)
+		return nil
+	}
+}
+
+// function to pick random int between N/4 and 3N/4
+func randomInterrupt(N int) int {
+	if N < 2 {
+		return 0
+	}
+	return rand.Intn(N/2) + N/4
+}
+
+func (roots CheckedRoots) Check(t *testing.T) {
 	// Let's also confirm that our root state nodes form the state root hash in the headers
 	for block, node := range roots {
 		require.Equal(t, block.Root(), crypto.Keccak256Hash(node),
@@ -49,17 +150,13 @@ func RunBuilderTests(
 	}
 }
 
-// Sorts contained state nodes, storage nodes, and IPLDs
+func normalize(diff *sdtypes.StateObject) normalizedStateDiff {
-func normalize(diff *sdtypes.StateObject) {
+	norm := normalizedStateDiff{
-	sort.Slice(diff.IPLDs, func(i, j int) bool {
+		BlockNumber: diff.BlockNumber,
-		return diff.IPLDs[i].CID < diff.IPLDs[j].CID
+		BlockHash:   diff.BlockHash,
-	})
+		Nodes:       make(map[string]sdtypes.StateLeafNode),
-	sort.Slice(diff.Nodes, func(i, j int) bool {
+		IPLDs:       make(map[string]sdtypes.IPLD),
-		return bytes.Compare(
+	}
-			diff.Nodes[i].AccountWrapper.LeafKey,
-			diff.Nodes[j].AccountWrapper.LeafKey,
-		) < 0
-	})
 	for _, node := range diff.Nodes {
 		sort.Slice(node.StorageDiff, func(i, j int) bool {
 			return bytes.Compare(
@@ -67,5 +164,10 @@ func normalize(diff *sdtypes.StateObject) {
 				node.StorageDiff[j].LeafKey,
 			) < 0
 		})
+		norm.Nodes[node.AccountWrapper.CID] = node
 	}
+	for _, ipld := range diff.IPLDs {
+		norm.IPLDs[ipld.CID] = ipld
+	}
+	return norm
 }

test_helpers/mocks/indexer.go

@@ -17,6 +17,7 @@
 package mocks
 
 import (
+	context "context"
 	"math/big"
 	"time"
 
@@ -52,6 +53,10 @@ func (sdi *StateDiffIndexer) PushBlock(block *types.Block, receipts types.Receip
 	return &batch{}, nil
 }
 
+func (sdi *StateDiffIndexer) PushHeader(batch interfaces.Batch, header *types.Header, reward, td *big.Int) (string, error) {
+	return "", nil
+}
+
 func (sdi *StateDiffIndexer) PushStateNode(txi interfaces.Batch, stateNode sdtypes.StateLeafNode, headerID string) error {
 	return nil
 }
@@ -80,10 +85,21 @@ func (sdi *StateDiffIndexer) ClearWatchedAddresses() error {
 	return nil
 }
 
+func (sdi *StateDiffIndexer) BeginTx(number *big.Int, ctx context.Context) interfaces.Batch {
+	return &batch{}
+}
+
 func (sdi *StateDiffIndexer) Close() error {
 	return nil
 }
 
-func (tx *batch) Submit(err error) error {
+func (tx *batch) RollbackOnFailure(error) {}
+
+func (tx *batch) Submit() error {
 	return nil
 }
+
+// batch.BlockNumber
+func (tx *batch) BlockNumber() string {
+	return "0"
+}

types/types.go

@@ -53,7 +53,7 @@ type StateLeafNode struct {
 	StorageDiff    []StorageLeafNode
 }
 
-// StorageLeafNode holds the data for a single storage diff node leaf node
+// StorageLeafNode holds the data for a single storage diff leaf node
 type StorageLeafNode struct {
 	Removed bool
 	Value   []byte

utils/iterator.go

@@ -7,24 +7,9 @@ import (
 	"github.com/ethereum/go-ethereum/trie"
 )
 
-type symmDiffIterator struct {
+type SymmDiffIterator struct {
-	a, b        iterState // Nodes returned are those in b - a and a - b (keys only)
+	a, b iterState // Nodes returned are those in b - a and a - b (keys only)
-	yieldFromA  bool      // Whether next node comes from a
+	SymmDiffState
-	count       int       // Number of nodes scanned on either trie
-	eqPathIndex int       // Count index of last pair of equal paths, to detect an updated key
-}
-
-// NewSymmetricDifferenceIterator constructs a trie.NodeIterator that iterates over the symmetric difference
-// of elements in a and b, i.e., the elements in a that are not in b, and vice versa.
-// Returns the iterator, and a pointer to an integer recording the number of nodes seen.
-func NewSymmetricDifferenceIterator(a, b trie.NodeIterator) (*symmDiffIterator, *int) {
-	it := &symmDiffIterator{
-		a: iterState{a, true},
-		b: iterState{b, true},
-		// common paths are detected by a distance <=1 from this index, so put it out of reach
-		eqPathIndex: -2,
-	}
-	return it, &it.count
 }
 
 // pairs an iterator with a cache of its valid status
@@ -33,66 +18,93 @@ type iterState struct {
 	valid bool
 }
 
+// SymmDiffState exposes state specific to symmetric difference iteration, which is not accessible
+// from the NodeIterator interface. This includes the number of nodes seen, whether the current key
+// is common to both A and B, and whether the current node is sourced from A or B.
+type SymmDiffState struct {
+	yieldFromA  bool // Whether next node comes from a
+	count       int  // Number of nodes scanned on either trie
+	eqPathIndex int  // Count index of last pair of equal paths, to detect an updated key
+}
+
+// NewSymmetricDifferenceIterator constructs a trie.NodeIterator that iterates over the symmetric difference
+// of elements in a and b, i.e., the elements in a that are not in b, and vice versa.
+// Returns the iterator, and a pointer to an auxiliary object for accessing the state not exposed by the NodeIterator interface recording the number of nodes seen.
+func NewSymmetricDifferenceIterator(a, b trie.NodeIterator) *SymmDiffIterator {
+	it := &SymmDiffIterator{
+		a: iterState{a, true},
+		b: iterState{b, true},
+		// common paths are detected by a distance <=1 between count and this index, so we start at -2
+		SymmDiffState: SymmDiffState{eqPathIndex: -2},
+	}
+	return it
+}
+
 func (st *iterState) Next(descend bool) bool {
 	st.valid = st.NodeIterator.Next(descend)
 	return st.valid
 }
 
-func (it *symmDiffIterator) curr() *iterState {
+// FromA returns true if the current node is sourced from A.
+func (it *SymmDiffState) FromA() bool {
+	return it.yieldFromA
+}
+
+// CommonPath returns true if a node with the current path exists in each sub-iterator - i.e. it
+// represents an updated node.
+func (it *SymmDiffState) CommonPath() bool {
+	return it.count-it.eqPathIndex <= 1
+}
+
+// Count returns the number of nodes seen.
+func (it *SymmDiffState) Count() int {
+	return it.count
+}
+
+func (it *SymmDiffIterator) curr() *iterState {
 	if it.yieldFromA {
 		return &it.a
 	}
 	return &it.b
 }
 
-// FromA returns true if the current node is sourced from A.
+func (it *SymmDiffIterator) Hash() common.Hash {
-func (it *symmDiffIterator) FromA() bool {
-	return it.yieldFromA
-}
-
-// CommonPath returns true if a node with the current path exists in each sub-iterator - i.e. it
-// represents an updated node.
-func (it *symmDiffIterator) CommonPath() bool {
-	return it.count-it.eqPathIndex <= 1
-}
-
-func (it *symmDiffIterator) Hash() common.Hash {
 	return it.curr().Hash()
 }
 
-func (it *symmDiffIterator) Parent() common.Hash {
+func (it *SymmDiffIterator) Parent() common.Hash {
 	return it.curr().Parent()
 }
 
-func (it *symmDiffIterator) Leaf() bool {
+func (it *SymmDiffIterator) Leaf() bool {
 	return it.curr().Leaf()
 }
 
-func (it *symmDiffIterator) LeafKey() []byte {
+func (it *SymmDiffIterator) LeafKey() []byte {
 	return it.curr().LeafKey()
 }
 
-func (it *symmDiffIterator) LeafBlob() []byte {
+func (it *SymmDiffIterator) LeafBlob() []byte {
 	return it.curr().LeafBlob()
 }
 
-func (it *symmDiffIterator) LeafProof() [][]byte {
+func (it *SymmDiffIterator) LeafProof() [][]byte {
 	return it.curr().LeafProof()
 }
 
-func (it *symmDiffIterator) Path() []byte {
+func (it *SymmDiffIterator) Path() []byte {
 	return it.curr().Path()
 }
 
-func (it *symmDiffIterator) NodeBlob() []byte {
+func (it *SymmDiffIterator) NodeBlob() []byte {
 	return it.curr().NodeBlob()
 }
 
-func (it *symmDiffIterator) AddResolver(resolver trie.NodeResolver) {
+func (it *SymmDiffIterator) AddResolver(resolver trie.NodeResolver) {
 	panic("not implemented")
 }
 
-func (it *symmDiffIterator) Next(bool) bool {
+func (it *SymmDiffIterator) Next(bool) bool {
 	// NodeIterators start in a "pre-valid" state, so the first Next advances to a valid node.
 	if it.count == 0 {
 		if it.a.Next(true) {
@@ -110,7 +122,7 @@ func (it *symmDiffIterator) Next(bool) bool {
 	return it.a.valid || it.b.valid
 }
 
-func (it *symmDiffIterator) seek() {
+func (it *SymmDiffIterator) seek() {
 	// Invariants:
 	// - At the end of the function, the sub-iterator with the lexically lesser path
 	// points to the next element
@@ -151,7 +163,7 @@ func (it *symmDiffIterator) seek() {
 	}
 }
 
-func (it *symmDiffIterator) Error() error {
+func (it *SymmDiffIterator) Error() error {
 	if err := it.a.Error(); err != nil {
 		return err
 	}
@@ -172,3 +184,9 @@ func compareNodes(a, b trie.NodeIterator) int {
 	}
 	return 0
 }
+
+// AlwaysBState returns a dummy SymmDiffState that indicates all elements are from B, and have no
+// common paths with A. This is equivalent to a diff against an empty A.
+func AlwaysBState() SymmDiffState {
+	return SymmDiffState{yieldFromA: false, eqPathIndex: -2}
+}

utils/iterator_test.go

@@ -45,36 +45,33 @@ func TestSymmetricDifferenceIterator(t *testing.T) {
 	t.Run("with no difference", func(t *testing.T) {
 		db := trie.NewDatabase(rawdb.NewMemoryDatabase())
 		triea := trie.NewEmpty(db)
-		di, count := utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), triea.NodeIterator(nil))
+		di := utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), triea.NodeIterator(nil))
 		for di.Next(true) {
 			t.Errorf("iterator should not yield any elements")
 		}
-		assert.Equal(t, 0, *count)
+		assert.Equal(t, 0, di.Count())
 
 		triea.MustUpdate([]byte("foo"), []byte("bar"))
-		di, count = utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), triea.NodeIterator(nil))
+		di = utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), triea.NodeIterator(nil))
 		for di.Next(true) {
 			t.Errorf("iterator should not yield any elements")
 		}
-		assert.Equal(t, 2, *count)
+		// two nodes visited: the leaf (value) and its parent
+		assert.Equal(t, 2, di.Count())
 
 		trieb := trie.NewEmpty(db)
-		di, count = utils.NewSymmetricDifferenceIterator(
+		di = utils.NewSymmetricDifferenceIterator(triea.NodeIterator([]byte("jars")), trieb.NodeIterator(nil))
-			triea.NodeIterator([]byte("jars")),
-			trieb.NodeIterator(nil))
 		for di.Next(true) {
-			t.Errorf("iterator should not yield any elements, but got key %s", di.Path())
+			t.Errorf("iterator should not yield any elements")
 		}
-		assert.Equal(t, 0, *count)
+		assert.Equal(t, 0, di.Count())
 
-		// // TODO will fail until merged: https://github.com/ethereum/go-ethereum/pull/27838
+		// TODO will fail until merged: https://github.com/ethereum/go-ethereum/pull/27838
-		// di, count = utils.NewSymmetricDifferenceIterator(
+		// di, aux = utils.NewSymmetricDifferenceIterator(triea.NodeIterator([]byte("food")), trieb.NodeIterator(nil))
-		// 	triea.NodeIterator([]byte("food")),
-		// 	trieb.NodeIterator(nil))
 		// for di.Next(true) {
-		// 	t.Errorf("iterator should not yield any elements, but got key %s", di.Path())
+		// 	t.Errorf("iterator should not yield any elements")
 		// }
-		// assert.Equal(t, 0, *count)
+		// assert.Equal(t, 0, di.Count())
 	})
 
 	t.Run("small difference", func(t *testing.T) {
@@ -85,7 +82,7 @@ func TestSymmetricDifferenceIterator(t *testing.T) {
 		trieb := trie.NewEmpty(dbb)
 		trieb.MustUpdate([]byte("foo"), []byte("bar"))
 
-		di, count := utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), trieb.NodeIterator(nil))
+		di := utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), trieb.NodeIterator(nil))
 		leaves := 0
 		for di.Next(true) {
 			if di.Leaf() {
@@ -96,10 +93,10 @@ func TestSymmetricDifferenceIterator(t *testing.T) {
 			}
 		}
 		assert.Equal(t, 1, leaves)
-		assert.Equal(t, 2, *count)
+		assert.Equal(t, 2, di.Count())
 
 		trieb.MustUpdate([]byte("quux"), []byte("bars"))
-		di, count = utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), trieb.NodeIterator([]byte("quux")))
+		di = utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), trieb.NodeIterator([]byte("quux")))
 		leaves = 0
 		for di.Next(true) {
 			if di.Leaf() {
@@ -110,7 +107,7 @@ func TestSymmetricDifferenceIterator(t *testing.T) {
 			}
 		}
 		assert.Equal(t, 1, leaves)
-		assert.Equal(t, 1, *count)
+		assert.Equal(t, 1, di.Count())
 	})
 
 	dba := trie.NewDatabase(rawdb.NewMemoryDatabase())
@@ -127,7 +124,7 @@ func TestSymmetricDifferenceIterator(t *testing.T) {
 	onlyA := make(map[string]string)
 	onlyB := make(map[string]string)
 	var deletions, creations []string
-	it, _ := utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), trieb.NodeIterator(nil))
+	it := utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), trieb.NodeIterator(nil))
 	for it.Next(true) {
 		if !it.Leaf() {
 			continue
@@ -208,7 +205,7 @@ func TestCompareDifferenceIterators(t *testing.T) {
 		pathsA = append(pathsA, itAonly.Path())
 	}
 
-	itSym, _ := utils.NewSymmetricDifferenceIterator(treeA.NodeIterator(nil), treeB.NodeIterator(nil))
+	itSym := utils.NewSymmetricDifferenceIterator(treeA.NodeIterator(nil), treeB.NodeIterator(nil))
 	var idxA, idxB int
 	for itSym.Next(true) {
 		if itSym.FromA() {

utils/log/log.go

@@ -17,6 +17,7 @@ func init() {
 	// The plugeth logger is only initialized with the geth runtime,
 	// but tests expect to have a logger available, so default to this.
 	DefaultLogger = TestLogger
+	TestLogger.SetLevel(int(log15.LvlInfo))
 }
 
 func Trace(m string, a ...interface{}) { DefaultLogger.Trace(m, a...) }