service-common/net/socket/websocket/websocket.go

package websocket

import (
	"bytes"
	"context"
	"git.hlsq.asia/mmorpg/service-common/net/socket"
	"github.com/panjf2000/gnet/v2"
	"github.com/panjf2000/gnet/v2/pkg/logging"
	"sync"
	"time"
)

// WSServer 实现GNet库接口
type WSServer struct {
	eng            gnet.Engine
	i              socket.ISocketServer
	logger         logging.Logger // 日志
	upgradeTimeout time.Duration  // 升级超时时间
	unUpgradeConn  sync.Map
}

func NewWSServer(i socket.ISocketServer, logger logging.Logger, upgradeTimeout time.Duration) *WSServer {
	if i == nil {
		return nil
	}
	return &WSServer{
		i:              i,
		logger:         logger,
		upgradeTimeout: upgradeTimeout,
		unUpgradeConn:  sync.Map{},
	}
}

func (s *WSServer) Run(addr string, multicore bool, numEventLoop int, tcpNoDelay gnet.TCPSocketOpt, readBufferCap, writeBufferCap int, lockOSThread bool, reusePort, ticker bool, logger logging.Logger) error {
	return gnet.Run(
		s,
		addr,
		gnet.WithMulticore(multicore),
		gnet.WithNumEventLoop(numEventLoop),
		gnet.WithTCPNoDelay(tcpNoDelay),
		gnet.WithReadBufferCap(readBufferCap),
		gnet.WithWriteBufferCap(writeBufferCap),
		gnet.WithLockOSThread(lockOSThread),
		gnet.WithReusePort(reusePort),
		gnet.WithTicker(ticker),
		gnet.WithLogger(logger),
	)
}

func (s *WSServer) Stop() error {
	return s.eng.Stop(context.Background())
}

func (s *WSServer) OnBoot(eng gnet.Engine) gnet.Action {
	s.eng = eng
	return gnet.None
}

func (s *WSServer) OnShutdown(_ gnet.Engine) {
}

func (s *WSServer) OnOpen(c gnet.Conn) ([]byte, gnet.Action) {
	ws := &WSConn{
		Conn:      c,
		isUpgrade: false,
		openTime:  time.Now().Unix(),
		buf:       bytes.Buffer{},
		logger:    s.logger,
		wsMessageBuf: wsMessageBuf{
			curHeader: nil,
			cachedBuf: bytes.Buffer{},
		},
		param:      make(map[string]interface{}),
		remoteAddr: c.RemoteAddr().String(),
	}
	c.SetContext(ws)
	s.unUpgradeConn.Store(c.RemoteAddr().String(), ws)
	d, a := s.i.OnOpen(ws)
	return d, (gnet.Action)(a)
}

// OnClose fires when a connection has been closed.
// The parameter err is the last known connection error.
func (s *WSServer) OnClose(c gnet.Conn, err error) (action gnet.Action) {
	s.unUpgradeConn.Delete(c.RemoteAddr().String())
	ws, ok := c.Context().(*WSConn)
	if ok {
		ws.isClose = true
		ws.logger.Warnf("connection close, err: %v", err)
		return gnet.Action(s.i.OnClose(ws, err))
	}
	return
}

// OnTraffic fires when a local socket receives data from the peer.
func (s *WSServer) OnTraffic(c gnet.Conn) gnet.Action {
	ws := c.Context().(*WSConn)
	if action := ws.readBytesBuf(c); action != gnet.None {
		return action
	}

	if !ws.isUpgrade {
		data, action := ws.upgrade()
		if len(data) > 0 {
			s.unUpgradeConn.Delete(c.RemoteAddr().String())
			action = gnet.Action(s.i.OnHandShake(ws))
			if action == gnet.None {
				if err := ws.Conn.AsyncWrite(data, nil); err != nil {
					ws.logger.Errorf("OnTraffic upgrade AsyncWrite err: %v", err)
					if err = ws.Close(); err != nil {
						ws.logger.Errorf("OnTraffic upgrade Close error: %v", err)
					}
				}
			}
		}
		return action
	}

	if msg, err := ws.readWsMessages(); err != nil {
		ws.logger.Errorf("read ws messages err: %v", err)
		return gnet.Close
	} else if msg != nil {
		for _, m := range msg {
			if socket.OpCode(m.OpCode) == socket.OpPong {
				s.i.OnPong(ws)
				continue
			}
			if socket.OpCode(m.OpCode) == socket.OpClose {
				return gnet.Close
			}
			if socket.OpCode(m.OpCode) == socket.OpPing {
				continue
			}
			a := s.i.OnMessage(ws, m.Payload)
			if gnet.Action(a) != gnet.None {
				return gnet.Action(a)
			}
		}
	}
	return gnet.None
}

// OnTick fires immediately after the engine starts and will fire again
// following the duration specified by the delay return value.
func (s *WSServer) OnTick() (delay time.Duration, action gnet.Action) {
	now := time.Now().Unix()

	delConn := make([]string, 0)
	s.unUpgradeConn.Range(func(key, value interface{}) bool {
		k, ok := key.(string)
		if !ok {
			return true
		}

		v, ok := value.(*WSConn)
		if !ok {
			return true
		}
		if now-v.openTime >= int64(s.upgradeTimeout.Seconds()) {
			delConn = append(delConn, k)
		}
		return true
	})

	for _, k := range delConn {
		wsConn, _ := s.unUpgradeConn.LoadAndDelete(k)
		if wsConn == nil {
			continue
		}

		v, ok := wsConn.(*WSConn)
		if !ok {
			continue
		}
		if err := v.Close(); err != nil {
			v.logger.Errorf("upgrade ws time out close socket error: %v", err)
		}
	}
	d, a := s.i.OnTick()
	delay = d
	action = gnet.Action(a)

	return
}