natpass/code/client/pool/conn.go

package pool

import (
	"context"
	"natpass/code/network"
	"natpass/code/utils"
	"strings"
	"sync"
	"time"

	"github.com/lwch/logging"
)

// Conn pool connection
type Conn struct {
	sync.RWMutex
	Idx          uint32
	ReadTimeout  time.Duration
	WriteTimeout time.Duration
	parent       *Pool
	conn         *network.Conn
	read         map[string]chan *network.Msg // link id => channel
	unknownRead  chan *network.Msg            // read message without link
	write        chan *network.Msg            // link id => channel
}

func newConn(parent *Pool, conn *network.Conn, idx uint32) *Conn {
	ret := &Conn{
		Idx:          idx,
		ReadTimeout:  parent.cfg.ReadTimeout,
		WriteTimeout: parent.cfg.WriteTimeout,
		parent:       parent,
		conn:         conn,
		read:         make(map[string]chan *network.Msg),
		unknownRead:  make(chan *network.Msg),
		write:        make(chan *network.Msg),
	}
	logging.Info("new connection: %s-%d", ret.parent.cfg.ID, ret.Idx)
	ctx, cancel := context.WithCancel(context.Background())
	go ret.loopRead(cancel)
	go ret.loopWrite(cancel)
	go ret.keepalive(ctx)
	return ret
}

func (conn *Conn) hasLink(id string) bool {
	conn.RLock()
	defer conn.RUnlock()
	_, ok := conn.read[id]
	return ok
}

// AddLink attach read message
func (conn *Conn) AddLink(id string) {
	logging.Info("add link %s from %d", id, conn.Idx)
	conn.Lock()
	if _, ok := conn.read[id]; !ok {
		conn.read[id] = make(chan *network.Msg)
	}
	conn.Unlock()
}

// RemoveLink detach read message
func (conn *Conn) RemoveLink(id string) {
	logging.Info("remove link %s from %d", id, conn.Idx)
	conn.Lock()
	ch := conn.read[id]
	if ch != nil {
		close(ch)
	}
	delete(conn.read, id)
	conn.Unlock()
}

// Close close connection
func (conn *Conn) Close() {
	conn.conn.Close()
	conn.Lock()
	for id, ch := range conn.read {
		close(ch)
		delete(conn.read, id)
	}
	conn.Unlock()
	if conn.unknownRead != nil {
		close(conn.unknownRead)
		conn.unknownRead = nil
	}
	if conn.write != nil {
		close(conn.write)
		conn.write = nil
	}
	conn.parent.onClose(conn.Idx)
	logging.Error("connection %s-%d closed", conn.parent.cfg.ID, conn.Idx)
}

func (conn *Conn) loopRead(cancel context.CancelFunc) {
	defer utils.Recover("loopRead")
	defer conn.Close()
	defer cancel()
	var timeout int
	for {
		msg, _, err := conn.conn.ReadMessage(conn.parent.cfg.ReadTimeout)
		if err != nil {
			if strings.Contains(err.Error(), "i/o timeout") {
				timeout++
				if timeout >= 60 {
					logging.Error("too many timeout times")
					return
				}
				continue
			}
			logging.Error("read message: %v", err)
			return
		}
		timeout = 0
		if msg.GetXType() == network.Msg_keepalive {
			continue
		}
		linkID := msg.GetLinkId()
		conn.RLock()
		ch := conn.read[linkID]
		conn.RUnlock()
		if ch == nil {
			ch = conn.unknownRead
		}
		select {
		case ch <- msg:
		case <-time.After(conn.WriteTimeout):
		}
	}
}

func (conn *Conn) loopWrite(cancel context.CancelFunc) {
	defer utils.Recover("loopWrite")
	defer conn.Close()
	defer cancel()
	for {
		msg := <-conn.write
		if msg == nil {
			return
		}
		msg.From = conn.parent.cfg.ID
		msg.FromIdx = conn.Idx
		err := conn.conn.WriteMessage(msg, conn.parent.cfg.WriteTimeout)
		if err != nil {
			logging.Error("write message error on %s-%d: %v",
				conn.parent.cfg.ID, conn.Idx, err)
			return
		}
	}
}

// ChanRead get read channel from link id
func (conn *Conn) ChanRead(id string) <-chan *network.Msg {
	conn.RLock()
	defer conn.RUnlock()
	return conn.read[id]
}

// Reset reset message next read
func (conn *Conn) Reset(id string, msg *network.Msg) {
	conn.RLock()
	ch := conn.read[id]
	conn.RUnlock()
	ch <- msg
}

// ChanUnknown get channel of unknown link id
func (conn *Conn) ChanUnknown() <-chan *network.Msg {
	return conn.unknownRead
}

func (conn *Conn) keepalive(ctx context.Context) {
	defer utils.Recover("keepalive")
	for {
		select {
		case <-ctx.Done():
			return
		case <-time.After(10 * time.Second):
			conn.SendKeepalive()
		}
	}
}