From 9541029303c3067aeef9d1be1caab00cf06f9732 Mon Sep 17 00:00:00 2001
From: Vsevolod Brekelov <vbrekelov@nspcc.ru>
Date: Mon, 16 Dec 2019 17:04:18 +0300
Subject: [PATCH] cli: add metrics to dump and restore calls

---
 cli/server/server.go | 38 +++++++++++++++++++++++++-------------
 1 file changed, 25 insertions(+), 13 deletions(-)

diff --git a/cli/server/server.go b/cli/server/server.go
index f3511ab01..d742e0d1c 100644
--- a/cli/server/server.go
+++ b/cli/server/server.go
@@ -134,6 +134,22 @@ func getCountAndSkipFromContext(ctx *cli.Context) (uint32, uint32) {
 	return count, skip
 }
 
+func initBCWithMetrics(cfg config.Config) (*core.Blockchain, *metrics.Service, *metrics.Service, error) {
+	chain, err := initBlockChain(cfg)
+	if err != nil {
+		return nil, nil, nil, cli.NewExitError(err, 1)
+	}
+	configureAddresses(cfg.ApplicationConfiguration)
+	prometheus := metrics.NewPrometheusService(cfg.ApplicationConfiguration.Prometheus)
+	pprof := metrics.NewPprofService(cfg.ApplicationConfiguration.Pprof)
+
+	go chain.Run()
+	go prometheus.Start()
+	go pprof.Start()
+
+	return chain, prometheus, pprof, nil
+}
+
 func dumpDB(ctx *cli.Context) error {
 	cfg, err := getConfigFromContext(ctx)
 	if err != nil {
@@ -154,11 +170,10 @@ func dumpDB(ctx *cli.Context) error {
 	defer outStream.Close()
 	writer := io.NewBinWriterFromIO(outStream)
 
-	chain, err := initBlockChain(cfg)
+	chain, prometheus, pprof, err := initBCWithMetrics(cfg)
 	if err != nil {
-		return cli.NewExitError(err, 1)
+		return err
 	}
-	go chain.Run()
 
 	chainHeight := chain.BlockHeight()
 	if skip+count > chainHeight {
@@ -182,9 +197,12 @@ func dumpDB(ctx *cli.Context) error {
 			return cli.NewExitError(err, 1)
 		}
 	}
+	pprof.ShutDown()
+	prometheus.ShutDown()
 	chain.Close()
 	return nil
 }
+
 func restoreDB(ctx *cli.Context) error {
 	cfg, err := getConfigFromContext(ctx)
 	if err != nil {
@@ -205,11 +223,10 @@ func restoreDB(ctx *cli.Context) error {
 	defer inStream.Close()
 	reader := io.NewBinReaderFromIO(inStream)
 
-	chain, err := initBlockChain(cfg)
+	chain, prometheus, pprof, err := initBCWithMetrics(cfg)
 	if err != nil {
 		return err
 	}
-	go chain.Run()
 
 	var allBlocks = reader.ReadU32LE()
 	if reader.Err != nil {
@@ -241,8 +258,9 @@ func restoreDB(ctx *cli.Context) error {
 			return cli.NewExitError(fmt.Errorf("failed to add block %d: %s", i, err), 1)
 		}
 	}
+	pprof.ShutDown()
+	prometheus.ShutDown()
 	chain.Close()
-
 	return nil
 }
 
@@ -271,23 +289,17 @@ func startServer(ctx *cli.Context) error {
 
 	serverConfig := network.NewServerConfig(cfg)
 
-	chain, err := initBlockChain(cfg)
+	chain, prometheus, pprof, err := initBCWithMetrics(cfg)
 	if err != nil {
 		return err
 	}
 
-	configureAddresses(cfg.ApplicationConfiguration)
 	server := network.NewServer(serverConfig, chain)
 	rpcServer := rpc.NewServer(chain, cfg.ApplicationConfiguration.RPC, server)
 	errChan := make(chan error)
-	prometheus := metrics.NewPrometheusService(cfg.ApplicationConfiguration.Prometheus)
-	pprof := metrics.NewPprofService(cfg.ApplicationConfiguration.Pprof)
 
-	go chain.Run()
 	go server.Start(errChan)
 	go rpcServer.Start(errChan)
-	go prometheus.Start()
-	go pprof.Start()
 
 	fmt.Println(logo())
 	fmt.Println(server.UserAgent)