stateroot: wait for the service to stop un Shutdown

pkg/services/stateroot/service.go

@@ -59,6 +59,7 @@ type (
 		maxRetries      int
 		relayExtensible RelayCallback
 		blockCh         chan *block.Block
+		stopCh          chan struct{}
 		done            chan struct{}
 	}
 )
@@ -79,6 +80,7 @@ func New(cfg config.StateRoot, sm *stateroot.Module, log *zap.Logger, bc Ledger,
 		log:             log,
 		incompleteRoots: make(map[uint32]*incompleteRoot),
 		blockCh:         make(chan *block.Block),
+		stopCh:          make(chan struct{}),
 		done:            make(chan struct{}),
 		timePerBlock:    time.Duration(bcConf.SecondsPerBlock) * time.Second,
 		maxRetries:      voteValidEndInc,

pkg/services/stateroot/validators.go

@@ -46,10 +46,11 @@ runloop:
 			s.srMtx.Lock()
 			delete(s.incompleteRoots, b.Index-voteValidEndInc)
 			s.srMtx.Unlock()
-		case <-s.done:
+		case <-s.stopCh:
 			break runloop
 		}
 	}
+	s.chain.UnsubscribeFromBlocks(s.blockCh)
 drainloop:
 	for {
 		select {
@@ -59,6 +60,7 @@ drainloop:
 		}
 	}
 	close(s.blockCh)
+	close(s.done)
 }
 
 // Shutdown stops the service.
@@ -66,8 +68,8 @@ func (s *service) Shutdown() {
 	if !s.started.CAS(true, false) {
 		return
 	}
-	s.chain.UnsubscribeFromBlocks(s.blockCh)
-	close(s.done)
+	close(s.stopCh)
+	<-s.done
 }
 
 func (s *service) signAndSend(r *state.MPTRoot) error {