ai_scheduler/internal/gateway/client.go

package gateway

import (
	errors "ai_scheduler/internal/data/error"
	"ai_scheduler/internal/data/model"
	"context"
	"log"
	"math/rand"
	"sync"
	"time"

	"github.com/google/uuid"

	"github.com/gofiber/websocket/v2"
)

var (
	ErrConnClosed = errors.SysErrf("连接不存在或已关闭")
	rng           = rand.New(rand.NewSource(time.Now().UnixNano()))
	idBuf         = make([]byte, 20)
)

type Client struct {
	id         string          // 客户端唯一ID
	conn       *websocket.Conn // WebSocket 连接
	session    string          // 会话ID
	key        string          // 应用密钥
	auth       string          // 用户凭证token
	codes      []string        // 用户权限code
	sysInfo    *model.AiSy     // 系统信息
	tasks      []model.AiTask  // 任务列表
	sysCode    string          // 系统编码
	Ctx        context.Context
	Cancel     context.CancelFunc
	LastActive time.Time
	mu         sync.Mutex
}

func NewClient(conn *websocket.Conn, ctx context.Context, cancel context.CancelFunc) *Client {

	return &Client{
		id:     generateClientID(),
		conn:   conn,
		Ctx:    ctx,
		Cancel: cancel,
		mu:     sync.Mutex{},
	}
}

// GetID 获取客户端的唯一ID
func (c *Client) GetID() string {
	return c.id
}

// GetConn 获取客户端的 WebSocket 连接
func (c *Client) GetConn() *websocket.Conn {
	return c.conn
}

// GetSession 获取会话ID
func (c *Client) GetSession() string {
	return c.session
}

// GetKey 获取应用密钥
func (c *Client) GetKey() string {
	return c.key
}

// GetAuth 获取用户凭证token
func (c *Client) GetAuth() string {
	return c.auth
}

// GetCodes 获取用户权限code
func (c *Client) GetCodes() []string {
	return c.codes
}

// 获取系统编码
func (c *Client) GetSysCode() string {
	return c.sysCode
}

// GetSysInfo 获取系统信息
func (c *Client) GetSysInfo() *model.AiSy {
	return c.sysInfo
}

// SetSysInfo 设置系统信息
func (c *Client) SetSysInfo(sysInfo *model.AiSy) {
	c.sysInfo = sysInfo
}

// GetTasks 获取任务列表
func (c *Client) GetTasks() []model.AiTask {
	return c.tasks
}

// SetTasks 设置任务列表
func (c *Client) SetTasks(tasks []model.AiTask) {
	c.tasks = tasks
}

// 设置用户权限code
func (c *Client) SetCodes(codes []string) {
	c.codes = codes
}

// Close 关闭客户端连接
func (c *Client) Close() {
	//c.mu.Lock()
	//defer c.mu.Unlock()
	if c.conn != nil {
		_ = c.conn.Close()
		c.conn = nil
	}
}

// SendFunc 发送消息到客户端
func (c *Client) SendFunc(msg []byte) error {
	return c.SendMessage(websocket.TextMessage, msg)
}

// 在Client结构体中添加更详细的日志
func (c *Client) SendMessage(msgType int, msg []byte) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.conn == nil {
		return ErrConnClosed
	}

	err := c.conn.WriteMessage(msgType, msg)
	if err != nil {
		log.Printf("发送消息失败: %v, 客户端ID: %s, 消息类型: %d",
			err, c.id, msgType)
	}
	return err
}

// 生成唯一的客户端ID
func generateClientID() string {
	return uuid.New().String()
	//// 1. 时间戳
	//timestamp := time.Now().UnixNano()
	//binary.BigEndian.PutUint64(idBuf[:8], uint64(timestamp))
	//
	//// 2. 随机数（4字节）
	//binary.BigEndian.PutUint32(idBuf[8:12], rng.Uint32())
	//
	//// 3. 十六进制编码
	//n := pkg.HexEncode(idBuf[:12], idBuf[12:])
	//return string(idBuf[12 : 12+n])
}

// 连接数据验证和收集
func (c *Client) DataAuth() (err error) {
	c.session = c.conn.Query("x-session", "")
	if len(c.session) == 0 {
		err = errors.SessionNotFound
		return
	}
	c.auth = c.conn.Query("x-authorization", "")
	if len(c.auth) == 0 {
		err = errors.AuthNotFound
		return
	}
	c.key = c.conn.Query("x-app-key", "")
	if len(c.key) == 0 {
		err = errors.KeyNotFound
		return
	}
	// 系统编码
	c.sysCode = c.conn.Query("x-sys-code", "")
	if len(c.sysCode) == 0 {
		err = errors.SysCodeNotFound
		return
	}
	return
}

// 总结：目前绝大多数浏览器不支持直接发送WebSocket Ping帧，因此在实际开发中，应该实现应用层ping机制作为主要心跳检测方案，todo 同时保留对未来可能的原生支持的兼容检测。
func (c *Client) InitHeartbeat(timeoutSecond time.Duration) {
	ticker := time.NewTicker(timeoutSecond * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			//*2是防止丢包，连续丢包两次，再加5s网络延迟容错
			if time.Since(c.LastActive) > (timeoutSecond*2*time.Second + 5) { // 5秒容错
				log.Println("心跳超时", "clientId", c.id)
				err := c.SendMessage(websocket.CloseMessage, []byte("Heartbeat timeout"))
				if err != nil {
					log.Println("发送心跳超时消息失败", err)
				}
				c.Close()
				return
			}
		case <-c.Ctx.Done():
			return
		}
	}
}

// 在Client结构体中添加ReadMessage方法
func (c *Client) ReadMessage() (messageType int, message []byte, err error) {
	if c.conn == nil {
		return 0, nil, ErrConnClosed
	}
	return c.conn.ReadMessage()
}