fix the issue sever cannot be stopped by control-c (#518)

iotexproject · Feb 6, 2019 · b480fe1 · b480fe1
1 parent 2a0264d
commit b480fe1
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/server/itx/server.go b/server/itx/server.go
@@ -43,6 +43,7 @@ type Server struct {
 	mainChainProtocol    *mainchain.Protocol
 	initializedSubChains map[uint32]bool
 	mutex                sync.RWMutex
+	subModuleCancel      context.CancelFunc
 }
 
 // NewServer creates a new server
@@ -111,18 +112,20 @@ func newServer(cfg config.Config, testing bool) (*Server, error) {
 
 // Start starts the server
 func (s *Server) Start(ctx context.Context) error {
-	if err := s.p2pAgent.Start(ctx); err != nil {
+	cctx, cancel := context.WithCancel(context.Background())
+	s.subModuleCancel = cancel
+	if err := s.p2pAgent.Start(cctx); err != nil {
 		return errors.Wrap(err, "error when starting P2P agent")
 	}
 	if err := s.rootChainService.Blockchain().AddSubscriber(s); err != nil {
 		return errors.Wrap(err, "error when starting sub-chain starter")
 	}
 	for _, cs := range s.chainservices {
-		if err := cs.Start(ctx); err != nil {
+		if err := cs.Start(cctx); err != nil {
 			return errors.Wrap(err, "error when starting blockchain")
 		}
 	}
-	if err := s.dispatcher.Start(ctx); err != nil {
+	if err := s.dispatcher.Start(cctx); err != nil {
 		return errors.Wrap(err, "error when starting dispatcher")
 	}
 
@@ -131,6 +134,7 @@ func (s *Server) Start(ctx context.Context) error {
 
 // Stop stops the server
 func (s *Server) Stop(ctx context.Context) error {
+	defer s.subModuleCancel()
 	if err := s.p2pAgent.Stop(ctx); err != nil {
 		return errors.Wrap(err, "error when stopping P2P agent")
 	}

diff --git a/server/main.go b/server/main.go
@@ -19,7 +19,6 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
-	"time"
 
 	_ "net/http/pprof"
 
@@ -52,6 +51,7 @@ func main() {
 	signal.Notify(stop, syscall.SIGTERM)
 	ctx, cancel := context.WithCancel(context.Background())
 	stopped := make(chan struct{})
+	livenessCtx, livenessCancel := context.WithCancel(context.Background())
 
 	cfg, err := config.New()
 	if err != nil {
@@ -71,11 +71,10 @@ func main() {
 		<-stopped
 
 		// liveness end
-		pctx, lc := context.WithTimeout(context.Background(), time.Second*3)
-		defer lc()
-		if err := probeSvr.Stop(pctx); err != nil {
+		if err := probeSvr.Stop(livenessCtx); err != nil {
 			log.L().Error("Error when stopping probe server.", zap.Error(err))
 		}
+		livenessCancel()
 	}()
 
 	// create and start the node
@@ -96,6 +95,7 @@ func main() {
 
 	itx.StartServer(ctx, svr, probeSvr, cfg)
 	close(stopped)
+	<-livenessCtx.Done()
 }
 
 func initLogger(cfg config.Config) {