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

- 日理万妓 2021-12-20 13:05 184阅读 0赞

代码如下：

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
    }