P2p first version

lib/mux/mux.go

@@ -0,0 +1,226 @@
+package mux
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"github.com/cnlh/nps/lib/pool"
+	"math"
+	"net"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+const (
+	MUX_PING_FLAG int32 = iota
+	MUX_NEW_CONN_OK
+	MUX_NEW_CONN_Fail
+	MUX_NEW_MSG
+	MUX_MSG_SEND_OK
+	MUX_NEW_CONN
+	MUX_PING
+	MUX_CONN_CLOSE
+)
+
+type Mux struct {
+	net.Listener
+	conn         net.Conn
+	connMap      *connMap
+	sendMsgCh    chan *msg  //write msg chan
+	sendStatusCh chan int32 //write read ok chan
+	newConnCh    chan *conn
+	id           int32
+	closeChan    chan struct{}
+	isClose      bool
+	sync.Mutex
+}
+
+func NewMux(c net.Conn) *Mux {
+	m := &Mux{
+		conn:         c,
+		connMap:      NewConnMap(),
+		sendMsgCh:    make(chan *msg),
+		sendStatusCh: make(chan int32),
+		id:           0,
+		closeChan:    make(chan struct{}),
+		newConnCh:    make(chan *conn),
+		isClose:      false,
+	}
+	//read session by flag
+	go m.readSession()
+	//write session
+	go m.writeSession()
+	//ping
+	go m.ping()
+	return m
+}
+
+func (s *Mux) NewConn() (*conn, error) {
+	if s.isClose {
+		return nil, errors.New("the mux has closed")
+	}
+	conn := NewConn(s.getId(), s, s.sendMsgCh, s.sendStatusCh)
+	raw := bytes.NewBuffer([]byte{})
+	if err := binary.Write(raw, binary.LittleEndian, MUX_NEW_CONN); err != nil {
+		return nil, err
+	}
+	if err := binary.Write(raw, binary.LittleEndian, conn.connId); err != nil {
+		return nil, err
+	}
+	//it must be set before send
+	s.connMap.Set(conn.connId, conn)
+	if _, err := s.conn.Write(raw.Bytes()); err != nil {
+		return nil, err
+	}
+	select {
+	case <-conn.connStatusOkCh:
+		return conn, nil
+	case <-conn.connStatusFailCh:
+	}
+	return nil, errors.New("create connection fail，the server refused the connection")
+}
+
+func (s *Mux) Accept() (net.Conn, error) {
+	if s.isClose {
+		return nil, errors.New("accpet error,the conn has closed")
+	}
+	return <-s.newConnCh, nil
+}
+
+func (s *Mux) Addr() net.Addr {
+	return s.conn.LocalAddr()
+}
+
+func (s *Mux) ping() {
+	go func() {
+		ticker := time.NewTicker(time.Second * 5)
+		raw := bytes.NewBuffer([]byte{})
+		for {
+			select {
+			case <-ticker.C:
+			}
+			//Avoid going beyond the scope
+			if (math.MaxInt32 - s.id) < 10000 {
+				s.id = 0
+			}
+			raw.Reset()
+			binary.Write(raw, binary.LittleEndian, MUX_PING_FLAG)
+			binary.Write(raw, binary.LittleEndian, MUX_PING)
+			if _, err := s.conn.Write(raw.Bytes()); err != nil {
+				s.Close()
+				break
+			}
+		}
+	}()
+	select {
+	case <-s.closeChan:
+	}
+}
+
+func (s *Mux) writeSession() {
+	go func() {
+		raw := bytes.NewBuffer([]byte{})
+		for {
+			raw.Reset()
+			select {
+			case msg := <-s.sendMsgCh:
+				if msg.content == nil { //close
+					binary.Write(raw, binary.LittleEndian, MUX_CONN_CLOSE)
+					binary.Write(raw, binary.LittleEndian, msg.connId)
+					break
+				}
+				binary.Write(raw, binary.LittleEndian, MUX_NEW_MSG)
+				binary.Write(raw, binary.LittleEndian, msg.connId)
+				binary.Write(raw, binary.LittleEndian, int32(len(msg.content)))
+				binary.Write(raw, binary.LittleEndian, msg.content)
+			case connId := <-s.sendStatusCh:
+				binary.Write(raw, binary.LittleEndian, MUX_MSG_SEND_OK)
+				binary.Write(raw, binary.LittleEndian, connId)
+			}
+			if _, err := s.conn.Write(raw.Bytes()); err != nil {
+				s.Close()
+				break
+			}
+		}
+	}()
+	select {
+	case <-s.closeChan:
+	}
+}
+
+func (s *Mux) readSession() {
+	go func() {
+		raw := bytes.NewBuffer([]byte{})
+		for {
+			var flag, i int32
+			if binary.Read(s.conn, binary.LittleEndian, &flag) == nil {
+				if binary.Read(s.conn, binary.LittleEndian, &i) != nil {
+					break
+				}
+				switch flag {
+				case MUX_NEW_CONN: //new conn
+					conn := NewConn(i, s, s.sendMsgCh, s.sendStatusCh)
+					s.connMap.Set(i, conn) //it has been set before send ok
+					s.newConnCh <- conn
+					raw.Reset()
+					binary.Write(raw, binary.LittleEndian, MUX_NEW_CONN_OK)
+					binary.Write(raw, binary.LittleEndian, i)
+					s.conn.Write(raw.Bytes())
+					continue
+				case MUX_PING_FLAG: //ping
+					continue
+				}
+				if conn, ok := s.connMap.Get(i); ok {
+					switch flag {
+					case MUX_NEW_MSG: //new msg from remote conn
+						buf := pool.BufPoolCopy.Get().([]byte)
+						if n, err := ReadLenBytes(buf, s.conn); err == nil {
+							if !conn.isClose {
+								conn.readMsgCh <- buf[:n]
+							} else {
+								pool.PutBufPoolCopy(buf)
+							}
+						} else { //read len bytes error,the mux has broken
+							break
+						}
+					case MUX_MSG_SEND_OK: //the remote has read
+						conn.getStatusCh <- struct{}{}
+					case MUX_NEW_CONN_OK: //conn ok
+						conn.connStatusOkCh <- struct{}{}
+					case MUX_NEW_CONN_Fail:
+						conn.connStatusFailCh <- struct{}{}
+					case MUX_CONN_CLOSE: //close the connection
+						conn.Close()
+					}
+				}
+			} else {
+				break
+			}
+		}
+		s.Close()
+	}()
+	select {
+	case <-s.closeChan:
+	}
+}
+
+func (s *Mux) Close() error {
+	if s.isClose {
+		return errors.New("the mux has closed")
+	}
+	s.isClose = true
+	s.connMap.Close()
+	s.closeChan <- struct{}{}
+	s.closeChan <- struct{}{}
+	s.closeChan <- struct{}{}
+	close(s.closeChan)
+	close(s.sendMsgCh)
+	close(s.sendStatusCh)
+	return s.conn.Close()
+}
+
+//get new connId as unique flag
+func (s *Mux) getId() int32 {
+	return atomic.AddInt32(&s.id, 1)
+}