connect/server/hub.go

package server

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"

	"github.com/redis/go-redis/v9"
	"sikey.com/websocket/utils/zlog"
)

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
}

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

		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),
	}

	go hub.run()
	go hub.remotelyEvent() // 远程事件
	return hub
}

func (h *Hub) run() {
	for {
		ctx := context.Background()
		select {
		case client := <-h.Connect:
			h.clients[client.UserId] = client
			h.OnPublishConnect(ctx, client)
		case client := <-h.Disconnect:
			close(client.Send)
			delete(h.clients, client.UserId)
			h.OnPublishDisconnect(ctx, client)
		case message := <-h.Message:
			if client, ok := h.clients[message.receiver]; ok {
				if client.isRemotely {
					h.OnPublishMessage(ctx, message)
				} else {
					client.Send <- message
				}
			}
		}
	}
}

var (
	connectChannelEvent    = "client.event.connect"
	disconnectChannelEvent = "client.event.disconnect"
	messageChannelEvent    = "client.event.message"
)

func (h *Hub) remotelyEvent() {
	ctx := context.TODO()
	pubsub := h.rdb.PSubscribe(ctx, "client.event.*")
	defer pubsub.Close()

	for {
		rMsg, err := pubsub.ReceiveMessage(ctx)
		if err != nil {
			zlog.Error(err)
		}

		switch rMsg.Channel {
		case connectChannelEvent:
			var event = deserializeEvent([]byte(rMsg.Payload))
			if event.ServerId != h.serverId {
				h.Connect <- &Client{
					hub:        h,
					UserId:     event.UserId,
					isRemotely: true,
					Send:       make(chan *Message),
				}
			}
		case disconnectChannelEvent:
			var event = deserializeEvent([]byte(rMsg.Payload))
			if event.ServerId != h.serverId {
				h.Disconnect <- h.clients[event.UserId]
			}
		case messageChannelEvent:
			h.Message <- deserializeMessage([]byte(rMsg.Payload))
		}
		fmt.Println(rMsg)
	}
}

func deserializeEvent(bytes []byte) *PublishEvent {
	var event PublishEvent
	_ = json.Unmarshal(bytes, &event)
	return &event
}

type PublishEvent struct {
	ServerId string
	UserId   string
}

type PublishMessage struct {
	ServerId string
	message  *Message
}

func (h *Hub) OnPublishConnect(ctx context.Context, client *Client) error {
	if !client.isRemotely {
		event := &PublishEvent{UserId: client.UserId, ServerId: h.serverId}
		return h.rdb.Publish(ctx, connectChannelEvent, event).Err()
	}
	return nil
}

func (h *Hub) OnPublishDisconnect(ctx context.Context, client *Client) error {
	if !client.isRemotely {
		event := &PublishEvent{UserId: client.UserId, ServerId: h.serverId}
		return h.rdb.Publish(ctx, disconnectChannelEvent, event).Err()
	}
	return nil
}

func (h *Hub) OnPublishMessage(ctx context.Context, message *Message) error {
	err := h.rdb.Publish(ctx, messageChannelEvent, message).Err()
	if err != nil {
		zlog.Error(err)
		return err
	}
	return nil
}