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

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