Add progress counter to prometheus output (#76)

* Add a progress counter by checking the distance already traversed from the startPath to endPath in a bounded iterator vs the estimated number of iterations.

.gitignore

@@ -5,3 +5,4 @@ mocks/
 .vscode
 output_dir*/
 log_file
+recovery_file

pkg/prom/prom.go

@@ -70,6 +70,16 @@ func Init() {
 	})
 }
 
+func RegisterGaugeFunc(name string, function func() float64) {
+	promauto.NewGaugeFunc(
+		prometheus.GaugeOpts{
+			Namespace: namespace,
+			Subsystem: statsSubsystem,
+			Name:      name,
+			Help:      name,
+		}, function)
+}
+
 // RegisterDBCollector create metric collector for given connection
 func RegisterDBCollector(name string, db DBStatsGetter) {
 	if metrics {
@@ -111,3 +121,7 @@ func DecActiveIterCount() {
 		activeIteratorCount.Dec()
 	}
 }
+
+func Enabled() bool {
+	return metrics
+}

pkg/snapshot/config.go

@@ -165,6 +165,9 @@ func (c *ServiceConfig) Init() error {
 	viper.BindEnv(SNAPSHOT_END_HEIGHT_TOML, SNAPSHOT_END_HEIGHT)
 	viper.BindEnv(SNAPSHOT_MODE_TOML, SNAPSHOT_MODE)
 	viper.BindEnv(SNAPSHOT_WORKERS_TOML, SNAPSHOT_WORKERS)
+	viper.BindEnv(PROM_HTTP_TOML, PROM_HTTP)
+	viper.BindEnv(PROM_HTTP_ADDR_TOML, PROM_HTTP_ADDR)
+	viper.BindEnv(PROM_METRICS_TOML, PROM_METRICS)
 
 	viper.BindEnv(SNAPSHOT_ACCOUNTS_TOML, SNAPSHOT_ACCOUNTS)
 	var allowedAccounts []string

pkg/snapshot/service.go

@@ -202,9 +202,6 @@ func (s *Service) createSnapshot(ctx context.Context, it trie.NodeIterator, head
 	}
 	defer func() {
 		err = CommitOrRollback(tx, err)
-		if err != nil {
-			log.Errorf("CommitOrRollback failed: %s", err)
-		}
 	}()
 
 	// path (from recovery dump) to be seeked on recovery
@@ -224,28 +221,31 @@ func (s *Service) createSnapshot(ctx context.Context, it trie.NodeIterator, head
 		recoveredPath = append(recoveredPath, *seekedPath...)
 		endPath = i.endPath
 	} else {
-		return errors.New("untracked iterator")
+		err = errors.New("untracked iterator")
+		return err
 	}
 
-	return s.createSubTrieSnapshot(ctx, tx, nil, it, recoveredPath, seekedPath, endPath, headerID, height, seekingPaths)
+	tx, err = s.createSubTrieSnapshot(ctx, tx, nil, it, recoveredPath, seekedPath, endPath, headerID, height, seekingPaths)
+	return err
 }
 
 // createSubTrieSnapshot processes nodes at the next level of a trie using the given subtrie iterator
 // continually updating seekedPath with path of the latest processed node
 func (s *Service) createSubTrieSnapshot(ctx context.Context, tx Tx, prefixPath []byte, subTrieIt trie.NodeIterator,
-	recoveredPath []byte, seekedPath *[]byte, endPath []byte, headerID string, height *big.Int, seekingPaths [][]byte) error {
+	recoveredPath []byte, seekedPath *[]byte, endPath []byte, headerID string, height *big.Int, seekingPaths [][]byte) (Tx, error) {
 	prom.IncActiveIterCount()
 	defer prom.DecActiveIterCount()
 
 	// descend in the first loop iteration to reach first child node
+	var err error
 	descend := true
 	for {
 		select {
 		case <-ctx.Done():
-			return errors.New("ctx cancelled")
+			return tx, errors.New("ctx cancelled")
 		default:
 			if ok := subTrieIt.Next(descend); !ok {
-				return subTrieIt.Error()
+				return tx, subTrieIt.Error()
 			}
 
 			// to avoid descending further
@@ -257,15 +257,16 @@ func (s *Service) createSubTrieSnapshot(ctx context.Context, tx Tx, prefixPath [
 				// if node path is empty and prefix is nil, it's the root node
 				if prefixPath == nil {
 					// create snapshot of node, if it is a leaf this will also create snapshot of entire storage trie
-					if err := s.createNodeSnapshot(tx, subTrieIt, headerID, height, seekingPaths, prefixPath); err != nil {
-						return err
+					tx, err = s.createNodeSnapshot(tx, subTrieIt, headerID, height, seekingPaths, prefixPath)
+					if err != nil {
+						return tx, err
 					}
 					updateSeekedPath(seekedPath, subTrieIt.Path())
 				}
 
 				if ok := subTrieIt.Next(true); !ok {
 					// return if no further nodes available
-					return subTrieIt.Error()
+					return tx, subTrieIt.Error()
 				}
 			}
 
@@ -283,7 +284,7 @@ func (s *Service) createSubTrieSnapshot(ctx context.Context, tx Tx, prefixPath [
 				if trackedSubtrieIt, ok := subTrieIt.(*trackedIter); ok {
 					s.tracker.stopIter(trackedSubtrieIt)
 				}
-				return subTrieIt.Error()
+				return tx, subTrieIt.Error()
 			}
 
 			// skip the current node if it's before recovered path and not along the recovered path
@@ -307,8 +308,9 @@ func (s *Service) createSubTrieSnapshot(ctx context.Context, tx Tx, prefixPath [
 
 			// if the node is along paths of interest
 			// create snapshot of node, if it is a leaf this will also create snapshot of entire storage trie
-			if err := s.createNodeSnapshot(tx, subTrieIt, headerID, height, seekingPaths, prefixPath); err != nil {
-				return err
+			tx, err = s.createNodeSnapshot(tx, subTrieIt, headerID, height, seekingPaths, prefixPath)
+			if err != nil {
+				return tx, err
 			}
 			// update seeked path after node has been processed
 			updateSeekedPath(seekedPath, nodePath)
@@ -316,11 +318,12 @@ func (s *Service) createSubTrieSnapshot(ctx context.Context, tx Tx, prefixPath [
 			// create an iterator to traverse and process the next level of this subTrie
 			nextSubTrieIt, err := s.createSubTrieIt(nodePath, subTrieIt.Hash(), recoveredPath)
 			if err != nil {
-				return err
+				return tx, err
 			}
 			// pass on the seekedPath of the tracked concurrent iterator to be updated
-			if err := s.createSubTrieSnapshot(ctx, tx, nodePath, nextSubTrieIt, recoveredPath, seekedPath, endPath, headerID, height, seekingPaths); err != nil {
-				return err
+			tx, err = s.createSubTrieSnapshot(ctx, tx, nodePath, nextSubTrieIt, recoveredPath, seekedPath, endPath, headerID, height, seekingPaths)
+			if err != nil {
+				return tx, err
 			}
 		}
 	}
@@ -359,10 +362,10 @@ func (s *Service) createSubTrieIt(prefixPath []byte, hash common.Hash, recovered
 // createNodeSnapshot indexes the current node
 // entire storage trie is also indexed (if available)
 func (s *Service) createNodeSnapshot(tx Tx, it trie.NodeIterator, headerID string, height *big.Int,
-	watchedAddressesLeafPaths [][]byte, prefixPath []byte) error {
+	watchedAddressesLeafPaths [][]byte, prefixPath []byte) (Tx, error) {
 	tx, err := s.ipfsPublisher.PrepareTxForBatch(tx, s.maxBatchSize)
 	if err != nil {
-		return err
+		return tx, err
 	}
 
 	// index values by leaf key
@@ -371,25 +374,25 @@ func (s *Service) createNodeSnapshot(tx Tx, it trie.NodeIterator, headerID strin
 		// publish codehash => code mappings
 		// take storage snapshot
 		if err := s.processStateValueNode(it, headerID, height, prefixPath, watchedAddressesLeafPaths, tx); err != nil {
-			return err
+			return tx, err
 		}
 	} else { // trie nodes will be written to blockstore only
 		// reminder that this includes leaf nodes, since the geth iterator.Leaf() actually signifies a "value" node
 		// so this is also where we publish the IPLD block corresponding to the "value" nodes indexed above
 		if IsNullHash(it.Hash()) {
 			// skip null node
-			return nil
+			return tx, nil
 		}
 		nodeVal := make([]byte, len(it.NodeBlob()))
 		copy(nodeVal, it.NodeBlob())
 		if len(watchedAddressesLeafPaths) > 0 {
 			var elements []interface{}
 			if err := rlp.DecodeBytes(nodeVal, &elements); err != nil {
-				return err
+				return tx, err
 			}
 			ok, err := isLeaf(elements)
 			if err != nil {
-				return err
+				return tx, err
 			}
 			if ok {
 				// create the full node path as it.Path() doesn't include the path before subtrie root
@@ -398,18 +401,18 @@ func (s *Service) createNodeSnapshot(tx Tx, it trie.NodeIterator, headerID strin
 				valueNodePath := append(nodePath, partialPath...)
 				if !isWatchedAddress(watchedAddressesLeafPaths, valueNodePath) {
 					// skip this node
-					return nil
+					return tx, nil
 				}
 			}
 		}
 		nodeHash := make([]byte, len(it.Hash().Bytes()))
 		copy(nodeHash, it.Hash().Bytes())
 		if err := s.ipfsPublisher.PublishIPLD(ipld.Keccak256ToCid(ipld.MEthStateTrie, nodeHash), nodeVal, height, tx); err != nil {
-			return err
+			return tx, err
 		}
 	}
 
-	return it.Error()
+	return tx, it.Error()
 }
 
 // reminder: it.Leaf() == true when the iterator is positioned at a "value node" which is not something that actually exists in an MMPT
@@ -423,15 +426,6 @@ func (s *Service) processStateValueNode(it trie.NodeIterator, headerID string, h
 		return nil
 	}
 
-	// created vs updated is important for leaf nodes since we need to diff their storage
-	// so we need to map all changed accounts at B to their leafkey, since account can change pathes but not leafkey
-	var account types.StateAccount
-	accountRLP := make([]byte, len(it.LeafBlob()))
-	copy(accountRLP, it.LeafBlob())
-	if err := rlp.DecodeBytes(accountRLP, &account); err != nil {
-		return fmt.Errorf("error decoding account for leaf value at leaf key %x\nerror: %v", it.LeafKey(), err)
-	}
-
 	// since this is a "value node", we need to move up to the "parent" node which is the actual leaf node
 	// it should be in the fastcache since it necessarily was recently accessed to reach the current "node"
 	parentNodeRLP, err := s.stateDB.TrieDB().Node(it.Parent())
@@ -450,6 +444,15 @@ func (s *Service) processStateValueNode(it trie.NodeIterator, headerID string, h
 	encodedPath := trie.HexToCompact(valueNodePath)
 	leafKey := encodedPath[1:]
 
+	// created vs updated is important for leaf nodes since we need to diff their storage
+	// so we need to map all changed accounts at B to their leafkey, since account can change pathes but not leafkey
+	var account types.StateAccount
+	accountRLP := make([]byte, len(it.LeafBlob()))
+	copy(accountRLP, it.LeafBlob())
+	if err := rlp.DecodeBytes(accountRLP, &account); err != nil {
+		return fmt.Errorf("error decoding account for leaf value at leaf key %x\nerror: %v", leafKey, err)
+	}
+
 	// write codehash => code mappings if we have a contract
 	if !bytes.Equal(account.CodeHash, emptyCodeHash) {
 		codeHash := common.BytesToHash(account.CodeHash)

pkg/snapshot/service_test.go

@@ -3,6 +3,7 @@ package snapshot
 import (
 	"errors"
 	"fmt"
+	"math/big"
 	"math/rand"
 	"os"
 	"path/filepath"
@@ -11,6 +12,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/ipfs/go-cid"
+
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/statediff/indexer/models"
 	"github.com/golang/mock/gomock"
@@ -192,8 +195,6 @@ func TestAccountSelectiveSnapshot(t *testing.T) {
 			expectedStorageNodes.Store(keys, value)
 		}
 
-		var count int
-
 		pub, tx := makeMocks(t)
 		pub.EXPECT().PublishHeader(gomock.Eq(&fixt.Chain2_Block32_Header))
 		pub.EXPECT().BeginTx().Return(tx, nil).
@@ -206,7 +207,6 @@ func TestAccountSelectiveSnapshot(t *testing.T) {
 			gomock.Any(),
 			gomock.Eq(tx)).
 			Do(func(stateNode *models.StateNodeModel, _ snapt.Tx) error {
-				count++
 				if stateNode.BlockNumber != fixt.Chain2_Block32_Header.Number.String() {
 					t.Fatalf(unexpectedStateNodeErr, stateNode.StateKey)
 				}
@@ -261,7 +261,12 @@ func TestAccountSelectiveSnapshot(t *testing.T) {
 			}).
 			AnyTimes()
 		pub.EXPECT().PublishIPLD(gomock.Any(), gomock.Any(), gomock.Eq(fixt.Chain2_Block32_Header.Number), gomock.Eq(tx)).
-			AnyTimes()
+			Do(func(_ cid.Cid, _ []byte, height *big.Int, _ snapt.Tx) {
+				if height.String() != fixt.Chain2_Block32_Header.Number.String() {
+					t.Fatalf("wrong blockheight for ipld publish: %s", height.String())
+				}
+			}).
+			MaxTimes(len(fixt.Chain2_Block32_StateIPLDs) + len(fixt.Chain2_Block32_StorageIPLDs) + 1 + 2)
 
 		chainDataPath, ancientDataPath := fixt.GetChainDataPath("chain2data")
 		config := testConfig(chainDataPath, ancientDataPath)
@@ -318,6 +323,7 @@ func TestRecovery(t *testing.T) {
 		pub.EXPECT().PublishHeader(gomock.Eq(&fixt.Block1_Header))
 		pub.EXPECT().BeginTx().Return(tx, nil).MaxTimes(workers)
 		pub.EXPECT().PrepareTxForBatch(gomock.Any(), gomock.Any()).Return(tx, nil).AnyTimes()
+		tx.EXPECT().Rollback().MaxTimes(workers)
 		tx.EXPECT().Commit().MaxTimes(workers)
 		pub.EXPECT().PublishStateLeafNode(
 			gomock.Any(),
@@ -481,7 +487,8 @@ func TestAccountSelectiveRecovery(t *testing.T) {
 		pub.EXPECT().PublishHeader(gomock.Eq(&fixt.Chain2_Block32_Header))
 		pub.EXPECT().BeginTx().Return(tx, nil).Times(workers)
 		pub.EXPECT().PrepareTxForBatch(gomock.Any(), gomock.Any()).Return(tx, nil).AnyTimes()
-		tx.EXPECT().Commit().Times(workers)
+		tx.EXPECT().Commit().MaxTimes(workers)
+		tx.EXPECT().Rollback().MaxTimes(workers)
 		pub.EXPECT().PublishStateLeafNode(
 			gomock.Any(),
 			gomock.Eq(tx)).

pkg/snapshot/tracker.go

@@ -7,20 +7,45 @@ import (
 	"fmt"
 	"os"
 	"os/signal"
+	"sync"
+	"sync/atomic"
 	"syscall"
 
+	"github.com/cerc-io/ipld-eth-state-snapshot/pkg/prom"
+
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/trie"
 	iter "github.com/ethereum/go-ethereum/trie/concurrent_iterator"
 	log "github.com/sirupsen/logrus"
 )
 
+var trackedIterCount int32
+
 type trackedIter struct {
+	id   int32
+	mu   sync.Mutex
+	done atomic.Bool
 	trie.NodeIterator
 	tracker *iteratorTracker
 
-	seekedPath []byte // latest path seeked from the tracked iterator
+	seekedPath []byte // latest full node path seeked from the tracked iterator
+	startPath  []byte // startPath for the tracked iterator
 	endPath    []byte // endPath for the tracked iterator
+	lastPath   []byte // latest it.Path() (not the full node path) seeked
+}
+
+func (it *trackedIter) getLastPath() []byte {
+	it.mu.Lock()
+	defer it.mu.Unlock()
+
+	return it.lastPath
+}
+
+func (it *trackedIter) setLastPath(val []byte) {
+	it.mu.Lock()
+	defer it.mu.Unlock()
+
+	it.lastPath = val
 }
 
 func (it *trackedIter) Next(descend bool) bool {
@@ -32,6 +57,9 @@ func (it *trackedIter) Next(descend bool) bool {
 		} else {
 			log.Errorf("iterator stopped after tracker halted: path=%x", it.Path())
 		}
+		it.done.Store(true)
+	} else {
+		it.setLastPath(it.Path())
 	}
 	return ret
 }
@@ -82,12 +110,43 @@ func (tr *iteratorTracker) tracked(it trie.NodeIterator, recoveredPath []byte) (
 	// if the iterator being tracked is a PrefixBoundIterator, capture it's end path
 	// to be used in trie traversal
 	var endPath []byte
+	var startPath []byte
 	if boundedIter, ok := it.(*iter.PrefixBoundIterator); ok {
+		startPath = boundedIter.StartPath
 		endPath = boundedIter.EndPath
 	}
 
-	ret = &trackedIter{it, tr, iterSeekedPath, endPath}
+	ret = &trackedIter{
+		atomic.AddInt32(&trackedIterCount, 1),
+		sync.Mutex{},
+		atomic.Bool{},
+		it,
+		tr,
+		iterSeekedPath,
+		startPath,
+		endPath,
+		nil,
+	}
 	tr.startChan <- ret
+
+	if prom.Enabled() {
+		pathDepth := max(max(len(startPath), len(endPath)), 1)
+		totalSteps := estimateSteps(startPath, endPath, pathDepth)
+		prom.RegisterGaugeFunc(
+			fmt.Sprintf("tracked_iterator_%d", ret.id),
+			func() float64 {
+				if ret.done.Load() {
+					return 100.0
+				}
+				lastPath := ret.getLastPath()
+				if nil == lastPath {
+					return 0.0
+				}
+				remainingSteps := estimateSteps(lastPath, endPath, pathDepth)
+				return (float64(totalSteps) - float64(remainingSteps)) / float64(totalSteps) * 100.0
+			})
+	}
+
 	return
 }
 

pkg/snapshot/util.go

@@ -5,12 +5,11 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/ethereum/go-ethereum/statediff/indexer/database/sql/postgres"
-
 	"github.com/cerc-io/ipld-eth-state-snapshot/pkg/prom"
 	file "github.com/cerc-io/ipld-eth-state-snapshot/pkg/snapshot/file"
 	"github.com/cerc-io/ipld-eth-state-snapshot/pkg/snapshot/pg"
 	snapt "github.com/cerc-io/ipld-eth-state-snapshot/pkg/types"
+	"github.com/ethereum/go-ethereum/statediff/indexer/database/sql/postgres"
 )
 
 func NewPublisher(mode SnapshotMode, config *Config) (snapt.Publisher, error) {
@@ -53,6 +52,77 @@ func decrementPath(path []byte) bool {
 	return true
 }
 
+// Estimate the number of iterations necessary to step from start to end.
+func estimateSteps(start []byte, end []byte, depth int) uint64 {
+	// We see paths in several forms (nil, 0600, 06, etc.). We need to adjust them to a comparable form.
+	// For nil, start and end indicate the extremes of 0x0 and 0x10.  For differences in depth, we often see a
+	// start/end range on a bounded iterator specified like 0500:0600, while the value returned by it.Path() may
+	// be shorter, like 06.  Since our goal is to estimate how many steps it would take to move from start to end,
+	// we want to perform the comparison at a stable depth, since to move from 05 to 06 is only 1 step, but
+	// to move from 0500:06 is 16.
+	normalizePathRange := func(start []byte, end []byte, depth int) ([]byte, []byte) {
+		if 0 == len(start) {
+			start = []byte{0x0}
+		}
+		if 0 == len(end) {
+			end = []byte{0x10}
+		}
+		normalizedStart := make([]byte, depth)
+		normalizedEnd := make([]byte, depth)
+		for i := 0; i < depth; i++ {
+			if i < len(start) {
+				normalizedStart[i] = start[i]
+			}
+			if i < len(end) {
+				normalizedEnd[i] = end[i]
+			}
+		}
+		return normalizedStart, normalizedEnd
+	}
+
+	// We have no need to handle negative exponents, so uints are fine.
+	pow := func(x uint64, y uint) uint64 {
+		if 0 == y {
+			return 1
+		}
+		ret := x
+		for i := uint(0); i < y; i++ {
+			ret *= x
+		}
+		return x
+	}
+
+	// Fix the paths.
+	start, end = normalizePathRange(start, end, depth)
+
+	// No negative distances, if the start is already >= end, the distance is 0.
+	if bytes.Compare(start, end) >= 0 {
+		return 0
+	}
+
+	// Subtract each component, right to left, carrying over if necessary.
+	difference := make([]byte, len(start))
+	var carry byte = 0
+	for i := len(start) - 1; i >= 0; i-- {
+		result := end[i] - start[i] - carry
+		if result > 0xf && i > 0 {
+			result &= 0xf
+			carry = 1
+		} else {
+			carry = 0
+		}
+		difference[i] = result
+	}
+
+	// Calculate the result.
+	var ret uint64 = 0
+	for i := 0; i < len(difference); i++ {
+		ret += uint64(difference[i]) * pow(16, uint(len(difference)-i-1))
+	}
+
+	return ret
+}
+
 // https://github.com/ethereum/go-ethereum/blob/master/trie/encoding.go#L97
 func keybytesToHex(str []byte) []byte {
 	l := len(str)*2 + 1
@@ -86,3 +156,11 @@ func checkUpperPathBound(nodePath, endPath []byte) bool {
 
 	return bytes.Compare(nodePath, endPath) <= 0
 }
+
+func max(a int, b int) int {
+	if a > b {
+		return a
+	}
+
+	return b
+}