Golang websocket服务端/客户端示例

0.安装依赖

go get github.com/gorilla/websocket

1.服务端

package utils

import (
    "log"
    "net/http"

    "github.com/gorilla/websocket"
)

// 从http升级到websocket,这是websocket的设计标准
var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}

// 处理单个websocket链接
func handleConn(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println("升级Websocket失败", err.Error())
    }
    for {
        // 读取消息
        _, msg, err := conn.ReadMessage()
        if err != nil {
            log.Println("读取失败:", err)
            break
        }
        log.Printf("收到消息: %s", msg)

        // 回复响应
        if err := conn.WriteMessage(websocket.TextMessage, []byte("服务端收到消息: "+string(msg))); err != nil {
            log.Println("写入失败:", err)
            break
        }
    }
}

func StartWsServer() {
    http.HandleFunc("/ws", handleConn)
    // 启动ws服务,http.ListenAnaServe会自动分配handleConn到新的协程
    http.ListenAndServe(":10002", nil)
}

2.客户端

package utils

import (
    "log"
    "time"

    "github.com/gorilla/websocket"
)

// 写消息
type writeReq struct {
    messageType int
    data        []byte
}

func listenResponse(conn *websocket.Conn) {
    for {
        _, msg, err := conn.ReadMessage()
        if err != nil {
            log.Println("读取消息失败:", err)
            break
        }
        log.Println("收到服务端消息:", string(msg))
    }
}

func sendLoop(conn *websocket.Conn, writeChan chan writeReq) {
    ticker := time.NewTicker(1 * time.Second)
    defer ticker.Stop()

    for {
        select {
        case req := <-writeChan:
            if err := conn.WriteMessage(req.messageType, req.data); err != nil {
                log.Println("发送消息失败:", err)
                return
            }
        case <-ticker.C:
            if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
                log.Println("心跳发送失败:", err)
                return
            }
        }
    }
}

// 这里的websocket连接,并发读是安全的,并发写是不安全的
func StartWsClient() {
    uri := "ws://127.0.0.1:10002/ws"
    conn, _, err := websocket.DefaultDialer.Dial(uri, nil)
    if err != nil {
        log.Fatal("链接服务端失败:", err.Error())
    }
    defer conn.Close()

    writeChan := make(chan writeReq, 10)
    go listenResponse(conn)
    go sendLoop(conn, writeChan)

    // 发送测试消息
    writeChan <- writeReq{websocket.TextMessage, []byte("你好")}
    writeChan <- writeReq{websocket.TextMessage, []byte("你好")}
    writeChan <- writeReq{websocket.TextMessage, []byte("你好")}

    <-make(chan struct{}) // 阻塞主协程
}