connect/server/hub.go

package server

import (
	"sync"

	"sikey.com/websocket/utils/zlog"

	"github.com/redis/go-redis/v9"
)

type HubConfig struct {
	ConnectSize    int
	DisconnectSize int
	MessageSize    int
	ServerId       string // serverId 服务器ID

	Rdb redis.UniversalClient
}

type Hub struct {
	serverId string // serverId 服务器ID
	rdb      redis.UniversalClient

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

	Message chan *Message

	exchange *Exchange
}

func NewHub(cfg HubConfig) *Hub {
	hub := &Hub{
		serverId: cfg.ServerId,

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

		Connect:    make(chan *Client, cfg.ConnectSize),
		Disconnect: make(chan *Client, cfg.DisconnectSize),
		Message:    make(chan *Message, cfg.MessageSize),
	}
	hub.exchange = NewExchange(cfg.ServerId, cfg.Rdb)

	go hub.run()
	return hub
}

func (h *Hub) run() {
	for {
		select {
		case client := <-h.Connect:

			h.mutex.Lock()
			h.clients[client.UserId] = client
			h.exchange.OnPublishConnect(client)
			h.mutex.Unlock()

		case client := <-h.Disconnect:

			h.mutex.Lock()
			close(client.Send)
			delete(h.clients, client.UserId)
			h.exchange.OnPublishDisconnect(client)
			h.mutex.Unlock()

		case message := <-h.Message:

			h.mutex.RLock()
			if client, ok := h.clients[message.Receiver]; ok {
				if client.isRemotely {
					h.exchange.OnPublishMessage(client, message)
				} else {
					zlog.Info("message: ", message)
					client.Send <- message
				}
			}
			h.mutex.RUnlock()

		case conn := <-h.exchange.Connect:

			h.mutex.Lock()
			h.clients[conn.UserId] = &Client{
				isRemotely: true,
				hub:        h,
				Send:       make(chan *Message),
			}
			h.mutex.Unlock()

		case conn := <-h.exchange.Disconnect:

			h.mutex.Lock()
			if client, ok := h.clients[conn.UserId]; ok {
				close(client.Send)
				delete(h.clients, client.UserId)
			}
			h.mutex.Unlock()

		case message := <-h.exchange.Message:

			h.mutex.RLock()
			if client, ok := h.clients[message.Receiver]; ok {
				client.Send <- message
			}
			h.mutex.RUnlock()
		}
	}
}

func (h *Hub) GetClients() map[string]*Client {
	h.mutex.RLock()
	defer h.mutex.RUnlock()
	return h.clients
}