dy-admin-temp/internal/pcmserver/bus/service/home_client.go

package service

import (
	"dy-admin/internal/pcmserver/global"
	"dy-admin/internal/pcmserver/pkg/code"
	systemService "dy-admin/internal/pcmserver/sys/service"
	"dy-admin/pkg/jwt"
	"dy-admin/pkg/log"
	"encoding/json"
	"fmt"
	"github.com/gorilla/websocket"
	"go.uber.org/zap"
	"time"
)

type HomeClient struct {
	//当前客户端连接
	clientConn *websocket.Conn
	//客户端连接地址
	address string
	//写入当前连接的消息通道
	send chan *HomeMsg
	// 连接时的deptID
	deptId int
	// 本部门以及子部门
	deptWithSon map[int]struct{}
}

func NewHomeClient(conn *websocket.Conn) (*HomeClient, error) {
	if conn == nil {
		return nil, fmt.Errorf("websocket conn is nil")
	}

	hc := &HomeClient{
		clientConn: conn,
		address:    conn.RemoteAddr().String(),
		send:       make(chan *HomeMsg, 1024),
	}
	homeClientHub.clientRegister <- hc

	return hc, nil
}

func (hc *HomeClient) ReadMessage() {
	defer func() {
		// 断开连接,关闭连接
		if hc.clientConn != nil {
			_ = hc.clientConn.Close()
			homeClientHub.clientUnRegister <- hc
		}
	}()

	for {
		msgType, msg, err := hc.clientConn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Warn("readMessage err", zap.String("address", hc.address), zap.Error(err))
			} else {
				log.Info("readMessage conn closed", zap.String("address", hc.address))
			}
			break
		}
		log.Info("read message", zap.Int("msgType", msgType), zap.String("msg", string(msg)))
		receiveData := &HomeMsg{}
		err = json.Unmarshal(msg, receiveData)
		if err != nil {
			log.Error("msg unmarshal err", zap.String("address", hc.address), zap.Error(err))
		}
		switch receiveData.MsgType {
		case ping:
			sendData := &HomeMsg{DeptID: 0, MsgType: pong}

			hc.send <- sendData
		case wsLeaveReq:
			// 解析数据
			var leaveReq WSLeaveReq
			err = json.Unmarshal([]byte(receiveData.Data), &leaveReq)
			if err != nil {
				log.Error(code.ErrCodeMap[code.ErrTokenInvalidCode].Error())
				continue
			}
			// 发送请假详情
			err = sendLeaveInfo(leaveReq.StaffID)
			if err != nil {
				log.Error(err.Error())
				continue
			}
		case clientToken:
			// 解析客户端发来的token
			claims, err := global.JWT.ParseToken(receiveData.Data)
			if err != nil {
				if err == jwt.TokenExpired {
					sendData := &HomeMsg{DeptID: 0, MsgType: clientToken, Data: code.ErrCodeMap[code.ErrTokenExpiredCode].Error()}
					hc.send <- sendData
					log.Error(code.ErrCodeMap[code.ErrTokenExpiredCode].Error())
					continue
				}
				sendData := &HomeMsg{DeptID: 0, MsgType: clientToken, Data: code.ErrCodeMap[code.ErrTokenInvalidCode].Error()}
				hc.send <- sendData
				log.Error(code.ErrCodeMap[code.ErrTokenInvalidCode].Error())
				continue
			}
			deptIds, err := systemService.ServicesGroupApp.GetSonDeptIds(claims.DeptID)
			if err != nil {
				sendData := &HomeMsg{DeptID: 0, MsgType: clientToken, Data: err.Error()}
				hc.send <- sendData
				log.Error(err.Error())
				continue
			}
			deptMap := make(map[int]struct{})
			for _, dep := range deptIds {
				deptMap[dep] = struct{}{}
			}
			hc.deptId = claims.DeptID
			hc.deptWithSon = deptMap

			// 推送首页数据
			hc.PushHomePageData()
		}
	}
}

func (hc *HomeClient) WriteMessage() {
	defer func() {
		if hc.clientConn != nil {
			err := hc.clientConn.Close()
			if err != nil {
				log.Warn("writeMessage defer in  client err", zap.String("address", hc.address), zap.Error(err))
			}
			log.Info("writeMessage defer in  client conn closed", zap.String("address", hc.address))
		}
	}()

	for {
		select {
		case homeMsg, ok := <-hc.send:
			if !ok {
				return
			}
			message, err := json.Marshal(homeMsg)
			if err != nil {
				log.Error("json.Marshal", zap.Error(err))
				continue
			}
			//将消息管道的消息，写入当前连接
			err = hc.clientConn.WriteMessage(websocket.BinaryMessage, message)
			if err != nil {
				log.Error("clientConn WriteMessage error", zap.String("address", hc.address), zap.Error(err))
				return
			}
			log.Debug("send message success", zap.String("address", hc.address), zap.String("msg", string(message)))
		}
	}
}

func (hc *HomeClient) SendHello() {
	for {
		hc.send <- &HomeMsg{
			DeptID:  0,
			MsgType: ping,
			Data:    "hello",
		}
		time.Sleep(time.Second * 10)
	}
}

func (hc *HomeClient) PushHomePageData() {
	if hc.deptId == 0 {
		return
	}
	var err error
	// 部门树数据
	if err = sendDeptTree(hc.deptId, false, hc); err != nil {
		log.Error("sendDeptTree", zap.Error(err))
	}

	// 获取所有子部门
	deptIds, err := systemService.ServicesGroupApp.DeptService.GetSonDeptIds(hc.deptId)
	if err != nil {
		log.Error("GetSonDeptIds", zap.Error(err))
	}

	if deptIds == nil || len(deptIds) == 0 {
		return
	}

	// 发送每个部门的柜子信息
	for _, id := range deptIds {
		// 发送每个部门的存物统计
		if err = sendDeptStoneInfo(id, false, hc); err != nil {
			log.Error("sendDeptStoneInfo", zap.Error(err))
		}
		// 发送每个部门的柜子连接统计
		if err = sendCabinetConnect(id, false, hc); err != nil {
			log.Error("sendCabinetConnect", zap.Error(err))
		}
		// 发送每个部门的柜子信息
		if err = sendCabinetInfo(id, false, hc); err != nil {
			log.Error("sendCabinetInfo", zap.Error(err))
		}
		// 发送每个部门的预警信息
		if err = sendDeptWarnLog(id, false, hc); err != nil {
			log.Error("sendDeptWarnLog", zap.Error(err))
		}
	}

	log.Info("PushHomePageData success", zap.Int("dept_id", hc.deptId))
}