push-server/interval/server/ws/hub.go

package ws

import (
	"context"
	"log"

	"git.jinshen.cn/remilia/push-server/interval/protocol"
	"github.com/coder/websocket"
)

// Hub is the central message distribution hub.
type Hub struct {
	// ClientID -> *Client
	clientsByID map[string]*Client
	// topic -> clientID -> *Client
	//
	// Invariant:
	// - clientsByTopic contains only topics with at least one active subscriber.
	// - A topic key must not exist if it has zero clients.
	clientsByTopic map[protocol.Topic]map[string]*Client
	// clientID -> topic -> struct{}
	topicsByClients map[string]map[protocol.Topic]struct{}

	register    chan *Client
	unregister  chan *Client
	subscribe   chan protocol.Subscription
	unsubscribe chan protocol.Subscription
	broadcast   chan protocol.BroadcastMessage
}

func NewHub() *Hub {
	return &Hub{
		clientsByID:     make(map[string]*Client),
		clientsByTopic:  make(map[protocol.Topic]map[string]*Client),
		topicsByClients: make(map[string]map[protocol.Topic]struct{}),

		register:    make(chan *Client),
		unregister:  make(chan *Client),
		subscribe:   make(chan protocol.Subscription, 8),
		unsubscribe: make(chan protocol.Subscription, 8),
		broadcast:   make(chan protocol.BroadcastMessage, 64),
	}
}

func (h *Hub) RegisterClient(client *Client) {
	h.register <- client
}

func (h *Hub) UnregisterClient(client *Client) {
	h.unregister <- client
}

func (h *Hub) Subscribe(sub protocol.Subscription) {
	h.subscribe <- sub
}

func (h *Hub) Unsubscribe(sub protocol.Subscription) {
	h.unsubscribe <- sub
}

func (h *Hub) BroadcastMessage(ctx context.Context, msg protocol.BroadcastMessage) error {
	select {
	case h.broadcast <- msg:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}

}

func (h *Hub) Run(ctx context.Context) {
	log.Println("Hub is running")
	for {
		select {
		case c := <-h.register:
			h.onRegister(c)

		case c := <-h.unregister:
			h.onUnregister(c)

		case c := <-h.subscribe:
			h.onSubscribe(c)

		case s := <-h.unsubscribe:
			h.onUnsubscribe(s)

		case msg := <-h.broadcast:
			h.onBroadcast(msg)

		case <-ctx.Done():
			h.shutdown()
			return
		}
	}
}

// Get a client by its ClientID
func (h *Hub) getClient(id string) (*Client, bool) {
	c, ok := h.clientsByID[id]
	return c, ok
}

// Create a new entry for the client in topicsByClients map when it registers.
func (h *Hub) onRegister(c *Client) {
	h.clientsByID[c.ID] = c
	h.topicsByClients[c.ID] = make(map[protocol.Topic]struct{})
	log.Printf("Current clientList: %v\n", h.clientsByID)
}

// Delete all topic subscriptions for the client when it unregisters.
func (h *Hub) onUnregister(c *Client) {
	topics := h.topicsByClients[c.ID]
	for t := range topics {
		if clients, ok := h.clientsByTopic[t]; ok {
			delete(clients, c.ID)
			if len(clients) == 0 {
				delete(h.clientsByTopic, t)
			}
		}
	}
	delete(h.topicsByClients, c.ID)
	delete(h.clientsByID, c.ID)
}

func (h *Hub) onSubscribe(s protocol.Subscription) {
	c, ok := h.getClient(s.ClientID)
	if !ok {
		// If the client does not exist, log an error and return.
		log.Printf("Subscribe failed: client %s not found", s.ClientID)
		return
	}
	if h.clientsByTopic[s.Topic] == nil {
		h.clientsByTopic[s.Topic] = make(map[string]*Client)
	}

	h.clientsByTopic[s.Topic][s.ClientID] = c
	h.topicsByClients[s.ClientID][s.Topic] = struct{}{}
	log.Printf("Client %s subscribed to topic %s", s.ClientID, s.Topic)
}

func (h *Hub) onUnsubscribe(s protocol.Subscription) {
	if clients, ok := h.clientsByTopic[s.Topic]; ok {
		delete(clients, s.ClientID)
		if len(clients) == 0 {
			delete(h.clientsByTopic, s.Topic)
		}
	}

	if topics, ok := h.topicsByClients[s.ClientID]; ok {
		delete(topics, s.Topic)
	}
}

func (h *Hub) onBroadcast(msg protocol.BroadcastMessage) {
	if !msg.Topic.Valid() {
		log.Printf("Broadcast failed: invalid topic")
		return
	}
	log.Printf("Receiving message for topic %s: %s", msg.Topic, string(msg.Payload))
	for _, c := range h.clientsByTopic[msg.Topic] {
		select {
		case c.SendChan <- msg.Payload:
			log.Printf("Sending message to client %s: %s", c.ID, string(msg.Payload))
		default:
			h.UnregisterClient(c)
			if c.Conn != nil {
				_ = c.Conn.Close(websocket.StatusPolicyViolation, "Slow consumer")
			}
		}
	}
}

func (h *Hub) shutdown() {
	for _, c := range h.clientsByID {
		close(c.SendChan)
		if c.Conn != nil {
			_ = c.Conn.Close(websocket.StatusNormalClosure, "Server shutdown")
		}
	}
}