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

     1  package nsqd
     2  
     3  import (
     4  	"encoding/binary"
     5  	"fmt"
     6  	"io"
     7  	"net"
     8  	"time"
     9  
    10  	"github.com/nsqio/go-nsq"
    11  	"github.com/nsqio/nsq/internal/lg"
    12  )
    13  
    14  // lookupPeer is a low-level type for connecting/reading/writing to nsqlookupd
    15  //
    16  // A lookupPeer instance is designed to connect lazily to nsqlookupd and reconnect
    17  // gracefully (i.e. it is all handled by the library).  Clients can simply use the
    18  // Command interface to perform a round-trip.
    19  type lookupPeer struct {
    20  	logf            lg.AppLogFunc
    21  	addr            string
    22  	conn            net.Conn
    23  	state           int32
    24  	connectCallback func(*lookupPeer)
    25  	maxBodySize     int64
    26  	Info            peerInfo
    27  }
    28  
    29  // peerInfo contains metadata for a lookupPeer instance (and is JSON marshalable)
    30  type peerInfo struct {
    31  	TCPPort          int    `json:"tcp_port"`
    32  	HTTPPort         int    `json:"http_port"`
    33  	Version          string `json:"version"`
    34  	BroadcastAddress string `json:"broadcast_address"`
    35  }
    36  
    37  // newLookupPeer creates a new lookupPeer instance connecting to the supplied address.
    38  //
    39  // The supplied connectCallback will be called *every* time the instance connects.
    40  func newLookupPeer(addr string, maxBodySize int64, l lg.AppLogFunc, connectCallback func(*lookupPeer)) *lookupPeer {
    41  	return &lookupPeer{
    42  		logf:            l,
    43  		addr:            addr,
    44  		state:           stateDisconnected,
    45  		maxBodySize:     maxBodySize,
    46  		connectCallback: connectCallback,
    47  	}
    48  }
    49  
    50  // Connect will Dial the specified address, with timeouts
    51  func (lp *lookupPeer) Connect() error {
    52  	lp.logf(lg.INFO, "LOOKUP connecting to %s", lp.addr)
    53  	conn, err := net.DialTimeout("tcp", lp.addr, time.Second)
    54  	if err != nil {
    55  		return err
    56  	}
    57  	lp.conn = conn
    58  	return nil
    59  }
    60  
    61  // String returns the specified address
    62  func (lp *lookupPeer) String() string {
    63  	return lp.addr
    64  }
    65  
    66  // Read implements the io.Reader interface, adding deadlines
    67  func (lp *lookupPeer) Read(data []byte) (int, error) {
    68  	lp.conn.SetReadDeadline(time.Now().Add(time.Second))
    69  	return lp.conn.Read(data)
    70  }
    71  
    72  // Write implements the io.Writer interface, adding deadlines
    73  func (lp *lookupPeer) Write(data []byte) (int, error) {
    74  	lp.conn.SetWriteDeadline(time.Now().Add(time.Second))
    75  	return lp.conn.Write(data)
    76  }
    77  
    78  // Close implements the io.Closer interface
    79  func (lp *lookupPeer) Close() error {
    80  	lp.state = stateDisconnected
    81  	if lp.conn != nil {
    82  		return lp.conn.Close()
    83  	}
    84  	return nil
    85  }
    86  
    87  // Command performs a round-trip for the specified Command.
    88  //
    89  // It will lazily connect to nsqlookupd and gracefully handle
    90  // reconnecting in the event of a failure.
    91  //
    92  // It returns the response from nsqlookupd as []byte
    93  func (lp *lookupPeer) Command(cmd *nsq.Command) ([]byte, error) {
    94  	initialState := lp.state
    95  	if lp.state != stateConnected {
    96  		err := lp.Connect()
    97  		if err != nil {
    98  			return nil, err
    99  		}
   100  		lp.state = stateConnected
   101  		_, err = lp.Write(nsq.MagicV1)
   102  		if err != nil {
   103  			lp.Close()
   104  			return nil, err
   105  		}
   106  		if initialState == stateDisconnected {
   107  			lp.connectCallback(lp)
   108  		}
   109  		if lp.state != stateConnected {
   110  			return nil, fmt.Errorf("lookupPeer connectCallback() failed")
   111  		}
   112  	}
   113  	if cmd == nil {
   114  		return nil, nil
   115  	}
   116  	_, err := cmd.WriteTo(lp)
   117  	if err != nil {
   118  		lp.Close()
   119  		return nil, err
   120  	}
   121  	resp, err := readResponseBounded(lp, lp.maxBodySize)
   122  	if err != nil {
   123  		lp.Close()
   124  		return nil, err
   125  	}
   126  	return resp, nil
   127  }
   128  
   129  func readResponseBounded(r io.Reader, limit int64) ([]byte, error) {
   130  	var msgSize int32
   131  
   132  	// message size
   133  	err := binary.Read(r, binary.BigEndian, &msgSize)
   134  	if err != nil {
   135  		return nil, err
   136  	}
   137  
   138  	if int64(msgSize) > limit {
   139  		return nil, fmt.Errorf("response body size (%d) is greater than limit (%d)",
   140  			msgSize, limit)
   141  	}
   142  
   143  	// message binary data
   144  	buf := make([]byte, msgSize)
   145  	_, err = io.ReadFull(r, buf)
   146  	if err != nil {
   147  		return nil, err
   148  	}
   149  
   150  	return buf, nil
   151  }