nginx-ui/server/modules/agent_server/hub.go

package agent_server

import (
	"nginx-ui/server/models"
)

// Hub maintains the set of active clients and broadcasts messages to the
// clients.
type Hub struct {
	// Registered clients.
	clientMap map[string]*WsClient
	// Inbound messages from the clients.
	broadcast chan []byte

	// Register requests from the clients.
	register chan *WsClient

	// Unregister requests from clients.
	unregister chan *WsClient
	handlers   map[string]func(c *WsClient, message *models.AgentData) (interface{}, error)
}

func newHub() *Hub {
	return &Hub{
		broadcast: make(chan []byte),
		register:  make(chan *WsClient),
		clientMap: make(map[string]*WsClient),
		handlers:  make(map[string]func(c *WsClient, message *models.AgentData) (interface{}, error)),
	}
}

func (h *Hub) SetMessageHandler(key string, handler func(c *WsClient, message *models.AgentData) (interface{}, error)) {
	h.handlers[key] = handler
}

func (h *Hub) FindClient(key string) (*WsClient, bool) {
	c, ok := h.clientMap[key]
	return c, ok
}

func (h *Hub) run() {
	for {
		select {
		case client := <-h.register:
			old, ok := h.clientMap[client.Token]
			if ok {
				old.Close()
			}
			h.clientMap[client.Token] = client
		case client := <-h.unregister:
			if _, ok := h.clientMap[client.Token]; ok {
				delete(h.clientMap, client.Token)
			}
			client.Close()
			delete(h.clientMap, client.Token)
		case message := <-h.broadcast:
			for _, client := range h.clientMap {
				select {
				case client.send <- message:
				default:
					close(client.send)
					delete(h.clientMap, client.Token)
				}
			}
		}
	}
}

var AgentHub = newHub()