slow_wild_ws/internal/server/websocketserver.go

package server

import (
	"context"
	"fmt"
	"net/http"
	"slowwildws/internal/config"
	"slowwildws/internal/constants"
	"slowwildws/internal/types"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/core/threading"
)

type WebsocketServer struct {
	sync.RWMutex // 读写锁
	connToUser   map[*ConnectionServer]int64
	userToConn   map[int64]*ConnectionServer
	upgrader     websocket.Upgrader
	logx.Logger
	opt *serverOption
	*threading.TaskRunner
	ctx           context.Context
	messageServer *MessageServer
}

func NewWebsockerServer(conf config.Config, opts ...ServerOptions) *WebsocketServer {
	opt := newServerOptions(conf, opts...)
	return &WebsocketServer{
		upgrader: websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		},
		Logger:        logx.WithContext(context.Background()),
		connToUser:    make(map[*ConnectionServer]int64),
		userToConn:    make(map[int64]*ConnectionServer),
		opt:           &opt,
		TaskRunner:    threading.NewTaskRunner(opt.concurrency),
		messageServer: NewMessageServer(),
	}
}

// 初始化一个连接信息
func (s *WebsocketServer) InitConn(ctx context.Context, conn *websocket.Conn, conf config.Config) *ConnectionServer {
	return NewConnectionServer(ctx, conf, conn)
}

// 把连接添加到map中
func (s *WebsocketServer) AddConn(conn *ConnectionServer, ctx context.Context) {
	uid := ctx.Value(constants.UserID).(int64)

	fmt.Println("获取到的用户id： ", uid)

	s.RWMutex.Lock()
	defer s.RWMutex.Unlock()
	// 验证用户是否登录过
	if c := s.userToConn[uid]; c != nil {
		c.Close()
	}
	s.connToUser[conn] = uid
	s.userToConn[uid] = conn
	// todo 是否应该再这里加上用户在线标识
}

// 获取存储的用户id
func (s *WebsocketServer) GetUserId(conn *ConnectionServer) int64 {
	s.RWMutex.RLock()
	defer s.RWMutex.RUnlock()

	if conn == nil {
		return 0
	}

	return s.connToUser[conn]
}

// 关闭这个连接服务
func (s *WebsocketServer) Close(conn *ConnectionServer) {
	s.RWMutex.Lock()
	defer s.RWMutex.Unlock()
	uid := s.connToUser[conn]
	if uid == 0 {
		// 已经关闭了
		return
	}
	delete(s.connToUser, conn)
	delete(s.userToConn, uid)
	conn.Close()
}

// 判断是否需要ack确认
func (s *WebsocketServer) IsAck(message *types.Message) bool {
	if message == nil {
		return s.opt.ack != constants.NoAck
	}
	if message.FormID == constants.SYSTEM_ROOT_UID {
		//超级管理员不需要这个ack确认机制，直接发
		return false
	}
	return s.opt.ack != constants.NoAck && message.FrameType != constants.FrameNoAck
}

// 读取消息的ack
func (s *WebsocketServer) ReadAck(conn *ConnectionServer) {
	for {
		select {
		case <-conn.done:
			s.Infof("close message ack uid %v", conn.Uid)
			return
		default:

		}
		conn.messageMu.Lock()
		if len(conn.readMessageList) == 0 {
			conn.messageMu.Unlock()
			// 没有消息的话让其休眠100毫秒再进行下一次判定
			time.Sleep(100 * time.Millisecond)
			continue
		}

		// 读取第一条消息
		message := conn.readMessageList[0]
		//判断ack的方式
		switch s.opt.ack {
		case constants.OnlyAck:
			// 直接给客户端回复
			s.messageServer.Send(&types.Message{
				FrameType: constants.FrameAck,
				Id:        message.Id,
				AckSeq:    message.AckSeq + 1,
			}, conn)
			// 进行业务处理，从队列中移除
			conn.readMessageList = conn.readMessageList[1:]
			conn.messageMu.Unlock()
			conn.message <- message
		case constants.RigorAck:
			if message.AckSeq == 0 {
				// 还没有确认
				conn.readMessageList[0].AckSeq++             //记录确认序号加1
				conn.readMessageList[0].AckTime = time.Now() // 记录确认时间
				s.messageServer.Send(&types.Message{
					FrameType: constants.FrameAck,
					Id:        message.Id,
					AckSeq:    message.AckSeq + 1,
				}, conn)
				s.Infof("message ack RigorAck send mid %v, seq %v, time %v", message.Id, message.AckSeq, message.AckTime)
				conn.messageMu.Unlock()
				continue
			}

			// 再验证

			// 1.客户端返回结果，再一次确认
			msgSeq := conn.readMessageSeq[message.Id]
			if msgSeq.AckSeq > message.AckSeq {
				// 进行业务处理，从队列中移除
				conn.readMessageList = conn.readMessageList[1:]
				conn.messageMu.Unlock()
				conn.message <- message
				s.Infof("message ack rigorAck success mid %v", message.Id)
				continue
			}
			// 2. 客户端没有确认，考虑是否再次发送确认消息
			val := s.opt.ackTimeout - time.Since(message.AckTime)
			fmt.Println("超时时间： ", val)
			// 2.1 超过结束确认，抛出错误
			if !message.AckTime.IsZero() && val <= 0 {
				delete(conn.readMessageSeq, message.Id)
				conn.readMessageList = conn.readMessageList[1:]
				conn.messageMu.Unlock()
				continue
			}
			// 2.2 未超过，超新发送
			conn.messageMu.Unlock()
			s.messageServer.Send(&types.Message{
				FrameType: constants.FrameAck,
				Id:        message.Id,
				AckSeq:    message.AckSeq,
			}, conn)
			time.Sleep(3 * time.Second)
		}
	}

}

// 任务处理,读取消息体中的method方法，判断是否需要执行
func (s *WebsocketServer) HandleWrite(conn *ConnectionServer) {
	for {
		select {
		case <-conn.done:
			//连接关闭
			return
		case message := <-conn.message:
			switch message.FrameType {
			case constants.FramePing:
				s.messageServer.Send(&types.Message{FrameType: constants.FramePing}, conn)
			case constants.FrameData, constants.FrameNoAck:
				if handler, ok := message.Method; ok {
					handler(s, conn, message)
				} else {
					s.messageServer.Send(&types.Message{FrameType: constants.FrameData, Data: []byte(fmt.Sprintf("不存在执行的方法 %v 请检查", message.Method))}, conn)
				}
			}
			if s.opt.ack != constants.NoAck {
				// 删除ack消息的序号记录
				conn.messageMu.Lock()
				delete(conn.readMessageSeq, message.Id)
				conn.messageMu.Unlock()
			}
		}

	}
}

func (s *WebsocketServer) Send(msg interface{}, conns ...*ConnectionServer) error {
	return s.messageServer.Send(msg, conns...)
}

func (s *WebsocketServer) NewErrMessage(err error) *types.Message {
	return &types.Message{
		FrameType: constants.FrameData,
		Data:      err.Error(),
	}
}