fix(zero): make zero shutdown cleanly

Author: xqqp

conn/pool.go

@@ -91,6 +91,18 @@ func (p *Pools) GetAll() []*Pool {
 	return pool
 }
 
+// RemoveAll removes all pool entries.
+func (p *Pools) RemoveAll() {
+	p.Lock()
+	defer p.Unlock()
+
+	for k, pool := range p.all {
+		glog.Warningf("CONN: Disconnecting from %s\n", k)
+		delete(p.all, k)
+		pool.shutdown()
+	}
+}
+
 // RemoveInvalid removes invalid nodes from the list of pools.
 func (p *Pools) RemoveInvalid(state *pb.MembershipState) {
 	// Keeps track of valid IP addresses, assigned to active nodes. We do this
@@ -277,11 +289,20 @@ func (p *Pool) MonitorHealth() {
 
 	// We might have lost connection to the destination. In that case, re-dial
 	// the connection.
-	reconnect := func() {
+	// Returns true, if reconnection was successful
+	reconnect := func() bool {
+		reconnectionTicker := time.NewTicker(time.Second)
+		defer reconnectionTicker.Stop()
 		for {
-			time.Sleep(time.Second)
+			select {
+			case <-p.closer.HasBeenClosed():
+				glog.Infof("CONN: Returning from MonitorHealth for %s", p.Addr)
+				return false
+			case <-reconnectionTicker.C:
+			}
+
 			if err := p.closer.Ctx().Err(); err != nil {
-				return
+				return false
 			}
 			ctx, cancel := context.WithTimeout(p.closer.Ctx(), 10*time.Second)
 			conn, err := grpc.NewClient(p.Addr, p.dialOpts...)
@@ -298,7 +319,7 @@ func (p *Pool) MonitorHealth() {
 				}
 				p.conn = conn
 				p.Unlock()
-				return
+				return true
 			}
 			glog.Errorf("CONN: Unable to connect with %s : %s\n", p.Addr, err)
 			if conn != nil {
@@ -309,19 +330,21 @@ func (p *Pool) MonitorHealth() {
 		}
 	}
 
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
 	for {
 		select {
 		case <-p.closer.HasBeenClosed():
 			glog.Infof("CONN: Returning from MonitorHealth for %s", p.Addr)
 			return
-		default:
-			err := p.listenToHeartbeat()
-			if err != nil {
-				reconnect()
+		case <-ticker.C:
+		}
+
+		err := p.listenToHeartbeat()
+		if err != nil {
+			if reconnect() {
 				glog.Infof("CONN: Re-established connection with %s.\n", p.Addr)
 			}
-			// Sleep for a bit before retrying.
-			time.Sleep(echoDuration)
 		}
 	}
 }

dgraph/cmd/zero/raft.go

@@ -879,7 +879,7 @@ func (n *node) Run() {
 	// snapshot can cause select loop to block while deleting entries, so run
 	// it in goroutine
 	readStateCh := make(chan raft.ReadState, 100)
-	closer := z.NewCloser(5)
+	closer := z.NewCloser(4)
 	defer func() {
 		closer.SignalAndWait()
 		n.closer.Done()

dgraph/cmd/zero/run.go

@@ -346,6 +346,8 @@ func run() {
 		st.node.closer.SignalAndWait()
 		// Stop all internal requests.
 		_ = grpcListener.Close()
+		// Stop all pools
+		conn.GetPools().RemoveAll()
 	}()
 
 	st.zero.closer.AddRunning(2)