Golang codec.NewEncoder函數代碼示例

// 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
}