connect/server/hub.go

package server

import (
	"sync"

	"go.uber.org/zap"
	"sikey.com/websocket/config"
)

type Hub struct {
	serverId string // serverId 服务器ID

	clients         map[string]*Client
	mutex           sync.RWMutex
	Connect         chan *Client
	Disconnect      chan *Client
	Message         chan *Message
	FirebaseMessage chan *FirebaseMessage

	FCM  *FirebaseMessageServer
	Nats *Nats
}

func NewHub(serverId string) *Hub {
	hub := &Hub{
		serverId: serverId,

		clients: make(map[string]*Client),
		mutex:   sync.RWMutex{},

		Connect:    make(chan *Client, config.Websocket.ConnectSize),
		Disconnect: make(chan *Client, config.Websocket.ConnectSize),
		Message:    make(chan *Message, config.Websocket.MessageSize),

		Nats: NewNats(config.Websocket.NatsUrl),
		FCM:  NewFirebaseMessageServer(),
	}

	go hub.run()
	return hub
}

func (h *Hub) run() {
	for {
		select {
		case client := <-h.Connect:
			h.mutex.Lock()
			h.clients[client.UserId] = client
			h.mutex.Unlock()
			if h.Nats != nil {
				h.Nats.Subscribe <- &subscriber{client: client}
			}
		case client := <-h.Disconnect:
			h.mutex.Lock()
			close(client.Send)
			delete(h.clients, client.UserId)
			h.mutex.Unlock()
			if h.Nats != nil {
				h.Nats.Unsubscribe <- &subscriber{client: client}
			}
		case message := <-h.Message:
			h.mutex.RLock()
			if client, ok := h.clients[message.Receiver]; ok {
				zap.L().Info("[conn] message to client",
					zap.String("receiver", message.Receiver),
					zap.String("request_id", message.RequestId))
				client.Send <- message
			} else {
				h.Nats.Send <- &natsMessage{userId: message.Receiver, message: message}
			}

			h.mutex.RUnlock()
			//  检查用户如果不在线,并且有 FCM token,则发送 FCM 消息
			// case firebaseMessage := <-h.FirebaseMessage:
			// 	if err := h.FCM.Send(context.Background(), firebaseMessage.token, firebaseMessage.message); err != nil {
			// 		zap.L().Error("unable to send fcm message", zap.Error(err))
			// 	}
		}

	}
}

func (h *Hub) GetClients() []string {
	h.mutex.RLock()
	defer h.mutex.RUnlock()

	var clients = make([]string, 0)
	for _, c := range h.clients {
		clients = append(clients, c.UserId)
	}
	return clients
}

func (h *Hub) getClientByUserId(userId string) *Client {
	h.mutex.RLock()
	defer h.mutex.RUnlock()

	if client, ok := h.clients[userId]; ok {
		return client
	}
	return nil
}