1、消息提示

app/websocket/hub.go

@@ -0,0 +1,63 @@
+package websocket
+
+import "sync"
+
+type Hub struct {
+	Clients    map[string]*Client
+	Register   chan *Client
+	Unregister chan *Client
+	mu         sync.RWMutex
+}
+
+func NewHub() *Hub {
+	return &Hub{
+		Clients:    make(map[string]*Client),
+		Register:   make(chan *Client),
+		Unregister: make(chan *Client),
+	}
+}
+
+func (h *Hub) Run() {
+	for {
+		select {
+		case client := <-h.Register:
+			h.mu.Lock()
+			h.Clients[client.ID] = client
+			h.mu.Unlock()
+		case client := <-h.Unregister:
+			h.mu.Lock()
+			if _, ok := h.Clients[client.ID]; ok {
+				delete(h.Clients, client.ID)
+				close(client.Send)
+			}
+			h.mu.Unlock()
+		}
+	}
+}
+
+func (h *Hub) SendToClient(id string, msg []byte) {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	if c, ok := h.Clients[id]; ok {
+		c.Send <- msg
+	}
+}
+
+func (h *Hub) SendToAll(ids []string, msg []byte) {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	for _, id := range ids {
+		if c, ok := h.Clients[id]; ok {
+			c.Send <- msg
+		}
+	}
+}
+
+func (h *Hub) Broadcast(msg []byte) {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	for _, client := range h.Clients {
+		client.Send <- msg
+	}
+}