iotlcm/server/internal/library/websocket/client.go

// Package websocket
// @Link  https://github.com/xiujiecn/xiu-admin
// @Copyright  Copyright (c) 2025 LiXiujie
// @Author  Lxj <li@xiujie.cn>
// @License  https://github.com/xiujiecn/xiu-admin/blob/master/LICENSE
package websocket

import (
	"context"
	"fmt"
	"runtime/debug"

	"git.5iot.com/iotlcm-public/iotlcm-admin-svr/utility/uuid32"

	"github.com/gogf/gf/v2/container/garray"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gorilla/websocket"
)

const (
	// 用户连接超时时间
	heartbeatExpirationTime = 6 * 60
)

// 用户登录
type login struct {
	UserId uint64
	Client *Client
}

// GetKey 读取客户端数据
func (l *login) GetKey() (key string) {
	key = GetUserKey(l.UserId)
	return
}

// Client 客户端连接
type Client struct {
	Addr          string           // 客户端地址
	ID            string           // 连接唯一标识
	Socket        *websocket.Conn  // 用户连接
	Send          chan *WResponse  // 待发送的数据
	SendClose     bool             // 发送是否关闭
	UserId        uint64           // 用户ID，用户登录以后才有
	FirstTime     uint64           // 首次连接事件
	HeartbeatTime uint64           // 用户上次心跳时间
	LoginTime     uint64           // 登录时间 登录以后才有
	isApp         bool             // 是否是app
	tags          *garray.StrArray // 标签
	Ctx           context.Context  // 上下文
	closeSignal   chan struct{}    // 关闭信号
}

// NewClient 初始化
func NewClient(addr string, socket *websocket.Conn, firstTime uint64, ctx context.Context) (client *Client) {
	client = &Client{
		Addr:          addr,
		ID:            uuid32.New(),
		Socket:        socket,
		Send:          make(chan *WResponse, 100),
		SendClose:     false,
		FirstTime:     firstTime,
		HeartbeatTime: firstTime,
		Ctx:           ctx,
		closeSignal:   make(chan struct{}, 1),
		tags:          garray.NewStrArray(true),
	}
	return
}

// 读取客户端数据
func (c *Client) read() {
	defer func() {
		if r := recover(); r != nil {
			fmt.Println("write stop", string(debug.Stack()), r)
		}
	}()

	defer func() {
		c.close()
	}()

	for {
		_, message, err := c.Socket.ReadMessage()
		if err != nil {
			g.Log().Info(c.Ctx, "websocket Client.read failed.", err)
			return
		}
		// 处理程序
		// fmt.Println(message)
		ProcessData(c, message)
	}
}

// 向客户端写数据
func (c *Client) write() {
	defer func() {
		if r := recover(); r != nil {
			fmt.Println("write stop", string(debug.Stack()), r)
		}
	}()
	defer func() {
		clientManager.Unregister <- c
		_ = c.Socket.Close()
	}()
	for {
		select {
		case <-c.closeSignal:
			g.Log().Infof(mctx, "websocket client exit, user:%+v", c.UserId)
			return
		case message, ok := <-c.Send:
			if !ok {
				g.Log().Info(c.Ctx, "websocket Client.write failed.")
				// 发送数据错误 关闭连接
				return
			}
			_ = c.Socket.WriteJSON(message)
		}
	}
}

// SendMsg 发送数据
func (c *Client) SendMsg(msg *WResponse) {
	if c == nil || c.SendClose {
		return
	}
	defer func() {
		if r := recover(); r != nil {
			fmt.Println("SendMsg stop:", r, string(debug.Stack()))
		}
	}()
	c.Send <- msg
}

// Heartbeat 心跳更新
func (c *Client) Heartbeat(currentTime uint64) {
	c.HeartbeatTime = currentTime
}

// IsHeartbeatTimeout 心跳是否超时
func (c *Client) IsHeartbeatTimeout(currentTime uint64) (timeout bool) {
	if c.HeartbeatTime+heartbeatExpirationTime <= currentTime {
		timeout = true
	}
	return
}

// 关闭客户端
func (c *Client) close() {
	g.Log().Info(c.Ctx, "websocket Client.close")
	if c.SendClose {
		return
	}
	c.SendClose = true
	c.closeSignal <- struct{}{}
}