Make sure of closing trie storage when server closes (0xPolygon#172)

blockchain/storage/leveldb/leveldb_test.go

@@ -19,6 +19,9 @@ func newStorage(t *testing.T) (storage.Storage, func()) {
 		t.Fatal(err)
 	}
 	close := func() {
+		if err := s.Close(); err != nil {
+			t.Fatal(err)
+		}
 		if err := os.RemoveAll(path); err != nil {
 			t.Fatal(err)
 		}

server/server.go

@@ -33,9 +33,10 @@ import (
 
 // Minimal is the central manager of the blockchain client
 type Server struct {
-	logger hclog.Logger
-	config *Config
-	state  state.State
+	logger       hclog.Logger
+	config       *Config
+	state        state.State
+	stateStorage itrie.Storage
 
 	consensus consensus.Consensus
 
@@ -101,6 +102,7 @@ func NewServer(logger hclog.Logger, config *Config) (*Server, error) {
 	if err != nil {
 		return nil, err
 	}
+	m.stateStorage = stateStorage
 
 	st := itrie.NewState(stateStorage)
 	m.state = st
@@ -391,6 +393,11 @@ func (s *Server) Close() {
 	if err := s.consensus.Close(); err != nil {
 		s.logger.Error("failed to close consensus", "err", err.Error())
 	}
+
+	// Close the state storage
+	if err := s.stateStorage.Close(); err != nil {
+		s.logger.Error("failed to close storage for trie", "err", err.Error())
+	}
 }
 
 // Entry is a backend configuration entry

state/immutable-trie/storage.go

@@ -29,6 +29,8 @@ type Storage interface {
 	Batch() Batch
 	SetCode(hash types.Hash, code []byte)
 	GetCode(hash types.Hash) ([]byte, bool)
+
+	Close() error
 }
 
 // KVStorage is a k/v storage on memory using leveldb
@@ -78,6 +80,10 @@ func (kv *KVStorage) Get(k []byte) ([]byte, bool) {
 	return data, true
 }
 
+func (kv *KVStorage) Close() error {
+	return kv.db.Close()
+}
+
 func NewLevelDBStorage(path string, logger hclog.Logger) (Storage, error) {
 	db, err := leveldb.OpenFile(path, nil)
 	if err != nil {
@@ -127,6 +133,10 @@ func (m *memStorage) Batch() Batch {
 	return &memBatch{db: &m.db}
 }
 
+func (m *memStorage) Close() error {
+	return nil
+}
+
 func (m *memBatch) Put(p, v []byte) {
 	buf := make([]byte, len(v))
 	copy(buf[:], v[:])