lone-services/pkg/utils/ws.go

package utils

import (
	"context"
	"sync"

	"github.com/zeromicro/go-zero/core/logx"
)

// Stream 通用双向流抽象接口
type Stream interface {
	Send(payload []byte) error
}

// WaitingSession 通用等待会话
type WaitingSession struct {
	Stream Stream
	Ctx    context.Context
	Cancel context.CancelFunc
	MsgCh  chan []byte
}

// WsSessionManager 会话管理器
type WsSessionManager struct {
	mu       sync.RWMutex
	sessions map[string]*WaitingSession // key: deviceCode
}

// 全局单例，logic直接用 WsSessionMgr
var WsSessionMgr = NewWsSessionManager()

func NewWsSessionManager() *WsSessionManager {
	return &WsSessionManager{
		sessions: make(map[string]*WaitingSession),
	}
}

// Register 注册会话
func (m *WsSessionManager) Register(deviceCode string, sess *WaitingSession) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.sessions[deviceCode] = sess
	logx.Infof("register ws session, deviceCode=%s", deviceCode)
}

// Pop 取出并删除会话
func (m *WsSessionManager) Pop(deviceCode string) *WaitingSession {
	m.mu.Lock()
	defer m.mu.Unlock()
	sess, ok := m.sessions[deviceCode]
	if !ok {
		return nil
	}
	delete(m.sessions, deviceCode)
	logx.Infof("pop ws session, deviceCode=%s", deviceCode)
	return sess
}

// Remove 移除会话
func (m *WsSessionManager) Remove(deviceCode string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.sessions, deviceCode)
	logx.Infof("remove ws session, deviceCode=%s", deviceCode)
}