MUX optimization

lib/mux/mux.go

@@ -22,38 +22,35 @@ const (
 	MUX_PING
 	MUX_CONN_CLOSE
 	MUX_PING_RETURN
+	MUX_STOP_WRITE
 	RETRY_TIME = 2 //Heart beat allowed fault tolerance times
 )
 
 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
-	pingOk       int
+	conn          net.Conn
+	connMap       *connMap
+	newConnCh     chan *conn
+	id            int32
+	closeChan     chan struct{}
+	IsClose       bool
+	pingOk        int
+	waitQueueSize int
 	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,
+		conn:          c,
+		connMap:       NewConnMap(),
+		id:            0,
+		closeChan:     make(chan struct{}),
+		newConnCh:     make(chan *conn),
+		IsClose:       false,
+		waitQueueSize: 10, //TODO :In order to be more efficient, this value can be dynamically generated according to the delay algorithm.
 	}
 	//read session by flag
 	go m.readSession()
-	//write session
-	go m.writeSession()
 	//ping
 	go m.ping()
 	return m
@@ -63,7 +60,7 @@ 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)
+	conn := NewConn(s.getId(), s)
 	raw := bytes.NewBuffer([]byte{})
 	if err := binary.Write(raw, binary.LittleEndian, MUX_NEW_CONN); err != nil {
 		return nil, err
@@ -76,10 +73,14 @@ func (s *Mux) NewConn() (*conn, error) {
 	if _, err := s.conn.Write(raw.Bytes()); err != nil {
 		return nil, err
 	}
+	//set a timer timeout 30 second
+	timer := time.NewTimer(time.Second * 30)
+	defer timer.Stop()
 	select {
 	case <-conn.connStatusOkCh:
 		return conn, nil
 	case <-conn.connStatusFailCh:
+	case <-timer.C:
 	}
 	return nil, errors.New("create connection fail，the server refused the connection")
 }
@@ -95,10 +96,24 @@ func (s *Mux) Addr() net.Addr {
 	return s.conn.LocalAddr()
 }
 
+func (s *Mux) sendInfo(flag int32, id int32, content []byte) error {
+	raw := bytes.NewBuffer([]byte{})
+	binary.Write(raw, binary.LittleEndian, flag)
+	binary.Write(raw, binary.LittleEndian, id)
+	if content != nil && len(content) > 0 {
+		binary.Write(raw, binary.LittleEndian, int32(len(content)))
+		binary.Write(raw, binary.LittleEndian, content)
+	}
+	if _, err := s.conn.Write(raw.Bytes()); err != nil || s.pingOk > RETRY_TIME {
+		s.Close()
+		return err
+	}
+	return nil
+}
+
 func (s *Mux) ping() {
 	go func() {
 		ticker := time.NewTicker(time.Second * 5)
-		raw := bytes.NewBuffer([]byte{})
 		for {
 			select {
 			case <-ticker.C:
@@ -107,11 +122,7 @@ func (s *Mux) ping() {
 			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.pingOk > RETRY_TIME {
-				s.Close()
+			if err := s.sendInfo(MUX_PING_FLAG, MUX_PING, nil); err != nil || s.pingOk > RETRY_TIME {
 				break
 			}
 			s.pingOk += 1
@@ -122,45 +133,9 @@ func (s *Mux) ping() {
 	}
 }
 
-func (s *Mux) writeSession() {
-	go func() {
-		raw := bytes.NewBuffer([]byte{})
-		for {
-			raw.Reset()
-			select {
-			case msg := <-s.sendMsgCh:
-				if msg == nil {
-					break
-				}
-				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() {
+	var buf []byte
 	go func() {
-		raw := bytes.NewBuffer([]byte{})
-		buf := pool.BufPoolCopy.Get().([]byte)
-		defer pool.PutBufPoolCopy(buf)
 		for {
 			var flag, i int32
 			var n int
@@ -171,24 +146,19 @@ func (s *Mux) readSession() {
 				}
 				switch flag {
 				case MUX_NEW_CONN: //new conn
-					conn := NewConn(i, s, s.sendMsgCh, s.sendStatusCh)
+					conn := NewConn(i, s)
 					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())
+					s.sendInfo(MUX_NEW_CONN_OK, i, nil)
 					continue
 				case MUX_PING_FLAG: //ping
-					raw.Reset()
-					binary.Write(raw, binary.LittleEndian, MUX_PING_RETURN)
-					binary.Write(raw, binary.LittleEndian, MUX_PING)
-					s.conn.Write(raw.Bytes())
+					s.sendInfo(MUX_PING_RETURN, MUX_PING, nil)
 					continue
 				case MUX_PING_RETURN:
 					s.pingOk -= 1
 					continue
 				case MUX_NEW_MSG:
+					buf = pool.GetBufPoolCopy()
 					if n, err = ReadLenBytes(buf, s.conn); err != nil {
 						break
 					}
@@ -196,20 +166,36 @@ func (s *Mux) readSession() {
 				if conn, ok := s.connMap.Get(i); ok && !conn.isClose {
 					switch flag {
 					case MUX_NEW_MSG: //new msg from remote conn
-						copy(conn.readBuffer, buf[:n])
-						conn.endRead = n
+						//insert wait queue
+						conn.waitQueue.Push(NewBufNode(buf, n))
+						//judge len if >xxx ,send stop
 						if conn.readWait {
+							conn.readWait = false
 							conn.readCh <- struct{}{}
 						}
+						if conn.waitQueue.Size() > s.waitQueueSize {
+							s.sendInfo(MUX_STOP_WRITE, conn.connId, nil)
+						}
+					case MUX_STOP_WRITE:
+						conn.stopWrite = true
 					case MUX_MSG_SEND_OK: //the remote has read
-						conn.getStatusCh <- struct{}{}
+						if conn.stopWrite {
+							conn.stopWrite = false
+							select {
+							case conn.getStatusCh <- struct{}{}:
+							default:
+							}
+						}
 					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()
+						go conn.Close()
+						s.connMap.Delete(i)
 					}
+				} else if flag == MUX_NEW_MSG {
+					pool.PutBufPoolCopy(buf)
 				}
 			} else {
 				break
@@ -231,9 +217,6 @@ func (s *Mux) Close() error {
 	s.closeChan <- struct{}{}
 	s.closeChan <- struct{}{}
 	s.closeChan <- struct{}{}
-	close(s.closeChan)
-	close(s.sendMsgCh)
-	close(s.sendStatusCh)
 	return s.conn.Close()
 }