ai_scheduler/internal/gateway/gateway.go

117 lines
2.3 KiB
Go

package gateway
import (
"errors"
"log"
"sync"
"time"
)
type Gateway struct {
mu sync.RWMutex
clients map[string]*Client // clientID -> Client
uidMap map[string][]string // uid -> []clientID
}
func NewGateway() *Gateway {
return &Gateway{
clients: make(map[string]*Client),
uidMap: make(map[string][]string),
}
}
func (g *Gateway) AddClient(c *Client) {
g.mu.Lock()
defer func() {
g.mu.Unlock()
//心跳开始计时
c.LastActive = time.Now()
log.Println("client connected:", c.GetID())
log.Println("客户端已连接")
}()
g.clients[c.GetID()] = c
}
func (g *Gateway) Cleanup(clientID string) {
g.mu.Lock()
defer func() {
if c, ex := g.clients[clientID]; ex {
delete(g.clients, clientID)
_ = c.conn.Close()
c.Cancel()
}
g.mu.Unlock()
log.Println("client disconnected:", clientID)
}()
for uid, list := range g.uidMap {
newList := []string{}
for _, cid := range list {
if cid != clientID {
newList = append(newList, cid)
}
}
g.uidMap[uid] = newList
}
}
func (g *Gateway) SendToAll(msg []byte) {
g.mu.RLock()
defer g.mu.RUnlock()
for _, c := range g.clients {
_ = c.SendFunc(msg)
}
}
func (g *Gateway) SendToClient(clientID string, msg []byte) error {
g.mu.RLock()
defer g.mu.RUnlock()
if c, ok := g.clients[clientID]; ok {
return c.SendFunc(msg)
}
return errors.New("client not found")
}
func (g *Gateway) BindUid(clientID, uid string) error {
g.mu.Lock()
defer g.mu.Unlock()
if _, ok := g.clients[clientID]; !ok {
return errors.New("client not found")
}
g.uidMap[uid] = append(g.uidMap[uid], clientID)
log.Printf("bind %s -> uid:%s\n", clientID, uid)
return nil
}
func (g *Gateway) SendToUid(uid string, msg []byte) {
g.mu.RLock()
defer g.mu.RUnlock()
if list, ok := g.uidMap[uid]; ok {
for _, cid := range list {
if c, ok := g.clients[cid]; ok {
_ = c.SendFunc(msg)
}
}
}
}
func (g *Gateway) ListClients() []string {
g.mu.RLock()
defer g.mu.RUnlock()
ids := make([]string, 0, len(g.clients))
for id := range g.clients {
ids = append(ids, id)
}
return ids
}
func (g *Gateway) ListUids() map[string][]string {
g.mu.RLock()
defer g.mu.RUnlock()
result := make(map[string][]string, len(g.uidMap))
for uid, list := range g.uidMap {
result[uid] = append([]string(nil), list...)
}
return result
}