Merge pull request #1123 from fjl/lean-blockchain-commands

cmd/geth: leaner blockchain commands

2 files changed, 108 insertions, 67 deletions

diff --git a/cmd/utils/cmd.go b/cmd/utils/cmd.go
index 39b4e46da..e5413973d 100644
--- a/cmd/utils/cmd.go
+++ b/cmd/utils/cmd.go
@@ -40,6 +40,10 @@ import (
 	"github.com/peterh/liner"
 )
 
+const (
+	importBatchSize = 2500
+)
+
 var interruptCallbacks = []func(os.Signal){}
 
 // Register interrupt handlers callbacks
@@ -125,10 +129,17 @@ func initDataDir(Datadir string) {
 	}
 }
 
-// Fatalf formats a message to standard output and exits the program.
+// Fatalf formats a message to standard error and exits the program.
+// The message is also printed to standard output if standard error
+// is redirected to a different file.
 func Fatalf(format string, args ...interface{}) {
-	fmt.Fprintf(os.Stderr, "Fatal: "+format+"\n", args...)
-	fmt.Fprintf(os.Stdout, "Fatal: "+format+"\n", args...)
+	w := io.MultiWriter(os.Stdout, os.Stderr)
+	outf, _ := os.Stdout.Stat()
+	errf, _ := os.Stderr.Stat()
+	if outf != nil && errf != nil && os.SameFile(outf, errf) {
+		w = os.Stderr
+	}
+	fmt.Fprintf(w, "Fatal: "+format+"\n", args...)
 	logger.Flush()
 	os.Exit(1)
 }
@@ -166,53 +177,86 @@ func FormatTransactionData(data string) []byte {
 	return d
 }
 
