From c880435c920f53c015d9012b487e86a512bcda1b Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Mon, 16 Mar 2020 12:00:22 +0300
Subject: [PATCH] cli: persist restored blocks on exit

Add grace context to execute defer's on exit.
---
 cli/server/dump.go   |  4 ----
 cli/server/server.go | 17 +++++++++++++++--
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/cli/server/dump.go b/cli/server/dump.go
index 0e38fa74e..dbba61963 100644
--- a/cli/server/dump.go
+++ b/cli/server/dump.go
@@ -111,10 +111,6 @@ func (d *dump) add(index uint32, batch *storage.MemBatch) {
 }
 
 func (d *dump) tryPersist(prefix string, index uint32) error {
-	if index%1000 != 0 {
-		return nil
-	}
-
 	f, err := createFile(prefix, index)
 	if err != nil {
 		return err
diff --git a/cli/server/server.go b/cli/server/server.go
index 259a57692..e3f1a5f3b 100644
--- a/cli/server/server.go
+++ b/cli/server/server.go
@@ -277,9 +277,19 @@ func restoreDB(ctx *cli.Context) error {
 		}
 	}
 
+	gctx := newGraceContext()
+	var lastIndex uint32
 	dump := newDump()
+	defer func() {
+		_ = dump.tryPersist(dumpDir, lastIndex)
+	}()
 
 	for ; i < skip+count; i++ {
+		select {
+		case <-gctx.Done():
+			return cli.NewExitError("cancelled", 1)
+		default:
+		}
 		bytes, err := readBlock(reader)
 		block := &block.Block{}
 		newReader := io.NewBinReaderFromBuf(bytes)
@@ -302,8 +312,11 @@ func restoreDB(ctx *cli.Context) error {
 		if dumpDir != "" {
 			batch := chain.LastBatch()
 			dump.add(block.Index, batch)
-			if err := dump.tryPersist(dumpDir, block.Index); err != nil {
-				return cli.NewExitError(fmt.Errorf("can't dump storage to file: %v", err), 1)
+			lastIndex = block.Index
+			if block.Index%1000 == 0 {
+				if err := dump.tryPersist(dumpDir, block.Index); err != nil {
+					return cli.NewExitError(fmt.Errorf("can't dump storage to file: %v", err), 1)
+				}
 			}
 		}
 	}