etcd网络层(六)——模拟etcd网络层的stream通道维护长连接

代码如下：

package main

import (
	"io"
	"log"
	"net"
	"fmt"
	"sync"
	"time"
	"strings"
	"strconv"
	"net/http"
	"encoding/json"
	"encoding/binary"
)

func main() {
	peerURL1 := "http://127.0.0.1:8081"
	peerURL2 := "http://127.0.0.1:8082"

	id1 := getPid(peerURL1)
	id2 := getPid(peerURL2)

	//开启节点1
	tr1 := &Transport{}
	tr1.Start(int64(id1))
	go func() {
		err := http.ListenAndServe(":8081",tr1.Handler())
		log.Fatal(err)
	}()



	//开启节点2
	tr2 := &Transport{}
	tr2.Start(int64(id2))
	go func() {
		err := http.ListenAndServe(":8082",tr2.Handler())
		log.Fatal(err)
	}()


	time.Sleep(time.Second*3)

	//节点1添加节点2
	tr1.AddPeer(int64(id2),peerURL2)

	//节点2添加节点1
	tr2.AddPeer(int64(id1),peerURL1)




	go func() {
		ticker := time.NewTicker(time.Second*10)
		defer ticker.Stop()
		for {
			select {
			case <- ticker.C:
				peers := tr1.GetPeer()
				for i := range peers {
					peers[i].send(&Message{MsgType:msgTypeProp, MsgBody:"Hello,I am tr1"})
				}
			}
		}
	}()


	go func() {
		ticker := time.NewTicker(time.Second*10)
		defer ticker.Stop()
		for {
			select {
			case <- ticker.C:
				peers := tr2.GetPeer()
				for i := range peers {
					peers[i].send(&Message{MsgType:msgTypeProp, MsgBody:"Hello,I am tr2"})
				}
			}
		}
	}()


	time.Sleep(time.Minute*10)


	tr1.Stop()
	tr2.Stop()
}



type Transport struct {
	ClusterID  int64
	ID          int64   // local member ID  当前节点自己的ID
	streamRt http.RoundTripper
	mu sync.RWMutex
	peers   map[int64]*peer    // peers map
}

func (tr *Transport) Start(id int64) error {
	tr.ID = id
	tr.streamRt = &http.Transport{
	Dial: (&net.Dialer{
		Timeout: 10*time.Second,
		// value taken from http.DefaultTransport
		KeepAlive: 30 * time.Second,
	}).Dial,
	}
	tr.peers = make(map[int64]*peer)

	return nil
}

func (tr *Transport) GetPeer() (result []*peer) {
	tr.mu.RLock()
	defer tr.mu.RUnlock()

	for k := range tr.peers{
		result = append(result,tr.peers[k])
	}

	return
}

func (tr *Transport) AddPeer(id int64, peerURL string) {
	tr.mu.RLock()
	if _,ok := tr.peers[id]; ok {
		tr.mu.RUnlock()
		return
	}

	tr.mu.RUnlock()

	peer := startPeer(tr,peerURL,tr.ID,id)

	tr.mu.Lock()
	tr.peers[id] = peer
	tr.mu.Unlock()
}

func (tr *Transport) Handler() http.Handler {
	streamHandler := newStreamHandler(tr,tr.ID,tr.ClusterID)
	mux := http.NewServeMux()
	mux.Handle("/raft/stream"+"/", streamHandler)
	return mux
}



type streamHandler struct {
	tr         *Transport   //关联的rafthttp.Transport实例
	id        int64  //当前节点ID
	cid      int64   //当前集群ID
}

func newStreamHandler(tr *Transport, id, cid int64) http.Handler  {
	return &streamHandler{
		tr:tr,
		id:id,
		cid:cid,
	}
}

func (h *streamHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	//请求参数校验，如Method是否是GET，检验集群ID
	if r.Method != "GET" {
		w.Header().Set("Allow", "GET")
		http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
		return
	}

	id := r.Header.Get("PeerID")
	if id == ""{
		w.Header().Set("PeerID", "Must")
		http.Error(w, "PeerID is not allow empty", http.StatusMethodNotAllowed)
		return
	}

	pid,_ := strconv.ParseUint(id,10,64)




	p,ok := h.tr.peers[int64(pid)]
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		time.Sleep(time.Second*5)
		return 
	}

	w.WriteHeader(http.StatusOK)  //返回状态码 200

	w.(http.Flusher).Flush()  //调用Flush()方法将响应数据发送到对端节点


   c := newCloseNotifier()
	conn := &outgoingConn{  //创建outgoingConn实例
		Writer:  w,
		Flusher: w.(http.Flusher),
		localID: h.tr.ID,
		Closer:c,
		peerID:  h.id,
	}
	p.attachOutgoingConn(conn)   //建立连接,将outgoingConn实例与对应的streamWriter实例绑定
	<- c.closeNotify()

}