-func ImportChain(chainmgr *core.ChainManager, fn string) error {
-	fmt.Printf("importing blockchain '%s'\n", fn)
-	fh, err := os.OpenFile(fn, os.O_RDONLY, os.ModePerm)
+func ImportChain(chain *core.ChainManager, fn string) error {
+	// Watch for Ctrl-C while the import is running.
+	// If a signal is received, the import will stop at the next batch.
+	interrupt := make(chan os.Signal, 1)
+	stop := make(chan struct{})
+	signal.Notify(interrupt, os.Interrupt)
+	defer signal.Stop(interrupt)
+	defer close(interrupt)
+	go func() {
+		if _, ok := <-interrupt; ok {
+			glog.Info("caught interrupt during import, will stop at next batch")
+		}
+		close(stop)
+	}()
+	checkInterrupt := func() bool {
+		select {
+		case <-stop:
+			return true
+		default:
+			return false
+		}
+	}
+
+	glog.Infoln("Importing blockchain", fn)
+	fh, err := os.Open(fn)
 	if err != nil {
 		return err
 	}
 	defer fh.Close()
-
-	chainmgr.Reset()
 	stream := rlp.NewStream(fh, 0)
-	var i, n int
-
-	batchSize := 2500
-	blocks := make(types.Blocks, batchSize)
 
-	for ; ; i++ {
-		var b types.Block
-		if err := stream.Decode(&b); err == io.EOF {
-			break
-		} else if err != nil {
-			return fmt.Errorf("at block %d: %v", i, err)
+	// Run actual the import.
+	blocks := make(types.Blocks, importBatchSize)
+	n := 0
+	for batch := 0; ; batch++ {
+		// Load a batch of RLP blocks.
+		if checkInterrupt() {
+			return fmt.Errorf("interrupted")
 		}
-
-		blocks[n] = &b
-		n++
-
-		if n == batchSize {
-			if _, err := chainmgr.InsertChain(blocks); err != nil {
-				return fmt.Errorf("invalid block %v", err)
+		i := 0
+		for ; i < importBatchSize; i++ {
+			var b types.Block
+			if err := stream.Decode(&b); err == io.EOF {
+				break
+			} else if err != nil {
+				return fmt.Errorf("at block %d: %v", n, err)
 			}
-			n = 0
-			blocks = make(types.Blocks, batchSize)
+			blocks[i] = &b
+			n++
+		}
+		if i == 0 {
+			break
+		}
+		// Import the batch.
+		if checkInterrupt() {
+			return fmt.Errorf("interrupted")
+		}
+		if hasAllBlocks(chain, blocks[:i]) {
+			glog.Infof("skipping batch %d, all blocks present [%x / %x]",
+				batch, blocks[0].Hash().Bytes()[:4], blocks[i-1].Hash().Bytes()[:4])
+			continue
+		}
+		if _, err := chain.InsertChain(blocks[:i]); err != nil {
+			return fmt.Errorf("invalid block %d: %v", n, err)
 		}
 	}
+	return nil
+}
 
-	if n > 0 {
-		if _, err := chainmgr.InsertChain(blocks[:n]); err != nil {
-			return fmt.Errorf("invalid block %v", err)
+func hasAllBlocks(chain *core.ChainManager, bs []*types.Block) bool {
+	for _, b := range bs {
+		if !chain.HasBlock(b.Hash()) {
+			return false
 		}
 	}
-
-	fmt.Printf("imported %d blocks\n", i)
-	return nil
+	return true
 }
 
 func ExportChain(chainmgr *core.ChainManager, fn string) error {
-	fmt.Printf("exporting blockchain '%s'\n", fn)
+	glog.Infoln("Exporting blockchain to", fn)
 	fh, err := os.OpenFile(fn, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, os.ModePerm)
 	if err != nil {
 		return err
@@ -221,6 +265,6 @@ func ExportChain(chainmgr *core.ChainManager, fn string) error {
 	if err := chainmgr.Export(fh); err != nil {
 		return err
 	}
-	fmt.Printf("exported blockchain\n")
+	glog.Infoln("Exported blockchain to", fn)
 	return nil
 }
diff --git a/cmd/utils/flags.go b/cmd/utils/flags.go
index 155110ddc..d319055b1 100644
--- a/cmd/utils/flags.go
+++ b/cmd/utils/flags.go
@@ -256,7 +256,8 @@ var (
 	}
 )
 
-func GetNAT(ctx *cli.Context) nat.Interface {
+// MakeNAT creates a port mapper from set command line flags.
+func MakeNAT(ctx *cli.Context) nat.Interface {
 	natif, err := nat.Parse(ctx.GlobalString(NATFlag.Name))
 	if err != nil {
 		Fatalf("Option %s: %v", NATFlag.Name, err)
@@ -264,7 +265,8 @@ func GetNAT(ctx *cli.Context) nat.Interface {
 	return natif
 }
 
-func GetNodeKey(ctx *cli.Context) (key *ecdsa.PrivateKey) {
+// MakeNodeKey creates a node key from set command line flags.
+func MakeNodeKey(ctx *cli.Context) (key *ecdsa.PrivateKey) {
 	hex, file := ctx.GlobalString(NodeKeyHexFlag.Name), ctx.GlobalString(NodeKeyFileFlag.Name)
 	var err error
 	switch {
@@ -282,21 +284,12 @@ func GetNodeKey(ctx *cli.Context) (key *ecdsa.PrivateKey) {
 	return key
 }
 
+// MakeEthConfig creates ethereum options from set command line flags.
 func MakeEthConfig(clientID, version string, ctx *cli.Context) *eth.Config {
-	// Set verbosity on glog
-	glog.SetV(ctx.GlobalInt(VerbosityFlag.Name))
-	glog.CopyStandardLogTo("INFO")
-	// Set the log type
-	//glog.SetToStderr(ctx.GlobalBool(LogToStdErrFlag.Name))
-	glog.SetToStderr(true)
-	// Set the log dir
-	glog.SetLogDir(ctx.GlobalString(LogFileFlag.Name))
-
 	customName := ctx.GlobalString(IdentityFlag.Name)
 	if len(customName) > 0 {
 		clientID += "/" + customName
 	}
-
 	return &eth.Config{
 		Name:               common.MakeName(clientID, version),
 		DataDir:            ctx.GlobalString(DataDirFlag.Name),
@@ -309,15 +302,15 @@ func MakeEthConfig(clientID, version string, ctx *cli.Context) *eth.Config {
 		LogJSON:            ctx.GlobalString(LogJSONFlag.Name),
 		Etherbase:          ctx.GlobalString(EtherbaseFlag.Name),
 		MinerThreads:       ctx.GlobalInt(MinerThreadsFlag.Name),
-		AccountManager:     GetAccountManager(ctx),
+		AccountManager:     MakeAccountManager(ctx),
 		VmDebug:            ctx.GlobalBool(VMDebugFlag.Name),
 		MaxPeers:           ctx.GlobalInt(MaxPeersFlag.Name),
 		MaxPendingPeers:    ctx.GlobalInt(MaxPendingPeersFlag.Name),
 		Port:               ctx.GlobalString(ListenPortFlag.Name),
-		NAT:                GetNAT(ctx),
+		NAT:                MakeNAT(ctx),
 		NatSpec:            ctx.GlobalBool(NatspecEnabledFlag.Name),
 		Discovery:          !ctx.GlobalBool(NoDiscoverFlag.Name),
-		NodeKey:            GetNodeKey(ctx),
+		NodeKey:            MakeNodeKey(ctx),
 		Shh:                ctx.GlobalBool(WhisperEnabledFlag.Name),
 		Dial:               true,
 		BootNodes:          ctx.GlobalString(BootnodesFlag.Name),
@@ -327,35 +320,39 @@ func MakeEthConfig(clientID, version string, ctx *cli.Context) *eth.Config {
 	}
 }
 
-func GetChain(ctx *cli.Context) (*core.ChainManager, common.Database, common.Database) {
-	dataDir := ctx.GlobalString(DataDirFlag.Name)
+// SetupLogger configures glog from the logging-related command line flags.
+func SetupLogger(ctx *cli.Context) {
+	glog.SetV(ctx.GlobalInt(VerbosityFlag.Name))
+	glog.CopyStandardLogTo("INFO")
+	glog.SetToStderr(true)
+	glog.SetLogDir(ctx.GlobalString(LogFileFlag.Name))
+}
 
-	blockDb, err := ethdb.NewLDBDatabase(filepath.Join(dataDir, "blockchain"))
-	if err != nil {
+// MakeChain creates a chain manager from set command line flags.
+func MakeChain(ctx *cli.Context) (chain *core.ChainManager, blockDB, stateDB, extraDB common.Database) {
+	dd := ctx.GlobalString(DataDirFlag.Name)
+	var err error
+	if blockDB, err = ethdb.NewLDBDatabase(filepath.Join(dd, "blockchain")); err != nil {
 		Fatalf("Could not open database: %v", err)
 	}
-
-	stateDb, err := ethdb.NewLDBDatabase(filepath.Join(dataDir, "state"))
-	if err != nil {
+	if stateDB, err = ethdb.NewLDBDatabase(filepath.Join(dd, "state")); err != nil {
 		Fatalf("Could not open database: %v", err)
 	}
-
-	extraDb, err := ethdb.NewLDBDatabase(filepath.Join(dataDir, "extra"))
-	if err != nil {
+	if extraDB, err = ethdb.NewLDBDatabase(filepath.Join(dd, "extra")); err != nil {
 		Fatalf("Could not open database: %v", err)
 	}
 
 	eventMux := new(event.TypeMux)
 	pow := ethash.New()
-	chainManager := core.NewChainManager(blockDb, stateDb, pow, eventMux)
-	txPool := core.NewTxPool(eventMux, chainManager.State, chainManager.GasLimit)
-	blockProcessor := core.NewBlockProcessor(stateDb, extraDb, pow, txPool, chainManager, eventMux)
-	chainManager.SetProcessor(blockProcessor)
-
-	return chainManager, blockDb, stateDb
+	chain = core.NewChainManager(blockDB, stateDB, pow, eventMux)
+	txpool := core.NewTxPool(eventMux, chain.State, chain.GasLimit)
+	proc := core.NewBlockProcessor(stateDB, extraDB, pow, txpool, chain, eventMux)
+	chain.SetProcessor(proc)
+	return chain, blockDB, stateDB, extraDB
 }
 
-func GetAccountManager(ctx *cli.Context) *accounts.Manager {
+// MakeChain creates an account manager from set command line flags.
+func MakeAccountManager(ctx *cli.Context) *accounts.Manager {
 	dataDir := ctx.GlobalString(DataDirFlag.Name)
 	ks := crypto.NewKeyStorePassphrase(filepath.Join(dataDir, "keystore"))
 	return accounts.NewManager(ks)