Fix signal handling (#3)

Reviewed-on: cerc-io/eth-ipfs-state-validator#3

cmd/validateTrie.go

@@ -70,26 +70,30 @@ func validateTrie() {
 	stateRootStr := viper.GetString("validator.stateRoot")
 	storageRootStr := viper.GetString("validator.storageRoot")
 	contractAddrStr := viper.GetString("validator.address")
+
+	if stateRootStr == "" {
+		logWithCommand.Fatal("must provide a state root for state trie validation")
+	}
+	stateRoot := common.HexToHash(stateRootStr)
+
 	traversal := strings.ToLower(viper.GetString("validator.type"))
 	switch traversal {
 	case "f", "full":
-		if stateRootStr == "" {
-			logWithCommand.Fatal("must provide a state root for full state validation")
-		}
-		stateRoot := common.HexToHash(stateRootStr)
+		logWithCommand.
+			WithField("root", stateRoot).
+			Debug("Validating full state")
 		if err = v.ValidateTrie(stateRoot); err != nil {
-			logWithCommand.Fatalf("State for root %s is not complete\r\nerr: %v", stateRoot.String(), err)
+			logWithCommand.Fatalf("Validation failed: %v", err)
 		}
-		logWithCommand.Infof("State for root %s is complete", stateRoot.String())
+		logWithCommand.Infof("State for root %s is complete", stateRoot)
 	case "state":
-		if stateRootStr == "" {
-			logWithCommand.Fatal("must provide a state root for state trie validation")
-		}
-		stateRoot := common.HexToHash(stateRootStr)
+		logWithCommand.
+			WithField("root", stateRoot).
+			Debug("Validating state trie")
 		if err = v.ValidateStateTrie(stateRoot); err != nil {
-			logWithCommand.Fatalf("State trie for root %s is not complete\r\nerr: %v", stateRoot.String(), err)
+			logWithCommand.Fatalf("Validation failed: %s", err)
 		}
-		logWithCommand.Infof("State trie for root %s is complete", stateRoot.String())
+		logWithCommand.Infof("State trie for root %s is complete", stateRoot)
 	case "storage":
 		if storageRootStr == "" {
 			logWithCommand.Fatal("must provide a storage root for storage trie validation")
@@ -97,16 +101,16 @@ func validateTrie() {
 		if contractAddrStr == "" {
 			logWithCommand.Fatal("must provide a contract address for storage trie validation")
 		}
-		if stateRootStr == "" {
-			logWithCommand.Fatal("must provide a state root for state trie validation")
-		}
 		storageRoot := common.HexToHash(storageRootStr)
 		addr := common.HexToAddress(contractAddrStr)
-		stateRoot := common.HexToHash(stateRootStr)
+		logWithCommand.
+			WithField("contract", addr).
+			WithField("storage root", storageRoot).
+			Debug("Validating storage trie")
 		if err = v.ValidateStorageTrie(stateRoot, addr, storageRoot); err != nil {
-			logWithCommand.Fatalf("Storage trie for contract %s and root %s not complete\r\nerr: %v", addr.String(), storageRoot.String(), err)
+			logWithCommand.Fatalf("Validation failed", err)
 		}
-		logWithCommand.Infof("Storage trie for contract %s and root %s is complete", addr.String(), storageRoot.String())
+		logWithCommand.Infof("Storage trie for contract %s and root %s is complete", addr, storageRoot)
 	default:
 		logWithCommand.Fatalf("Invalid traversal level: '%s'", traversal)
 	}

go.mod

@@ -3,7 +3,7 @@ module github.com/cerc-io/eth-ipfs-state-validator/v5
 go 1.19
 
 require (
-	github.com/cerc-io/eth-iterator-utils v0.1.0
+	github.com/cerc-io/eth-iterator-utils v0.1.1
 	github.com/cerc-io/ipfs-ethdb/v5 v5.0.0-alpha
 	github.com/cerc-io/ipld-eth-statedb v0.0.5-alpha
 	github.com/ethereum/go-ethereum v1.11.6
@@ -247,9 +247,9 @@ require (
 )
 
 replace (
-	github.com/cerc-io/eth-iterator-utils => git.vdb.to/cerc-io/eth-iterator-utils v0.1.0
-	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/cerc-io/ipld-eth-statedb => git.vdb.to/cerc-io/ipld-eth-statedb v0.0.6-alpha
-	github.com/cerc-io/plugeth-statediff => git.vdb.to/cerc-io/plugeth-statediff v0.1.1
+	github.com/cerc-io/plugeth-statediff => git.vdb.to/cerc-io/plugeth-statediff v0.1.4
 	github.com/ethereum/go-ethereum => git.vdb.to/cerc-io/plugeth v0.0.0-20230808125822-691dc334fab1
 )

go.sum

@@ -45,15 +45,15 @@ dmitri.shuralyov.com/html/belt v0.0.0-20180602232347-f7d459c86be0/go.mod h1:JLBr
 dmitri.shuralyov.com/service/change v0.0.0-20181023043359-a85b471d5412/go.mod h1:a1inKt/atXimZ4Mv927x+r7UpyzRUf4emIoiiSC2TN4=
 dmitri.shuralyov.com/state v0.0.0-20180228185332-28bcc343414c/go.mod h1:0PRwlb0D6DFvNNtx+9ybjezNCa8XF0xaYcETyp6rHWU=
 git.apache.org/thrift.git v0.0.0-20180902110319-2566ecd5d999/go.mod h1:fPE2ZNJGynbRyZ4dJvy6G277gSllfV2HJqblrnkyeyg=
-git.vdb.to/cerc-io/eth-iterator-utils v0.1.0 h1:4YALvTld0wfp9KV+pPfPcyJfzVCAkKSEhnxtJiFiNlE=
-git.vdb.to/cerc-io/eth-iterator-utils v0.1.0/go.mod h1:oYOOUbNpcuAUoFhfTg6Zab4PyrRIk8SSJ/i34WNE6Fw=
-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-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/ipld-eth-statedb v0.0.6-alpha h1:0YnoohjuK7w2JhIqLDDyVUNnu1RjyeDqqyhm6MojD74=
 git.vdb.to/cerc-io/ipld-eth-statedb v0.0.6-alpha/go.mod h1:cCQCfIUX5vTZBHeAfLa8wLUeLKO8kygDPm7Afc+MMI8=
 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-statediff v0.1.1 h1:qKygQwZO2vECySFpgfaiXwWn275jIV//pwsuNfeCXLo=
-git.vdb.to/cerc-io/plugeth-statediff v0.1.1/go.mod h1:61WPXOqJshEkxzalwyyLukTiIBG+Z7WyfvRdhVLfBF8=
+git.vdb.to/cerc-io/plugeth-statediff v0.1.4 h1:swDJDAk1/yu6MOHAvxeyZz+MS1H9FCmSWGQRswFxFEw=
+git.vdb.to/cerc-io/plugeth-statediff v0.1.4/go.mod h1:PX1bwTwZKzegQu6zYdpuJf1tHNmU0suapKOUUO2xIzs=
 github.com/AndreasBriese/bbloom v0.0.0-20180913140656-343706a395b7/go.mod h1:bOvUY6CB00SOBii9/FifXqc0awNKxLFCL/+pkDPuyl8=
 github.com/AndreasBriese/bbloom v0.0.0-20190306092124-e2d15f34fcf9/go.mod h1:bOvUY6CB00SOBii9/FifXqc0awNKxLFCL/+pkDPuyl8=
 github.com/AndreasBriese/bbloom v0.0.0-20190825152654-46b345b51c96 h1:cTp8I5+VIoKjsnZuH8vjyaysT/ses3EvZeaV/1UkF2M=
@@ -112,6 +112,7 @@ github.com/btcsuite/btcd v0.0.0-20190605094302-a0d1e3e36d50/go.mod h1:3J08xEfcug
 github.com/btcsuite/btcd v0.0.0-20190824003749-130ea5bddde3/go.mod h1:3J08xEfcugPacsc34/LKRU2yO7YmuT8yt28J8k2+rrI=
 github.com/btcsuite/btcd v0.20.1-beta/go.mod h1:wVuoA8VJLEcwgqHBwHmzLRazpKxTv13Px/pDuV7OomQ=
 github.com/btcsuite/btcd v0.21.0-beta/go.mod h1:ZSWyehm27aAuS9bvkATT+Xte3hjHZ+MRgMY/8NJ7K94=
+github.com/btcsuite/btcd v0.22.0-beta h1:LTDpDKUM5EeOFBPM8IXpinEcmZ6FWfNZbE3lfrfdnWo=
 github.com/btcsuite/btcd/btcec/v2 v2.3.2 h1:5n0X6hX0Zk+6omWcihdYvdAlGf2DfasC0GMf7DClJ3U=
 github.com/btcsuite/btcd/btcec/v2 v2.3.2/go.mod h1:zYzJ8etWJQIv1Ogk7OzpWjowwOdXY1W/17j2MW85J04=
 github.com/btcsuite/btcd/chaincfg/chainhash v1.0.2 h1:KdUfX2zKommPRa+PD0sWZUyXe9w277ABlgELO7H04IM=

pkg/validator.go

@@ -20,6 +20,9 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	"github.com/spf13/viper"
@@ -143,7 +146,7 @@ func (v *Validator) ValidateTrie(stateRoot common.Hash) error {
 	if err != nil {
 		return err
 	}
-	iterate := func(it trie.NodeIterator) error { return v.iterate(it, true) }
+	iterate := func(ctx context.Context, it trie.NodeIterator) error { return v.iterate(ctx, it, true) }
 	return iterateTracked(t, fmt.Sprintf(v.params.RecoveryFormat, fullTraversal), v.params.Workers, iterate)
 }
 
@@ -155,7 +158,7 @@ func (v *Validator) ValidateStateTrie(stateRoot common.Hash) error {
 	if err != nil {
 		return err
 	}
-	iterate := func(it trie.NodeIterator) error { return v.iterate(it, false) }
+	iterate := func(ctx context.Context, it trie.NodeIterator) error { return v.iterate(ctx, it, false) }
 	return iterateTracked(t, fmt.Sprintf(v.params.RecoveryFormat, stateTraversal), v.params.Workers, iterate)
 }
 
@@ -167,7 +170,7 @@ func (v *Validator) ValidateStorageTrie(stateRoot common.Hash, address common.Ad
 	if err != nil {
 		return err
 	}
-	iterate := func(it trie.NodeIterator) error { return v.iterate(it, false) }
+	iterate := func(ctx context.Context, it trie.NodeIterator) error { return v.iterate(ctx, it, false) }
 	return iterateTracked(t, fmt.Sprintf(v.params.RecoveryFormat, storageTraversal), v.params.Workers, iterate)
 }
 
@@ -181,12 +184,18 @@ func (v *Validator) Close() error {
 
 // Traverses one iterator fully
 // If storage = true, also traverse storage tries for each leaf.
-func (v *Validator) iterate(it trie.NodeIterator, storage bool) error {
+func (v *Validator) iterate(ctx context.Context, it trie.NodeIterator, storage bool) error {
 	// Iterate through entire state trie. it.Next() will return false when we have
 	// either completed iteration of the entire trie or run into an error (e.g. a
 	// missing node). If we are able to iterate through the entire trie without error
 	// then the trie is complete.
 	for it.Next(true) {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+
 		// This block adapted from geth - core/state/iterator.go
 		// If storage is not requested, or the state trie node is an internal entry, skip
 		if !storage || !it.Leaf() {
@@ -219,18 +228,22 @@ func (v *Validator) iterate(it trie.NodeIterator, storage bool) error {
 
 // Traverses each iterator in a separate goroutine.
 // Dumps to a recovery file on failure or interrupt.
-func iterateTracked(tree state.Trie, recoveryFile string, iterCount uint, fn func(trie.NodeIterator) error) error {
-	ctx, cancelCtx := context.WithCancel(context.Background())
+func iterateTracked(
+	tree state.Trie,
+	recoveryFile string,
+	iterCount uint,
+	fn func(context.Context, trie.NodeIterator) error,
+) error {
 	tracker := tracker.New(recoveryFile, iterCount)
-	tracker.CaptureSignal(cancelCtx)
 	halt := func() {
-		if err := tracker.HaltAndDump(); err != nil {
+		log.Errorf("writing recovery file: %s", recoveryFile)
+		if err := tracker.CloseAndSave(); err != nil {
 			log.Errorf("failed to write recovery file: %v", err)
 		}
 	}
 
 	// attempt to restore from recovery file if it exists
-	iters, err := tracker.Restore(tree.NodeIterator)
+	iters, _, err := tracker.Restore(tree.NodeIterator)
 	if err != nil {
 		return err
 	}
@@ -243,14 +256,28 @@ func iterateTracked(tree state.Trie, recoveryFile string, iterCount uint, fn fun
 		for i, it := range iters {
 			iters[i] = tracker.Tracked(it)
 		}
+	} else {
+		log.Debugf("restored %d iterators from: %s", len(iters), recoveryFile)
 	}
 
+	ctx, cancel := context.WithCancel(context.Background())
 	g, ctx := errgroup.WithContext(ctx)
-	defer halt()
 
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		sig := <-sigChan
+		log.Errorf("Signal received (%v), stopping", sig)
+		cancel()
+	}()
+
+	defer halt()
 	for _, it := range iters {
 		func(it trie.NodeIterator) {
-			g.Go(func() error { return fn(it) })
+			g.Go(func() error {
+				return fn(ctx, it)
+			})
+
 		}(it)
 	}
 	return g.Wait()

test/compose.yml

@@ -5,7 +5,7 @@ services:
     restart: on-failure
     depends_on:
       - ipld-eth-db
-    image: git.vdb.to/cerc-io/ipld-eth-db/ipld-eth-db:v5.0.2-alpha
+    image: git.vdb.to/cerc-io/ipld-eth-db/ipld-eth-db:v5.0.5-alpha
     environment:
       DATABASE_USER: "vdbm"
       DATABASE_NAME: "cerc_testing"