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.

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/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
+}