1、消息提示

app/websocket/client.go

@@ -0,0 +1,84 @@
+package websocket
+
+import (
+	"context"
+	"log"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+type Client struct {
+	ID         string
+	Conn       *websocket.Conn
+	Send       chan []byte
+	Context    context.Context
+	CancelFunc context.CancelFunc
+}
+
+func (c *Client) Read(hub *Hub) {
+	defer func() {
+		hub.Unregister <- c
+		c.Conn.Close()
+		log.Printf("Client %s disconnected\n", c.ID)
+		c.CancelFunc()
+	}()
+
+	c.Conn.SetReadLimit(512)
+	c.Conn.SetReadDeadline(time.Now().Add(60 * time.Second))
+	c.Conn.SetPongHandler(func(string) error {
+		c.Conn.SetReadDeadline(time.Now().Add(60 * time.Second))
+		return nil
+	})
+
+	for {
+		select {
+		case <-c.Context.Done():
+			return
+		default:
+			_, message, err := c.Conn.ReadMessage()
+			if err != nil {
+				log.Printf("Read error from client %s: %v", c.ID, err)
+				return
+			}
+			log.Printf("Receive [%s]: %s", c.ID, message)
+			// 这里你可以把消息发给 hub.Broadcast 或业务处理
+		}
+	}
+}
+
+func (c *Client) Write() {
+	ticker := time.NewTicker(54 * time.Second) // 小于读超时，保证ping及时发
+	defer func() {
+		ticker.Stop()
+		c.Conn.Close()
+		c.CancelFunc()
+	}()
+
+	for {
+		select {
+		case <-c.Context.Done():
+			return
+		case msg, ok := <-c.Send:
+			if !ok {
+				// 通道关闭，结束写入
+				c.Conn.WriteMessage(websocket.CloseMessage, []byte{})
+				return
+			}
+
+			c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
+			err := c.Conn.WriteMessage(websocket.TextMessage, msg)
+			if err != nil {
+				log.Printf("Write error to client %s: %v", c.ID, err)
+				return
+			}
+		case <-ticker.C:
+			// 发送 ping
+			c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
+			if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+				log.Printf("Ping error to client %s: %v", c.ID, err)
+				return
+			}
+		}
+	}
+}