nsqd: fix stall in Exit() due to tcp-protocol producer connections

Consumer connections are closed when topics are closed,
but tcp-protocol publish connections are not,
so add tcpServer.CloseAll().

problem introduced by #1190

nsqd/nsqd.go

@@ -61,6 +61,7 @@ type NSQD struct {
 
 	lookupPeers atomic.Value
 
+	tcpServer     *tcpServer
 	tcpListener   net.Listener
 	httpListener  net.Listener
 	httpsListener net.Listener
@@ -255,9 +256,9 @@ func (n *NSQD) Main() error {
 		})
 	}
 
-	tcpServer := &tcpServer{ctx: ctx}
+	n.tcpServer = &tcpServer{ctx: ctx}
 	n.waitGroup.Wrap(func() {
-		exitFunc(protocol.TCPServer(n.tcpListener, tcpServer, n.logf))
+		exitFunc(protocol.TCPServer(n.tcpListener, n.tcpServer, n.logf))
 	})
 	httpServer := newHTTPServer(ctx, false, n.getOpts().TLSRequired == TLSRequired)
 	n.waitGroup.Wrap(func() {
@@ -423,6 +424,9 @@ func (n *NSQD) Exit() {
 	if n.tcpListener != nil {
 		n.tcpListener.Close()
 	}
+	if n.tcpServer != nil {
+		n.tcpServer.CloseAll()
+	}
 
 	if n.httpListener != nil {
 		n.httpListener.Close()

nsqd/tcp.go

@@ -3,12 +3,14 @@ package nsqd
 import (
 	"io"
 	"net"
+	"sync"
 
 	"github.com/nsqio/nsq/internal/protocol"
 )
 
 type tcpServer struct {
-	ctx *context
+	ctx   *context
+	conns sync.Map
 }
 
 func (p *tcpServer) Handle(clientConn net.Conn) {
@@ -41,9 +43,19 @@ func (p *tcpServer) Handle(clientConn net.Conn) {
 		return
 	}
 
+	p.conns.Store(clientConn.RemoteAddr(), clientConn)
+
 	err = prot.IOLoop(clientConn)
 	if err != nil {
 		p.ctx.nsqd.logf(LOG_ERROR, "client(%s) - %s", clientConn.RemoteAddr(), err)
-		return
 	}
+
+	p.conns.Delete(clientConn.RemoteAddr())
+}
+
+func (p *tcpServer) CloseAll() {
+	p.conns.Range(func(k, v interface{}) bool {
+		v.(net.Conn).Close()
+		return true
+	})
 }