package gateway import ( "context" "fmt" "go_dreamfactory/comm" "go_dreamfactory/pb" "sync" "sync/atomic" "go_dreamfactory/lego/sys/log" "go_dreamfactory/lego/utils/container/id" "github.com/gorilla/websocket" "google.golang.org/protobuf/proto" ) /* 用户代理对象 封装用户socket 对象 处理用户的消息读取 写入 关闭等操作 */ func newAgent(gateway IGateway, conn *websocket.Conn) *Agent { agent := &Agent{ gateway: gateway, wsConn: conn, sessionId: id.NewUUId(), uId: "", writeChan: make(chan *pb.UserMessage, 2), closeSignal: make(chan bool), state: 1, } agent.wg.Add(2) go agent.readLoop() go agent.writeLoop() return agent } //用户代理 type Agent struct { gateway IGateway wsConn *websocket.Conn sessionId string uId string writeChan chan *pb.UserMessage closeSignal chan bool state int32 //状态 0 关闭 1 运行 2 关闭中 wg sync.WaitGroup } func (this *Agent) readLoop() { defer this.wg.Done() var ( data []byte msg *pb.UserMessage = &pb.UserMessage{} err error ) locp: for { if _, data, err = this.wsConn.ReadMessage(); err != nil { log.Errorf("agent:%s uId:%s ReadMessage err:%v", this.sessionId, this.uId, err) go this.Close() break locp } if err = proto.Unmarshal(data, msg); err != nil { log.Errorf("agent:%s uId:%s Unmarshal err:%v", this.sessionId, this.uId, err) go this.Close() break locp } else { this.messageDistribution(msg) } } log.Debugf("agent:%s uId:%s readLoop end!", this.sessionId, this.uId) } func (this *Agent) writeLoop() { defer this.wg.Done() var ( data []byte err error ) locp: for { select { case <-this.closeSignal: break locp case msg, ok := <-this.writeChan: if ok { data, err = proto.Marshal(msg) if err = this.wsConn.WriteMessage(websocket.BinaryMessage, data); err != nil { log.Errorf("agent:%s uId:%d WriteMessage err:%v", this.sessionId, this.uId, err) go this.Close() } } else { go this.Close() } } } log.Debugf("agent:%s uId:%s writeLoop end!", this.sessionId, this.uId) } func (this *Agent) SessionId() string { return this.sessionId } func (this *Agent) IP() string { return this.wsConn.RemoteAddr().String() } func (this *Agent) UserId() string { return this.uId } func (this *Agent) Build(uId string) { this.uId = uId } func (this *Agent) UnBuild() { this.uId = "" } func (this *Agent) WriteMsg(msg *pb.UserMessage) (err error) { if atomic.LoadInt32(&this.state) != 1 { return } this.writeChan <- msg return } //外部代用关闭 func (this *Agent) Close() { if !atomic.CompareAndSwapInt32(&this.state, 1, 2) { return } this.wsConn.Close() this.closeSignal <- true this.wg.Wait() atomic.StoreInt32(&this.state, 0) this.gateway.DisConnect(this) } //分发用户消息 func (this *Agent) messageDistribution(msg *pb.UserMessage) { reply := &pb.RPCMessageReply{} log.Debugf("agent:%s uId:%s MessageDistribution msg:%s.%s", this.sessionId, this.uId, msg.MainType, msg.SubType) if err := this.gateway.Service().RpcCallByType("worker", string(comm.Rpc_GatewayRoute), context.Background(), &pb.AgentMessage{ Ip: this.IP(), UserSessionId: this.sessionId, UserId: this.uId, GatewayServiceId: this.gateway.Service().GetId(), Method: fmt.Sprintf("%s.%s", msg.MainType, msg.SubType), Message: msg.Data, }, reply); err != nil { log.Errorf("agent:%s uId:%s MessageDistribution err:%v", this.sessionId, this.uId, err) } else { log.Debugf("agent:%s uId:%s MessageDistribution reply:%v", this.sessionId, this.uId, reply) } }