nps/component/client/client.go

package client

import (
	"ehang.io/nps/core/action"
	"ehang.io/nps/core/handler"
	"ehang.io/nps/core/process"
	"ehang.io/nps/core/rule"
	"ehang.io/nps/lib/enet"
	"ehang.io/nps/lib/logger"
	"go.uber.org/zap"
	"io"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

type Client struct {
	controlLn    net.Listener
	dataLn       net.Listener
	lastPongTime time.Time
	mux          *http.ServeMux
	ticker       *time.Ticker
	closeCh      chan struct{}
	closed       int32
	wg           sync.WaitGroup
}

func NewClient(controlLn, dataLn net.Listener) *Client {
	return &Client{
		controlLn: controlLn,
		dataLn:    dataLn,
		mux:       &http.ServeMux{},
		ticker:    time.NewTicker(time.Second * 5),
		closeCh:   make(chan struct{}, 0),
	}
}

func (c *Client) ping(writer http.ResponseWriter, request *http.Request) {
	c.lastPongTime = time.Now()
	_, err := io.WriteString(writer, "pong")
	if err != nil {
		logger.Warn("write pong error", zap.Error(err))
		return
	}
	logger.Debug("write pong success")
}

func (c *Client) Run() {
	c.mux.HandleFunc("/ping", c.ping)
	c.wg.Add(3)
	go c.handleControlConn()
	go c.handleDataConn()
	go c.checkPing()
	c.wg.Wait()
}

func (c *Client) HasPong() bool {
	return time.Now().Sub(c.lastPongTime).Seconds() < 10
}

func (c *Client) checkPing() {
	for {
		select {
		case <-c.ticker.C:
			if !c.lastPongTime.IsZero() && time.Now().Sub(c.lastPongTime).Seconds() > 15 && c.controlLn != nil {
				logger.Debug("close connection", zap.Time("lastPongTime", c.lastPongTime), zap.Time("now", time.Now()))
				_ = c.controlLn.Close()
			}
		case <-c.closeCh:
			c.wg.Done()
			return
		}
	}
}

func (c *Client) handleDataConn() {
	h := &handler.DefaultHandler{}
	ac := &action.LocalAction{}
	err := ac.Init()
	if err != nil {
		logger.Warn("init action failed", zap.Error(err))
		return
	}
	appPr := &process.PbAppProcessor{}
	_ = appPr.Init(ac)
	h.AddRule(&rule.Rule{Handler: h, Process: appPr, Action: ac})

	pingPr := &process.PbPingProcessor{}
	_ = appPr.Init(ac)
	h.AddRule(&rule.Rule{Handler: h, Process: pingPr, Action: ac})

	var conn net.Conn
	for {
		conn, err = c.dataLn.Accept()
		if err != nil {
			logger.Error("accept connection failed", zap.Error(err))
			break
		}
		go func(conn net.Conn) {
			_, err = h.HandleConn(nil, enet.NewReaderConn(conn))
			if err != nil {
				logger.Warn("process failed", zap.Error(err))
				return
			}
		}(conn)
	}
	c.wg.Done()
	c.Close()
}

func (c *Client) handleControlConn() {
	err := http.Serve(c.controlLn, c.mux)
	if err != nil {
		logger.Error("http error", zap.Error(err))
	}
	c.wg.Done()
	c.Close()
}

func (c *Client) Close() {
	if atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
		c.closeCh <- struct{}{}
		c.ticker.Stop()
		_ = c.controlLn.Close()
		_ = c.dataLn.Close()
	}
}