Repository: nsqio/nsq
File: nsqlookupd/lookup_protocol_v1.go (Go)

Code before:
	"github.com/nsqio/nsq/internal/protocol"
	"github.com/nsqio/nsq/internal/version"
)

type LookupProtocolV1 struct {
	ctx      *Context
	exitChan chan struct{}
}

func (p *LookupProtocolV1) IOLoop(conn net.Conn) error {
	var err error
	var line string

	client := NewClientV1(conn)
	reader := bufio.NewReader(client)
	readChan := make(chan string)
	errChan := make(chan error)

	for {
		exitLoop := false

		// do this read in a goroutine so we can exit this loop if needed
		go func() {
			line, e := reader.ReadString('\n')
			readChan <- line
			errChan <- e
		}()

		// now wait until we either get an exit signal or a readline completes
		select {
		case <-p.exitChan:
			exitLoop = true
			break
		case line = <-readChan:
			err = <-errChan
			break
		}

		if exitLoop || err != nil {
			break
		}

		line = strings.TrimSpace(line)
		params := strings.Split(line, " ")

		var response []byte
		response, err = p.Exec(client, reader, params)
		if err != nil {
			ctx := ""
			if parentErr := err.(protocol.ChildErr).Parent(); parentErr != nil {
				ctx = " - " + parentErr.Error()

Review comment (suggestion):
I think it would be better to instead find a way to call `conn.Close()` when `exitChan` indicates. Here's one idea:

```go
doneChan := make(chan interface{})
defer close(doneChan)
go func() {
        select {
        case <-p.exitChan:
                conn.Close()
        case <-doneChan:
                break
        }
}
for {
        line, err = reader.ReadString('\n')
        ...
```

Here's a different way that does not involve an extra goroutine and channel for each connection:

```diff
--- a/nsqlookupd/nsqlookupd.go
+++ b/nsqlookupd/nsqlookupd.go
@@ -66,6 +66,7 @@ func (l *NSQLookupd) Main() error {
        tcpServer := &tcpServer{ctx: ctx}
        l.waitGroup.Wrap(func() {
                exitFunc(protocol.TCPServer(l.tcpListener, tcpServer, l.logf))
+               tcpServer.CloseAll()
        })
        httpServer := newHTTPServer(ctx)
        l.waitGroup.Wrap(func() {

--- a/nsqlookupd/tcp.go
+++ b/nsqlookupd/tcp.go
@@ -3,12 +3,14 @@ package nsqlookupd
 import (
        "io"
        "net"
+       "sync"
 
        "github.com/nsqio/nsq/internal/protocol"
 )
 
 type tcpServer struct {
        ctx *Context
+       conns sync.Map
 }
 
 func (p *tcpServer) Handle(clientConn net.Conn) {
@@ -41,9 +43,20 @@ 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()
+               p.conns.Delete(k)
+               return true
+       })
 }
```

but I'm not sure if the overhead of the `sync.Map()` is something to worry about. probably not :)

Code after:
	"github.com/nsqio/nsq/internal/protocol"
	"github.com/nsqio/nsq/internal/version"
)

type LookupProtocolV1 struct {
	ctx *Context
}

func (p *LookupProtocolV1) IOLoop(conn net.Conn) error {
	var err error
	var line string

	client := NewClientV1(conn)
	reader := bufio.NewReader(client)
	for {
		line, err = reader.ReadString('\n')
		if err != nil {
			break
		}

		line = strings.TrimSpace(line)
		params := strings.Split(line, " ")

		var response []byte
		response, err = p.Exec(client, reader, params)
		if err != nil {
			ctx := ""
			if parentErr := err.(protocol.ChildErr).Parent(); parentErr != nil {
				ctx = " - " + parentErr.Error()
			}
			p.ctx.nsqlookupd.logf(LOG_ERROR, "[%s] - %s%s", client, err, ctx)

			_, sendErr := protocol.SendResponse(client, []byte(err.Error()))
			if sendErr != nil {
				p.ctx.nsqlookupd.logf(LOG_ERROR, "[%s] - %s%s", client, sendErr, ctx)
				break
			}

			// errors of type FatalClientErr should forceably close the connection
			if _, ok := err.(*protocol.FatalClientErr); ok {
				break
			}
			continue
		}

		if response != nil {
			_, err = protocol.SendResponse(client, response)
			if err != nil {
				break
			}
		}