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(), } }