Avoid t.Log data races flagged up by the race detector

This commit picks nsqio#1190 and backports
it to this version of NSQ to avoid data races in t.Log flagged up by the
race detector

internal/protocol/tcp_server.go

@@ -4,6 +4,7 @@ import (
 	"net"
 	"runtime"
 	"strings"
+	"sync"
 
 	"github.com/nsqio/nsq/internal/lg"
 )
@@ -15,6 +16,8 @@ type TCPHandler interface {
 func TCPServer(listener net.Listener, handler TCPHandler, logf lg.AppLogFunc) {
 	logf(lg.INFO, "TCP: listening on %s", listener.Addr())
 
+	var wg sync.WaitGroup
+
 	for {
 		clientConn, err := listener.Accept()
 		if err != nil {
@@ -29,8 +32,16 @@ func TCPServer(listener net.Listener, handler TCPHandler, logf lg.AppLogFunc) {
 			}
 			break
 		}
-		go handler.Handle(clientConn)
+
+		wg.Add(1)
+		go func() {
+			handler.Handle(clientConn)
+			wg.Done()
+		}()
 	}
 
+	// wait to return until all handler goroutines complete
+	wg.Wait()
+
 	logf(lg.INFO, "TCP: closing %s", listener.Addr())
 }

nsqlookupd/lookup_protocol_v1_test.go

@@ -38,6 +38,8 @@ func testIOLoopReturnsClientErr(t *testing.T, fakeConn test.FakeNetConn) {
 
 	prot := &LookupProtocolV1{ctx: &Context{nsqlookupd: New(opts)}}
 
+	prot.ctx.nsqlookupd.tcpServer = &tcpServer{ctx: prot.ctx}
+
 	errChan := make(chan error)
 	testIOLoop := func() {
 		errChan <- prot.IOLoop(fakeConn)

nsqlookupd/nsqlookupd.go

@@ -18,6 +18,7 @@ type NSQLookupd struct {
 	opts         *Options
 	tcpListener  net.Listener
 	httpListener net.Listener
+	tcpServer    *tcpServer
 	waitGroup    util.WaitGroupWrapper
 	DB           *RegistrationDB
 }
@@ -53,9 +54,9 @@ func (l *NSQLookupd) Main() {
 	l.Lock()
 	l.tcpListener = tcpListener
 	l.Unlock()
-	tcpServer := &tcpServer{ctx: ctx}
+	l.tcpServer = &tcpServer{ctx: ctx}
 	l.waitGroup.Wrap(func() {
-		protocol.TCPServer(tcpListener, tcpServer, l.logf)
+		protocol.TCPServer(tcpListener, l.tcpServer, l.logf)
 	})
 
 	httpListener, err := net.Listen("tcp", l.opts.HTTPAddress)
@@ -89,6 +90,10 @@ func (l *NSQLookupd) Exit() {
 		l.tcpListener.Close()
 	}
 
+	if l.tcpServer != nil {
+		l.tcpServer.CloseAll()
+	}
+
 	if l.httpListener != nil {
 		l.httpListener.Close()
 	}

nsqlookupd/tcp.go

@@ -3,12 +3,14 @@ package nsqlookupd
 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.nsqlookupd.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
+	})
 }