// getConn is used to get a connection from the pool.
func (n *NetworkTransport) getConn(target string) (*netConn, error) {
// Check for a pooled conn
if conn := n.getPooledConn(target); conn != nil {
return conn, nil
}
// Dial a new connection
conn, err := n.stream.Dial(target, n.timeout)
if err != nil {
return nil, err
}
// Wrap the conn
netConn := &netConn{
target: target,
conn: conn,
r: bufio.NewReader(conn),
w: bufio.NewWriter(conn),
}
// Setup encoder/decoders
netConn.dec = codec.NewDecoder(netConn.r, &codec.MsgpackHandle{})
netConn.enc = codec.NewEncoder(netConn.w, &codec.MsgpackHandle{})
// Done
return netConn, nil
}
// sendTCPUserMsg is used to send a TCP userMsg to another host
func (m *Memberlist) sendTCPUserMsg(to net.Addr, sendBuf []byte) error {
dialer := net.Dialer{Timeout: m.config.TCPTimeout}
conn, err := dialer.Dial("tcp", to.String())
if err != nil {
return err
}
defer conn.Close()
bufConn := bytes.NewBuffer(nil)
if err := bufConn.WriteByte(byte(userMsg)); err != nil {
return err
}
// Send our node state
header := userMsgHeader{UserMsgLen: len(sendBuf)}
hd := codec.MsgpackHandle{}
enc := codec.NewEncoder(bufConn, &hd)
if err := enc.Encode(&header); err != nil {
return err
}
if _, err := bufConn.Write(sendBuf); err != nil {
return err
}
return m.rawSendMsgTCP(conn, bufConn.Bytes())
}
// handleSnapshotRequest reads the request from the conn and dispatches it. This
// will be called from a goroutine after an incoming stream is determined to be
// a snapshot request.
func (s *Server) handleSnapshotRequest(conn net.Conn) error {
var args structs.SnapshotRequest
dec := codec.NewDecoder(conn, &codec.MsgpackHandle{})
if err := dec.Decode(&args); err != nil {
return fmt.Errorf("failed to decode request: %v", err)
}
var reply structs.SnapshotResponse
snap, err := s.dispatchSnapshotRequest(&args, conn, &reply)
if err != nil {
reply.Error = err.Error()
goto RESPOND
}
defer func() {
if err := snap.Close(); err != nil {
s.logger.Printf("[ERR] consul: Failed to close snapshot: %v", err)
}
}()
RESPOND:
enc := codec.NewEncoder(conn, &codec.MsgpackHandle{})
if err := enc.Encode(&reply); err != nil {
return fmt.Errorf("failed to encode response: %v", err)
}
if snap != nil {
if _, err := io.Copy(conn, snap); err != nil {
return fmt.Errorf("failed to stream snapshot: %v", err)
}
}
return nil
}
// joins the raft leader and sets up infrastructure for
// processing commands
// can return ErrNotLeader
func newConnToLeader(conn net.Conn, advertiseAddr string, lg *log.Logger) (*connToLeader, error) {
// send join command
h := &codec.MsgpackHandle{}
ret := &connToLeader{
c: conn,
e: codec.NewEncoder(conn, h),
d: codec.NewDecoder(conn, h),
l: new(sync.Mutex),
lg: lg,
pending: make(chan *commandCallback, 64),
}
join := &joinReq{
PeerAddr: advertiseAddr,
}
err := ret.e.Encode(join)
if err != nil {
ret.c.Close()
return nil, err
}
joinResp := &joinResp{}
err = ret.d.Decode(joinResp)
if err != nil {
ret.lg.Printf("Error connecting to leader at %s : %s", conn.RemoteAddr().String(), err)
ret.c.Close()
return nil, err
}
go ret.readResponses()
return ret, nil
}
func extractEvent(rec map[string]interface{}) (EventRecord, error) {
// encode to a byte stream
buf := bytes.NewBuffer(nil)
handle := codec.MsgpackHandle{RawToString: true, WriteExt: true}
encoder := codec.NewEncoder(buf, &handle)
if err := encoder.Encode(rec); err != nil {
log.Printf("encoding failed")
return nil, err
}
// decode back to appropriate struc
switch rec["Event"] {
case "member-join":
fallthrough
case "member-leave":
fallthrough
case "member-failed":
fallthrough
case "member-update":
fallthrough
case "member-reap":
return decodeMemberEventRecord(buf.Bytes())
case "user":
return decodeUserEventRecord(buf.Bytes())
case "query":
return decodeQueryEventRecord(buf.Bytes())
default:
return nil, fmt.Errorf("unhandled event type: %s", rec["Event"])
}
}
// Encode writes an encoded object to a new bytes buffer.
func encodeMsgPack(in interface{}) (*bytes.Buffer, error) {
buf := bytes.NewBuffer(nil)
hd := codec.MsgpackHandle{}
enc := codec.NewEncoder(buf, &hd)
err := enc.Encode(in)
return buf, err
}
// listen is a long running routine that listens for new clients
func (i *AgentRPC) listen() {
for {
conn, err := i.listener.Accept()
if err != nil {
if i.stop {
return
}
i.logger.Printf("[ERR] agent.rpc: Failed to accept client: %v", err)
continue
}
i.logger.Printf("[INFO] agent.rpc: Accepted client: %v", conn.RemoteAddr())
// Wrap the connection in a client
client := &rpcClient{
name: conn.RemoteAddr().String(),
conn: conn,
reader: bufio.NewReader(conn),
writer: bufio.NewWriter(conn),
}
client.dec = codec.NewDecoder(client.reader, msgpackHandle)
client.enc = codec.NewEncoder(client.writer, msgpackHandle)
// Register the client
i.Lock()
if !i.stop {
i.clients[client.name] = client
go i.handleClient(client)
} else {
conn.Close()
}
i.Unlock()
}
}
// sendLocalState is invoked to send our local state over a tcp connection
func (m *Memberlist) sendLocalState(conn net.Conn, join bool) error {
// Setup a deadline
conn.SetDeadline(time.Now().Add(m.config.TCPTimeout))
// Prepare the local node state
m.nodeLock.RLock()
localNodes := make([]pushNodeState, len(m.nodes))
for idx, n := range m.nodes {
localNodes[idx].Name = n.Name
localNodes[idx].Addr = n.Addr
localNodes[idx].Port = n.Port
localNodes[idx].Incarnation = n.Incarnation
localNodes[idx].State = n.State
localNodes[idx].Meta = n.Meta
localNodes[idx].Vsn = []uint8{
n.PMin, n.PMax, n.PCur,
n.DMin, n.DMax, n.DCur,
}
}
m.nodeLock.RUnlock()
// Get the delegate state
var userData []byte
if m.config.Delegate != nil {
userData = m.config.Delegate.LocalState(join)
}
// Create a bytes buffer writer
bufConn := bytes.NewBuffer(nil)
// Send our node state
header := pushPullHeader{Nodes: len(localNodes), UserStateLen: len(userData), Join: join}
hd := codec.MsgpackHandle{}
enc := codec.NewEncoder(bufConn, &hd)
// Begin state push
if _, err := bufConn.Write([]byte{byte(pushPullMsg)}); err != nil {
return err
}
if err := enc.Encode(&header); err != nil {
return err
}
for i := 0; i < header.Nodes; i++ {
if err := enc.Encode(&localNodes[i]); err != nil {
return err
}
}
// Write the user state as well
if userData != nil {
if _, err := bufConn.Write(userData); err != nil {
return err
}
}
// Get the send buffer
return m.rawSendMsgTCP(conn, bufConn.Bytes())
}
// encodeTags
func (t *tribe) encodeTags(tags map[string]string) []byte {
var buf bytes.Buffer
enc := codec.NewEncoder(&buf, &codec.MsgpackHandle{})
if err := enc.Encode(tags); err != nil {
panic(fmt.Sprintf("Failed to encode tags: %v", err))
}
return buf.Bytes()
}
// Encode writes an encoded object to a new bytes buffer
func encode(msgType messageType, in interface{}) (*bytes.Buffer, error) {
buf := bytes.NewBuffer(nil)
buf.WriteByte(uint8(msgType))
hd := codec.MsgpackHandle{}
enc := codec.NewEncoder(buf, &hd)
err := enc.Encode(in)
return buf, err
}
func encodeMessage(t messageType, msg interface{}) ([]byte, error) {
buf := bytes.NewBuffer(nil)
buf.WriteByte(uint8(t))
handle := codec.MsgpackHandle{}
encoder := codec.NewEncoder(buf, &handle)
err := encoder.Encode(msg)
return buf.Bytes(), err
}
func encodeFilter(f filterType, filt interface{}) ([]byte, error) {
buf := bytes.NewBuffer(nil)
buf.WriteByte(uint8(f))
handle := codec.MsgpackHandle{}
encoder := codec.NewEncoder(buf, &handle)
err := encoder.Encode(filt)
return buf.Bytes(), err
}
// NewCodec returns a MsgpackCodec that can be used as either
// a Client or Server rpc Codec. It also provides controls for
// enabling and disabling buffering for both reads and writes.
func NewCodec(bufReads, bufWrites bool, conn io.ReadWriteCloser) *MsgpackCodec {
cc := &MsgpackCodec{
conn: conn,
}
if bufReads {
cc.bufR = bufio.NewReader(conn)
cc.dec = codec.NewDecoder(cc.bufR, msgpackHandle)
} else {
cc.dec = codec.NewDecoder(cc.conn, msgpackHandle)
}
if bufWrites {
cc.bufW = bufio.NewWriter(conn)
cc.enc = codec.NewEncoder(cc.bufW, msgpackHandle)
} else {
cc.enc = codec.NewEncoder(cc.conn, msgpackHandle)
}
return cc
}
func MockEncoder(obj interface{}) []byte {
buf := bytes.NewBuffer(nil)
encoder := codec.NewEncoder(buf, msgpackHandle)
err := encoder.Encode(obj)
if err != nil {
panic(err)
}
return buf.Bytes()
}
// SnapshotRPC is a streaming client function for performing a snapshot RPC
// request to a remote server. It will create a fresh connection for each
// request, send the request header, and then stream in any data from the
// reader (for a restore). It will then parse the received response header, and
// if there's no error will return an io.ReadCloser (that you must close) with
// the streaming output (for a snapshot). If the reply contains an error, this
// will always return an error as well, so you don't need to check the error
// inside the filled-in reply.
func SnapshotRPC(pool *ConnPool, dc string, addr net.Addr,
args *structs.SnapshotRequest, in io.Reader, reply *structs.SnapshotResponse) (io.ReadCloser, error) {
conn, hc, err := pool.Dial(dc, addr)
if err != nil {
return nil, err
}
// keep will disarm the defer on success if we are returning the caller
// our connection to stream the output.
var keep bool
defer func() {
if !keep {
conn.Close()
}
}()
// Write the snapshot RPC byte to set the mode, then perform the
// request.
if _, err := conn.Write([]byte{byte(rpcSnapshot)}); err != nil {
return nil, fmt.Errorf("failed to write stream type: %v", err)
}
// Push the header encoded as msgpack, then stream the input.
enc := codec.NewEncoder(conn, &codec.MsgpackHandle{})
if err := enc.Encode(&args); err != nil {
return nil, fmt.Errorf("failed to encode request: %v", err)
}
if _, err := io.Copy(conn, in); err != nil {
return nil, fmt.Errorf("failed to copy snapshot in: %v", err)
}
// Our RPC protocol requires support for a half-close in order to signal
// the other side that they are done reading the stream, since we don't
// know the size in advance. This saves us from having to buffer just to
// calculate the size.
if hc != nil {
if err := hc.CloseWrite(); err != nil {
return nil, fmt.Errorf("failed to half close snapshot connection: %v", err)
}
} else {
return nil, fmt.Errorf("snapshot connection requires half-close support")
}
// Pull the header decoded as msgpack. The caller can continue to read
// the conn to stream the remaining data.
dec := codec.NewDecoder(conn, &codec.MsgpackHandle{})
if err := dec.Decode(reply); err != nil {
return nil, fmt.Errorf("failed to decode response: %v", err)
}
if reply.Error != "" {
return nil, errors.New(reply.Error)
}
keep = true
return conn, nil
}
// See raft.SnapshotSink.
func (m *MockSnapshot) Persist(sink raft.SnapshotSink) error {
hd := codec.MsgpackHandle{}
enc := codec.NewEncoder(sink, &hd)
if err := enc.Encode(m.logs[:m.maxIndex]); err != nil {
sink.Cancel()
return err
}
sink.Close()
return nil
}
// AckPayload is called to produce a payload to send back in response to a ping
// request.
func (p *pingDelegate) AckPayload() []byte {
var buf bytes.Buffer
// The first byte is the version number, forming a simple header.
version := []byte{PingVersion}
buf.Write(version)
// The rest of the message is the serialized coordinate.
enc := codec.NewEncoder(&buf, &codec.MsgpackHandle{})
if err := enc.Encode(p.serf.coordClient.GetCoordinate()); err != nil {
log.Printf("[ERR] serf: Failed to encode coordinate: %v\n", err)
}
return buf.Bytes()
}
// encodeTags is used to encode a tag map
func (s *Serf) encodeTags(tags map[string]string) []byte {
// Support role-only backwards compatibility
if s.ProtocolVersion() < 3 {
role := tags["role"]
return []byte(role)
}
// Use a magic byte prefix and msgpack encode the tags
var buf bytes.Buffer
buf.WriteByte(tagMagicByte)
enc := codec.NewEncoder(&buf, &codec.MsgpackHandle{})
if err := enc.Encode(tags); err != nil {
panic(fmt.Sprintf("Failed to encode tags: %v", err))
}
return buf.Bytes()
}
func (s *consulSnapshot) Persist(sink raft.SnapshotSink) error {
defer metrics.MeasureSince([]string{"consul", "fsm", "persist"}, time.Now())
// Register the nodes
encoder := codec.NewEncoder(sink, msgpackHandle)
// Write the header
header := snapshotHeader{
LastIndex: s.state.LastIndex(),
}
if err := encoder.Encode(&header); err != nil {
sink.Cancel()
return err
}
if err := s.persistNodes(sink, encoder); err != nil {
sink.Cancel()
return err
}
if err := s.persistSessions(sink, encoder); err != nil {
sink.Cancel()
return err
}
if err := s.persistACLs(sink, encoder); err != nil {
sink.Cancel()
return err
}
if err := s.persistKVs(sink, encoder); err != nil {
sink.Cancel()
return err
}
if err := s.persistTombstones(sink, encoder); err != nil {
sink.Cancel()
return err
}
if err := s.persistPreparedQueries(sink, encoder); err != nil {
sink.Cancel()
return err
}
return nil
}
func (s *nomadSnapshot) Persist(sink raft.SnapshotSink) error {
defer metrics.MeasureSince([]string{"nomad", "fsm", "persist"}, time.Now())
// Register the nodes
encoder := codec.NewEncoder(sink, structs.MsgpackHandle)
// Write the header
header := snapshotHeader{}
if err := encoder.Encode(&header); err != nil {
sink.Cancel()
return err
}
// Write the time table
sink.Write([]byte{byte(TimeTableSnapshot)})
if err := s.timetable.Serialize(encoder); err != nil {
sink.Cancel()
return err
}
// Write all the data out
if err := s.persistIndexes(sink, encoder); err != nil {
sink.Cancel()
return err
}
if err := s.persistNodes(sink, encoder); err != nil {
sink.Cancel()
return err
}
if err := s.persistJobs(sink, encoder); err != nil {
sink.Cancel()
return err
}
if err := s.persistEvals(sink, encoder); err != nil {
sink.Cancel()
return err
}
if err := s.persistAllocs(sink, encoder); err != nil {
sink.Cancel()
return err
}
if err := s.persistPeriodicLaunches(sink, encoder); err != nil {
sink.Cancel()
return err
}
return nil
}
// AckPayload is called to produce a payload to send back in response to a ping
// request. In this case we send back a legit ping response with a bad coordinate.
func (p *pingDimensionMetaDelegate) AckPayload() []byte {
var buf bytes.Buffer
// The first byte is the version number, forming a simple header.
version := []byte{PingVersion}
buf.Write(version)
// Make a bad coordinate with the wrong number of dimensions.
coord := coordinate.NewCoordinate(coordinate.DefaultConfig())
coord.Vec = make([]float64, 2*len(coord.Vec))
// The rest of the message is the serialized coordinate.
enc := codec.NewEncoder(&buf, &codec.MsgpackHandle{})
if err := enc.Encode(coord); err != nil {
p.t.Fatalf("err: %v", err)
}
return buf.Bytes()
}
// ClientFromConfig is used to create a new RPC client given the
// configuration object. This will return a client, or an error if
// the connection could not be established.
func ClientFromConfig(c *Config) (*RPCClient, error) {
// Setup the defaults
if c.Timeout == 0 {
c.Timeout = DefaultTimeout
}
// Try to dial to serf
conn, err := net.DialTimeout("tcp", c.Addr, c.Timeout)
if err != nil {
return nil, err
}
// Create the client
client := &RPCClient{
seq: 0,
timeout: c.Timeout,
conn: conn.(*net.TCPConn),
reader: bufio.NewReader(conn),
writer: bufio.NewWriter(conn),
dispatch: make(map[uint64]seqHandler),
shutdownCh: make(chan struct{}),
}
client.dec = codec.NewDecoder(client.reader,
&codec.MsgpackHandle{RawToString: true, WriteExt: true})
client.enc = codec.NewEncoder(client.writer,
&codec.MsgpackHandle{RawToString: true, WriteExt: true})
go client.listen()
// Do the initial handshake
if err := client.handshake(); err != nil {
client.Close()
return nil, err
}
// Do the initial authentication if needed
if c.AuthKey != "" {
if err := client.auth(c.AuthKey); err != nil {
client.Close()
return nil, err
}
}
return client, err
}
// listen is a long running routine that listens for new clients
func (i *AgentIPC) listen() {
for {
conn, err := i.listener.Accept()
if err != nil {
if i.stop {
return
}
i.logger.Printf("[ERR] agent.ipc: Failed to accept client: %v", err)
continue
}
i.logger.Printf("[INFO] agent.ipc: Accepted client: %v", conn.RemoteAddr())
metrics.IncrCounter([]string{"agent", "ipc", "accept"}, 1)
// Wrap the connection in a client
client := &IPCClient{
name: conn.RemoteAddr().String(),
conn: conn,
reader: bufio.NewReader(conn),
writer: bufio.NewWriter(conn),
eventStreams: make(map[uint64]*eventStream),
pendingQueries: make(map[uint64]*serf.Query),
}
client.dec = codec.NewDecoder(client.reader,
&codec.MsgpackHandle{RawToString: true, WriteExt: true})
client.enc = codec.NewEncoder(client.writer,
&codec.MsgpackHandle{RawToString: true, WriteExt: true})
if err != nil {
i.logger.Printf("[ERR] agent.ipc: Failed to create decoder: %v", err)
conn.Close()
continue
}
// Register the client
i.Lock()
if !i.stop {
i.clients[client.name] = client
go i.handleClient(client)
} else {
conn.Close()
}
i.Unlock()
}
}
// handleConn is used to handle an inbound connection for its lifespan.
func (n *NetworkTransport) handleConn(conn net.Conn) {
defer conn.Close()
r := bufio.NewReader(conn)
w := bufio.NewWriter(conn)
dec := codec.NewDecoder(r, &codec.MsgpackHandle{})
enc := codec.NewEncoder(w, &codec.MsgpackHandle{})
for {
if err := n.handleCommand(r, dec, enc); err != nil {
if err != io.EOF {
n.logger.Printf("[ERR] raft-net: Failed to decode incoming command: %v", err)
}
return
}
if err := w.Flush(); err != nil {
n.logger.Printf("[ERR] raft-net: Failed to flush response: %v", err)
return
}
}
}
func TestTimeTable_SerializeDeserialize(t *testing.T) {
tt := NewTimeTable(time.Second, time.Minute)
// Witness some data
start := time.Now()
plusOne := start.Add(time.Minute)
plusTwo := start.Add(2 * time.Minute)
plusFive := start.Add(5 * time.Minute)
plusThirty := start.Add(30 * time.Minute)
plusHour := start.Add(60 * time.Minute)
tt.Witness(2, start)
tt.Witness(10, plusOne)
tt.Witness(20, plusTwo)
tt.Witness(30, plusFive)
tt.Witness(40, plusThirty)
tt.Witness(50, plusHour)
var buf bytes.Buffer
enc := codec.NewEncoder(&buf, msgpackHandle)
err := tt.Serialize(enc)
if err != nil {
t.Fatalf("err: %v", err)
}
dec := codec.NewDecoder(&buf, msgpackHandle)
tt2 := NewTimeTable(time.Second, time.Minute)
err = tt2.Deserialize(dec)
if err != nil {
t.Fatalf("err: %v", err)
}
if !reflect.DeepEqual(tt.table, tt2.table) {
t.Fatalf("bad: %#v %#v", tt, tt2)
}
}
// NewRPCClient is used to create a new RPC client given the address.
// This will properly dial, handshake, and start listening
func NewRPCClient(addr string) (*RPCClient, error) {
var conn net.Conn
var err error
if envAddr := os.Getenv("CONSUL_RPC_ADDR"); envAddr != "" {
addr = envAddr
}
// Try to dial to agent
mode := "tcp"
if strings.HasPrefix(addr, "/") {
mode = "unix"
}
if conn, err = net.Dial(mode, addr); err != nil {
return nil, err
}
// Create the client
client := &RPCClient{
seq: 0,
conn: conn,
reader: bufio.NewReader(conn),
writer: bufio.NewWriter(conn),
dispatch: make(map[uint64]seqHandler),
shutdownCh: make(chan struct{}),
}
client.dec = codec.NewDecoder(client.reader, msgpackHandle)
client.enc = codec.NewEncoder(client.writer, msgpackHandle)
go client.listen()
// Do the initial handshake
if err := client.handshake(); err != nil {
client.Close()
return nil, err
}
return client, err
}
// encodeMsgPack is used to encode an object with msgpack
func encodeMsgPack(msg interface{}) ([]byte, error) {
var buf bytes.Buffer
err := codec.NewEncoder(&buf, msgpackHandle).Encode(msg)
return buf.Bytes(), err
}
// sendLocalState is invoked to send our local state over a tcp connection
func (m *Memberlist) sendLocalState(conn net.Conn, join bool) error {
// Setup a deadline
conn.SetDeadline(time.Now().Add(m.config.TCPTimeout))
// Prepare the local node state
m.nodeLock.RLock()
localNodes := make([]pushNodeState, len(m.nodes))
for idx, n := range m.nodes {
localNodes[idx].Name = n.Name
localNodes[idx].Addr = n.Addr
localNodes[idx].Port = n.Port
localNodes[idx].Incarnation = n.Incarnation
localNodes[idx].State = n.State
localNodes[idx].Meta = n.Meta
localNodes[idx].Vsn = []uint8{
n.PMin, n.PMax, n.PCur,
n.DMin, n.DMax, n.DCur,
}
}
m.nodeLock.RUnlock()
// Get the delegate state
var userData []byte
if m.config.Delegate != nil {
userData = m.config.Delegate.LocalState(join)
}
// Create a bytes buffer writer
bufConn := bytes.NewBuffer(nil)
// Send our node state
header := pushPullHeader{Nodes: len(localNodes), UserStateLen: len(userData), Join: join}
hd := codec.MsgpackHandle{}
enc := codec.NewEncoder(bufConn, &hd)
// Begin state push
if _, err := bufConn.Write([]byte{byte(pushPullMsg)}); err != nil {
return err
}
if err := enc.Encode(&header); err != nil {
return err
}
for i := 0; i < header.Nodes; i++ {
if err := enc.Encode(&localNodes[i]); err != nil {
return err
}
}
// Write the user state as well
if userData != nil {
if _, err := bufConn.Write(userData); err != nil {
return err
}
}
// Get the send buffer
sendBuf := bufConn.Bytes()
// Check if compresion is enabled
if m.config.EnableCompression {
compBuf, err := compressPayload(bufConn.Bytes())
if err != nil {
m.logger.Printf("[ERROR] memberlist: Failed to compress local state: %v", err)
} else {
sendBuf = compBuf.Bytes()
}
}
// Check if encryption is enabled
if m.config.EncryptionEnabled() {
crypt, err := m.encryptLocalState(sendBuf)
if err != nil {
m.logger.Printf("[ERROR] memberlist: Failed to encrypt local state: %v", err)
return err
}
sendBuf = crypt
}
// Write out the entire send buffer
metrics.IncrCounter([]string{"memberlist", "tcp", "sent"}, float32(len(sendBuf)))
if _, err := conn.Write(sendBuf); err != nil {
return err
}
return nil
}
// Encode is used to encode a MsgPack object with type prefix
func Encode(t MessageType, msg interface{}) ([]byte, error) {
var buf bytes.Buffer
buf.WriteByte(uint8(t))
err := codec.NewEncoder(&buf, msgpackHandle).Encode(msg)
return buf.Bytes(), err
}
func TestTCPPushPull(t *testing.T) {
m := GetMemberlist(t)
defer m.Shutdown()
m.nodes = append(m.nodes, &nodeState{
Node: Node{
Name: "Test 0",
Addr: net.ParseIP(m.config.BindAddr),
Port: uint16(m.config.BindPort),
},
Incarnation: 0,
State: stateSuspect,
StateChange: time.Now().Add(-1 * time.Second),
})
addr := fmt.Sprintf("%s:%d", m.config.BindAddr, m.config.BindPort)
conn, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("unexpected err %s", err)
}
defer conn.Close()
localNodes := make([]pushNodeState, 3)
localNodes[0].Name = "Test 0"
localNodes[0].Addr = net.ParseIP(m.config.BindAddr)
localNodes[0].Port = uint16(m.config.BindPort)
localNodes[0].Incarnation = 1
localNodes[0].State = stateAlive
localNodes[1].Name = "Test 1"
localNodes[1].Addr = net.ParseIP(m.config.BindAddr)
localNodes[1].Port = uint16(m.config.BindPort)
localNodes[1].Incarnation = 1
localNodes[1].State = stateAlive
localNodes[2].Name = "Test 2"
localNodes[2].Addr = net.ParseIP(m.config.BindAddr)
localNodes[2].Port = uint16(m.config.BindPort)
localNodes[2].Incarnation = 1
localNodes[2].State = stateAlive
// Send our node state
header := pushPullHeader{Nodes: 3}
hd := codec.MsgpackHandle{}
enc := codec.NewEncoder(conn, &hd)
// Send the push/pull indicator
conn.Write([]byte{byte(pushPullMsg)})
if err := enc.Encode(&header); err != nil {
t.Fatalf("unexpected err %s", err)
}
for i := 0; i < header.Nodes; i++ {
if err := enc.Encode(&localNodes[i]); err != nil {
t.Fatalf("unexpected err %s", err)
}
}
// Read the message type
var msgType messageType
if err := binary.Read(conn, binary.BigEndian, &msgType); err != nil {
t.Fatalf("unexpected err %s", err)
}
var bufConn io.Reader = conn
msghd := codec.MsgpackHandle{}
dec := codec.NewDecoder(bufConn, &msghd)
// Check if we have a compressed message
if msgType == compressMsg {
var c compress
if err := dec.Decode(&c); err != nil {
t.Fatalf("unexpected err %s", err)
}
decomp, err := decompressBuffer(&c)
if err != nil {
t.Fatalf("unexpected err %s", err)
}
// Reset the message type
msgType = messageType(decomp[0])
// Create a new bufConn
bufConn = bytes.NewReader(decomp[1:])
// Create a new decoder
dec = codec.NewDecoder(bufConn, &hd)
}
// Quit if not push/pull
if msgType != pushPullMsg {
t.Fatalf("bad message type")
}
if err := dec.Decode(&header); err != nil {
t.Fatalf("unexpected err %s", err)
}
// Allocate space for the transfer
remoteNodes := make([]pushNodeState, header.Nodes)
// Try to decode all the states
for i := 0; i < header.Nodes; i++ {
if err := dec.Decode(&remoteNodes[i]); err != nil {
t.Fatalf("unexpected err %s", err)
}
}
if len(remoteNodes) != 1 {
t.Fatalf("bad response")
}
n := &remoteNodes[0]
if n.Name != "Test 0" {
t.Fatalf("bad name")
}
if bytes.Compare(n.Addr, net.ParseIP(m.config.BindAddr)) != 0 {
t.Fatal("bad addr")
}
if n.Incarnation != 0 {
t.Fatal("bad incarnation")
}
if n.State != stateSuspect {
t.Fatal("bad state")
}
}