type closeNotifier struct {
	done chan struct{}
}

func newCloseNotifier() *closeNotifier {
	return &closeNotifier{
		done: make(chan struct{}),
	}
}

func (n *closeNotifier) Close() error {
	close(n.done)
	return nil
}

func (n *closeNotifier) closeNotify() <-chan struct{} { return n.done }



func (tr *Transport) Stop()  {
	tr.mu.RLock()
	defer tr.mu.RUnlock()
	for _,v := range tr.peers{
		v.stop()
	}
	tr.peers = nil
}


const (
	msgTypeHeartbeat  =  "01"  //心跳
	msgTypeProp = "02"  //prop消息
)


type Message struct {
	MsgType string
	MsgBody string
}


type peer struct {
	localID int64   //当前节点ID
	// id of the remote raft peer node
	id  int64   //该peer实例对应的节点ID，对端ID

	writer         *streamWriter   //负责向Stream消息通道中写消息

	msgAppReader   *streamReader  //负责从Stream消息通道中读消息

	msgc chan *Message
	stopc  chan struct{}
}

func startPeer(t *Transport, peerURL string, localID,peerID int64) *peer {
	pr := & peer{
		localID:localID,
		id:peerID,
		writer:newstreamWriter(localID,peerID),
		msgAppReader:newstreamReader(localID,peerID,peerURL,t),
		msgc : make(chan *Message,20),
		stopc : make(chan struct{}),
	}

	go func() {
		for msg := range pr.msgc {
			select {
			case pr.writer.msgc <- msg:
			default:
				log.Printf("write to writer error msg is %v",msg)
			}
		}
	}()


	return pr
}

func (pr *peer) stop()  {
	close(pr.stopc)
}

func (pr *peer) send(msg *Message) bool{
	select {
	case pr.msgc <- msg:
		return true
	default:
		return false
	}
}

func (pr *peer) attachOutgoingConn(conn *outgoingConn) {
	select {
	case pr.writer.connc <- conn:

	default:
		log.Printf("attachOutgoingConn error")
	}
}


type streamWriter struct {
	localID int64   //本端的ID
	peerID  int64  //对端节点的ID

	closer  io.Closer    //负责关闭底层的长连接

	mu sync.Mutex

	enc  *messageEncoder
	msgc  chan *Message   //Peer会将待发送的消息写入到该通道，streamWriter则从该通道中读取消息并发送出去
	connc chan *outgoingConn   //通过该通道获取当前streamWriter实例关联的底层网络连接，  outgoingConn其实是对网络连接的一层封装，其中记录了当前连接使用的协议版本，以及用于关闭连接的Flusher和Closer等信息。
	stopc chan struct{}
}

func newstreamWriter(localID,peerID int64)  *streamWriter{
	sw := &streamWriter{
		localID:localID,
		peerID:peerID,
		msgc:make(chan *Message,20),
		connc:make(chan *outgoingConn),
		stopc:make(chan struct{}),
	}

	go sw.run()
	return sw
}


func (sw *streamWriter) writec() (chan<- *Message) {
	sw.mu.Lock()
	defer sw.mu.Unlock()
	return sw.msgc
}

