btcd,netsync: add stopatheight flag and logic

When the flag stopatheight is set to something greater than zero,
btcd stops immediately after processing the height provided.

Also, we won't let the node start if the current height is equal
or greater than the provided stopatheight.

Author: allocz

btcd.go

@@ -265,6 +265,10 @@ func btcdMain(serverChan chan<- *server) error {
 	// Create server and start it.
 	server, err := newServer(cfg.Listeners, cfg.AgentBlacklist,
 		cfg.AgentWhitelist, db, activeNetParams.Params, interrupt)
+	// Don't start the server if we reached stopHeight
+	if err == errStopHeightReached {
+		return nil
+	}
 	if err != nil {
 		// TODO: this logging could do with some beautifying.
 		btcdLog.Errorf("Unable to start server on %v: %v",

config.go

@@ -172,6 +172,7 @@ type config struct {
 	SigNet               bool          `long:"signet" description:"Use the signet test network"`
 	SigNetChallenge      string        `long:"signetchallenge" description:"Connect to a custom signet network defined by this challenge instead of using the global default signet test network -- Can be specified multiple times"`
 	SigNetSeedNode       []string      `long:"signetseednode" description:"Specify a seed node for the signet network instead of using the global default signet network seed nodes"`
+	StopAtHeight         int           `long:"stopatheight" description:"Stop immediately after the block at the specified height is processed"`
 	TestNet3             bool          `long:"testnet" description:"Use the test network (version 3)"`
 	TestNet4             bool          `long:"testnet4" description:"Use the test network (version 4)"`
 	TorIsolation         bool          `long:"torisolation" description:"Enable Tor stream isolation by randomizing user credentials for each connection."`

netsync/interface.go

@@ -33,6 +33,7 @@ type Config struct {
 	Chain        *blockchain.BlockChain
 	TxMemPool    *mempool.TxPool
 	ChainParams  *chaincfg.Params
+	StopHeight   int32
 
 	DisableCheckpoints bool
 	MaxPeers           int

netsync/manager.go

@@ -6,8 +6,10 @@ package netsync
 
 import (
 	"math/rand"
+	"os"
 	"sync"
 	"sync/atomic"
+	"syscall"
 	"time"
 
 	"github.com/btcsuite/btcd/blockchain"
@@ -183,6 +185,7 @@ type SyncManager struct {
 	chain          *blockchain.BlockChain
 	txMemPool      *mempool.TxPool
 	chainParams    *chaincfg.Params
+	stopHeight     int32
 	progressLogger *blockProgressLogger
 	msgChan        chan interface{}
 	wg             sync.WaitGroup
@@ -715,6 +718,11 @@ func (sm *SyncManager) handleBlockMsg(bmsg *blockMsg) {
 		}
 	}
 
+	// Do not start processing the block if we are shutting down
+	if atomic.LoadInt32(&sm.shutdown) > 0 {
+		return
+	}
+
 	// Check if the block is eligible for less validation since the headers
 	// have already been verified to link together and are valid up to the
 	// next checkpoint.
@@ -826,6 +834,35 @@ func (sm *SyncManager) handleBlockMsg(bmsg *blockMsg) {
 		}
 	}
 
+	// Avoid processing the next block if we are at or after stopHeight
+	if sm.stopHeight > 0 && heightUpdate >= sm.stopHeight {
+		shutdown := atomic.LoadInt32(&sm.shutdown)
+		if shutdown < 1 {
+			atomic.AddInt32(&sm.shutdown, 1)
+		}
+
+		log.Infof("Reached stop height, stopping.")
+
+		// btcd will stop before processing any new block, so let's
+		// write all the cache coins to disk now
+		sm.chain.FlushUtxoCache(blockchain.FlushRequired)
+
+		// TODO(allocz): if we just call sm.Stop(), we will get stuck in
+		// a state were the process still runing waiting indefinitely
+		// for other goroutines to stop, so maybe a global stop function
+		// would be useful to register from where the stop call came
+		// from and also notify all goroutines to finish their job as
+		// soon as possible.
+		pid := os.Getpid()
+		proc, err := os.FindProcess(pid)
+		if err != nil {
+			log.Errorf("Error finding current process: %v",
+				err)
+			return
+		}
+		proc.Signal(syscall.SIGTERM)
+	}
+
 	// If we are not in the initial block download mode, it's a good time to
 	// periodically flush the blockchain cache because we don't expect new
 	// blocks immediately.  After that, there is nothing more to do.
@@ -1628,6 +1665,7 @@ func New(config *Config) (*SyncManager, error) {
 		chain:           config.Chain,
 		txMemPool:       config.TxMemPool,
 		chainParams:     config.ChainParams,
+		stopHeight:      config.StopHeight,
 		rejectedTxns:    make(map[chainhash.Hash]struct{}),
 		requestedTxns:   make(map[chainhash.Hash]struct{}),
 		requestedBlocks: make(map[chainhash.Hash]struct{}),

server.go

@@ -60,6 +60,12 @@ const (
 	connectionRetryInterval = time.Second * 5
 )
 
+var (
+	// errStopHeightReached is returned from newServer function when
+	// best height is greater or equal stopAtHeight value
+	errStopHeightReached = errors.New("stop height reached")
+)
+
 var (
 	// userAgentName is the user agent name and is used to help identify
 	// ourselves to other bitcoin peers.
@@ -296,7 +302,7 @@ type serverPeer struct {
 	addressesMtx   sync.RWMutex
 	knownAddresses lru.Cache
 	banScore       connmgr.DynamicBanScore
-	quit chan struct{}
+	quit           chan struct{}
 
 	// Closed by verAckOnce when OnVerAck fires.
 	verAckCh   chan struct{}
@@ -2995,6 +3001,14 @@ func newServer(listenAddrs, agentBlacklist, agentWhitelist []string,
 	if err != nil {
 		return nil, err
 	}
+	if sh := int32(cfg.StopAtHeight); sh > 0 {
+		bsh := s.chain.BestSnapshot().Height
+		if bsh >= sh {
+			btcdLog.Infof("Stopping because stopatheight is "+
+				"%d and chainstate best height is %d", sh, bsh)
+			return nil, errStopHeightReached
+		}
+	}
 
 	// Search for a FeeEstimator state in the database. If none can be found
 	// or if it cannot be loaded, create a new one.
@@ -3058,6 +3072,7 @@ func newServer(listenAddrs, agentBlacklist, agentWhitelist []string,
 		Chain:              s.chain,
 		TxMemPool:          s.txMemPool,
 		ChainParams:        s.chainParams,
+		StopHeight:         int32(cfg.StopAtHeight),
 		DisableCheckpoints: cfg.DisableCheckpoints,
 		MaxPeers:           cfg.MaxPeers,
 		FeeEstimator:       s.feeEstimator,