GoLang的websocket编程

GoLang的websocket编程

后端

package main

import (
	"fmt"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"reflect"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true }, // 解决跨域
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil) // HTTP 升级为 WebSocket
	if err != nil {
		log.Println("升级失败:", err)
		return
	}
	defer conn.Close()
	for {
		msgType, msg, err := conn.ReadMessage() // 读取消息
		if err != nil {
			log.Println("读取错误:", err)
			break
		}
		fmt.Println(msg)
		fmt.Println(reflect.TypeOf(msg))
		fmt.Println(reflect.TypeOf(msgType))
		log.Printf("收到: %s", msg)
		// 原样返回消息（实现 Echo 服务）
		if err := conn.WriteMessage(msgType, msg); err != nil {
			log.Println("发送错误:", err)
			break
		}
	}
}
func main() {
	http.HandleFunc("/ws", handleWebSocket)
	log.Fatal(http.ListenAndServe(":8080", nil))
}

1、核心代码实现

// 将字符串转换为字节切片，明确指定为文本消息类型
if err := conn.WriteMessage(websocket.TextMessage, []byte(msg)); err != nil {
    log.Printf("消息推送失败: %v", err)
    // 可选：返回错误信息给调用方
    return fmt.Errorf("推送失败: %w", err)
}

// websocket.TextMessage：适用于 UTF-8 编码的字符串（如 JSON/XML）
// websocket.BinaryMessage：适用于二进制数据（如图片/音视频）
// 字符串消息必须使用 TextMessage 类型

2、错误处理优化：

if err := conn.WriteMessage(websocket.TextMessage, []byte(msg)); err != nil {
    // 识别特定错误类型
    if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
        return "连接已正常关闭"
    }
    // 记录原始错误日志
    log.Printf("WS_ERR: %v | 消息: %s", err, msg) 
    return "服务繁忙，请重试"
}

// 区分连接关闭和发送失败场景
// 生产环境建议用结构化日志（如 JSON 格式）

消息推送：

1、设置字符串返回类型

// 假设msg是字符串类型
message := "这是要返回的字符串消息"

// 将字符串转换为字节切片
msgBytes := []byte(message)

// 发送文本消息
if err := conn.WriteMessage(websocket.TextMessage, msgBytes); err != nil {
    // 错误处理
    log.Println("发送消息失败:", err)
    // 可以选择断开连接或进行其他处理
    return
}

2、设置JSON返回类型

type Response struct {
    Code int    `json:"code"`
    Data string `json:"data"`
}

func sendJSON(conn *websocket.Conn, msg string) error {
    resp := Response{Code: 200, Data: msg}
    jsonData, _ := json.Marshal(resp)  // 序列化为 JSON
    return conn.WriteMessage(websocket.TextMessage, jsonData)
}

3、 广播消息实现

var clients = make(map[*websocket.Conn]bool) // 全局连接池

func broadcast(msg string) {
    for client := range clients {
        go func(c *websocket.Conn) { // 并发发送避免阻塞
            if err := c.WriteMessage(websocket.TextMessage, []byte(msg)); err != nil {
                delete(clients, c) // 移除失效连接
            }
        }(client)
    }
}

⚠️ 关键注意事项

1. 并发安全

   • 使用 sync.Mutex 保护共享资源（如 clients 连接池）
   • 通道传递消息避免竞态条件：

   type Message struct {
       Conn *websocket.Conn
       Text string
   }
   msgChan := make(chan Message, 100) // 缓冲队列

2. 性能优化

   • 大消息（> 64KB）启用分片传输：

   conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) // 设置超时
   writer, _ := conn.NextWriter(websocket.TextMessage)
   io.Copy(writer, bytes.NewReader(largeData)) // 流式写入
   writer.Close()

   • 启用压缩减少带宽：

   upgrader := websocket.Upgrader{
       EnableCompression: true, // 开启压缩
   }

3. 连接生命周期

   • 心跳机制维持长连接：

   go func() {
       for range time.Tick(30 * time.Second) {
           if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
               return // 心跳失败终止连接
           }
       }
   }()