github.com/nsqio/nsq@v1.3.0/nsqlookupd/tcp.go (about)

     1  package nsqlookupd
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  	"sync"
     7  
     8  	"github.com/nsqio/nsq/internal/protocol"
     9  )
    10  
    11  type tcpServer struct {
    12  	nsqlookupd *NSQLookupd
    13  	conns      sync.Map
    14  }
    15  
    16  func (p *tcpServer) Handle(conn net.Conn) {
    17  	p.nsqlookupd.logf(LOG_INFO, "TCP: new client(%s)", conn.RemoteAddr())
    18  
    19  	// The client should initialize itself by sending a 4 byte sequence indicating
    20  	// the version of the protocol that it intends to communicate, this will allow us
    21  	// to gracefully upgrade the protocol away from text/line oriented to whatever...
    22  	buf := make([]byte, 4)
    23  	_, err := io.ReadFull(conn, buf)
    24  	if err != nil {
    25  		p.nsqlookupd.logf(LOG_ERROR, "failed to read protocol version - %s", err)
    26  		conn.Close()
    27  		return
    28  	}
    29  	protocolMagic := string(buf)
    30  
    31  	p.nsqlookupd.logf(LOG_INFO, "CLIENT(%s): desired protocol magic '%s'",
    32  		conn.RemoteAddr(), protocolMagic)
    33  
    34  	var prot protocol.Protocol
    35  	switch protocolMagic {
    36  	case "  V1":
    37  		prot = &LookupProtocolV1{nsqlookupd: p.nsqlookupd}
    38  	default:
    39  		protocol.SendResponse(conn, []byte("E_BAD_PROTOCOL"))
    40  		conn.Close()
    41  		p.nsqlookupd.logf(LOG_ERROR, "client(%s) bad protocol magic '%s'",
    42  			conn.RemoteAddr(), protocolMagic)
    43  		return
    44  	}
    45  
    46  	client := prot.NewClient(conn)
    47  	p.conns.Store(conn.RemoteAddr(), client)
    48  
    49  	err = prot.IOLoop(client)
    50  	if err != nil {
    51  		p.nsqlookupd.logf(LOG_ERROR, "client(%s) - %s", conn.RemoteAddr(), err)
    52  	}
    53  
    54  	p.conns.Delete(conn.RemoteAddr())
    55  	client.Close()
    56  }
    57  
    58  func (p *tcpServer) Close() {
    59  	p.conns.Range(func(k, v interface{}) bool {
    60  		v.(protocol.Client).Close()
    61  		return true
    62  	})
    63  }