func (sw *streamWriter) run() {
	var (
	 	heartbeatC  <-chan time.Time

		flusher    http.Flusher  //负责刷新底层连接，将数据真正发送出去

		msgc  chan *Message
	)


	tickc := time.NewTicker(time.Second*7)  //发送心跳的定时器
	defer tickc.Stop()

	for {

		select {
		case msg := <- msgc:
			err := sw.enc.encode(msg)
			if err != nil {
				log.Printf("send to peer peerID is %d fail,error is %v",sw.peerID,err)
			}else {
				flusher.Flush()
				log.Printf("send to peer  peerID is %d success,  MsgType is :%s,MsgBody is:%s",sw.peerID,msg.MsgType,msg.MsgBody)
			}

		case <-heartbeatC: //向对端发送心跳消息
			err := sw.enc.encode(&Message{
				MsgType: msgTypeHeartbeat,
				MsgBody: time.Now().Format("2006-01-02 15:04:05"),
			})
			if err != nil {
				log.Printf("send to peer heartbeat data fail  peerID is %d,error is %v",sw.peerID,err)
			}else {
				flusher.Flush()
				log.Printf("send to peer heartbeat data success  peerID is %d ",sw.peerID)
			}
		case conn := <- sw.connc:
			sw.enc = &messageEncoder{w:conn.Writer}
			flusher = conn.Flusher
			sw.closer = conn.Closer

			heartbeatC,msgc = tickc.C,sw.msgc

		case <- sw.stopc:
			log.Println("msgWriter stop!")
			sw.closer.Close()
			return
		}
	}
}

func (sw *streamWriter) stop()  {
	close(sw.stopc)
}


type streamReader struct {
	localID  int64
	peerID int64  //对端节点的ID
	tr     *Transport   //关联的rafthttp.Transport实例
	peerURL  string  //对端URL


	mu     sync.Mutex
	closer io.Closer   //负责关闭底层的长连接

	done   chan struct{}
}

func newstreamReader(localID,peerID int64,peerURL string,tr *Transport) *streamReader {
	sr := &streamReader{
		localID:localID,
		peerID:peerID,
		peerURL:peerURL,
		tr :tr,
		done : make(chan struct{}),
	}
	go sr.run()

	return sr
}

func (sr *streamReader) run() {
	time.Sleep(time.Second*5)
	for {
		readColser,err := sr.dial()
		if err != nil {
			log.Printf("dial peer error,peerID is %d,err is %v",sr.peerID,err)
			time.Sleep(time.Second*10)
			continue
		}
		sr.closer = readColser

		err = sr.decodeLoop(readColser)
		if err != nil {
			log.Printf("decodeLoop error,peerID is %d, error is %v",sr.peerID,err)
		}
		sr.closer.Close()
	}
}

func (sr *streamReader) dial() (io.ReadCloser, error) {
	req, err := http.NewRequest("GET", sr.peerURL+"/raft/stream/dial", nil)
	if err != nil {
		return nil,err
	}

	req.Header.Set("PeerID",fmt.Sprintf("%d",sr.localID))

	resp,err := sr.tr.streamRt.RoundTrip(req)
	if err != nil {
		return nil,err
	}

	return resp.Body,nil
}

func (sr *streamReader) decodeLoop(rc io.ReadCloser) error {
	dec := &messageDecoder{rc}
	for {
		msg,err := dec.decode()
		if err != nil {
			log.Printf("read decodeLoop error,peerID is %d,err is %v",sr.peerID,err)
			continue
		}

		log.Printf("read from peer MsgType is %s,MsgBody is %s",msg.MsgType,msg.MsgBody)
	}
}


type outgoingConn struct {
	io.Writer
	http.Flusher
	io.Closer
	localID int64
	peerID  int64
}


type messageEncoder struct {
	w io.Writer
}

func (m *messageEncoder) encode(msg *Message) error {
	byts,err := json.Marshal(msg)
	if err != nil {
		return err
	}

	dataLen := make([]byte,8)
	binary.BigEndian.PutUint64(dataLen,uint64(len(byts)))

	sendData := append(dataLen,byts...)
	_,err = m.w.Write(sendData)
	if err != nil {
		return err
	}

	return nil
}

type messageDecoder struct {
	r io.Reader
}

func (dec *messageDecoder) decode() (*Message,error) {
	var m Message
	var l uint64
	if err := binary.Read(dec.r, binary.BigEndian, &l); err != nil {
		return nil, err
	}

	buf := make([]byte, int(l))
	if _, err := io.ReadFull(dec.r, buf); err != nil {
		return &m, err
	}
	err := json.Unmarshal(buf,&m)
	if err != nil {
		return nil,err
	}
	return &m, nil
}



func getPid(purl string)  int64 {
	index := strings.LastIndex(purl,":")
	if index > 0 {
		id,err := strconv.ParseInt(purl[index+1:],10,64)
		if err != nil {
			println(err)
		}
		return  id
	}

	return 0
}