Golang

关注公众号 jb51net

关闭
首页 > 脚本专栏 > Golang > Go语言WebSocket实时通信

Go语言WebSocket实时通信示例代码

作者：钟哩哩

在当今快速发展的互联网时代,实时通信技术已成为现代应用开发的核心需求,这篇文章主要介绍了Go语言WebSocket实时通信的相关资料,文中通过代码介绍的非常详细,需要的朋友可以参考下

一、WebSocket基础

WebSocket是一种在单个TCP连接上进行全双工通信的协议，允许服务器主动向客户端推送数据。

WebSocket特点

特性说明
全双工客户端和服务器可以同时发送数据
低延迟持久连接，避免HTTP握手开销
实时性服务器可以主动推送消息
跨域支持支持跨域通信
轻量级协议开销小

WebSocket握手过程

客户端                          服务器
  |                              |
  |--- HTTP请求(Upgrade) -------->|
  |  Connection: Upgrade         |
  |  Upgrade: websocket          |
  |  Sec-WebSocket-Key: xxx      |
  |                              |
  |<-- HTTP响应(101 Switching) --|
  |  Connection: Upgrade         |
  |  Upgrade: websocket          |
  |  Sec-WebSocket-Accept: xxx   |
  |                              |
  |======== WebSocket连接 =======|
  |<=======双向数据传输=========>|

二、Go语言WebSocket服务端

使用gorilla/websocket

go get github.com/gorilla/websocket

package main

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

var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool {
        return true // 允许所有来源
    },
}

func handler(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        fmt.Printf("Failed to upgrade connection: %v\n", err)
        return
    }
    defer conn.Close()
    
    for {
        // 读取消息
        messageType, p, err := conn.ReadMessage()
        if err != nil {
            fmt.Printf("Read error: %v\n", err)
            break
        }
        
        fmt.Printf("Received: %s\n", p)
        
        // 发送响应
        err = conn.WriteMessage(messageType, []byte("Server received: "+string(p)))
        if err != nil {
            fmt.Printf("Write error: %v\n", err)
            break
        }
    }
}

func main() {
    http.HandleFunc("/ws", handler)
    fmt.Println("Server starting on :8080")
    http.ListenAndServe(":8080", nil)
}

广播消息

package main

import (
    "net/http"
    "sync"
    
    "github.com/gorilla/websocket"
)

var (
    upgrader = websocket.Upgrader{
        ReadBufferSize:  1024,
        WriteBufferSize: 1024,
        CheckOrigin: func(r *http.Request) bool { return true },
    }
    
    clients   = make(map[*websocket.Conn]bool)
    clientsMu sync.Mutex
)

func broadcast(message []byte) {
    clientsMu.Lock()
    defer clientsMu.Unlock()
    
    for client := range clients {
        err := client.WriteMessage(websocket.TextMessage, message)
        if err != nil {
            fmt.Printf("Broadcast error: %v\n", err)
            client.Close()
            delete(clients, client)
        }
    }
}

func handler(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        fmt.Printf("Upgrade error: %v\n", err)
        return
    }
    defer conn.Close()
    
    // 注册客户端
    clientsMu.Lock()
    clients[conn] = true
    clientsMu.Unlock()
    
    // 清理客户端
    defer func() {
        clientsMu.Lock()
        delete(clients, conn)
        clientsMu.Unlock()
    }()
    
    for {
        _, message, err := conn.ReadMessage()
        if err != nil {
            fmt.Printf("Read error: %v\n", err)
            break
        }
        
        // 广播消息
        broadcast(message)
    }
}

三、WebSocket客户端

基础客户端

package main

import (
    "fmt"
    "net/url"
    "os"
    "time"
    
    "github.com/gorilla/websocket"
)

func main() {
    u := url.URL{Scheme: "ws", Host: "localhost:8080", Path: "/ws"}
    
    conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
    if err != nil {
        fmt.Printf("Dial error: %v\n", err)
        os.Exit(1)
    }
    defer conn.Close()
    
    // 发送消息
    err = conn.WriteMessage(websocket.TextMessage, []byte("Hello, Server!"))
    if err != nil {
        fmt.Printf("Write error: %v\n", err)
        return
    }
    
    // 接收响应
    _, message, err := conn.ReadMessage()
    if err != nil {
        fmt.Printf("Read error: %v\n", err)
        return
    }
    
    fmt.Printf("Received: %s\n", message)
}

并发客户端

package main

import (
    "fmt"
    "net/url"
    "sync"
    "time"
    
    "github.com/gorilla/websocket"
)

func clientWorker(id int, wg *sync.WaitGroup) {
    defer wg.Done()
    
    u := url.URL{Scheme: "ws", Host: "localhost:8080", Path: "/ws"}
    conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
    if err != nil {
        fmt.Printf("Client %d: Dial error: %v\n", id, err)
        return
    }
    defer conn.Close()
    
    // 发送消息
    message := fmt.Sprintf("Hello from client %d", id)
    err = conn.WriteMessage(websocket.TextMessage, []byte(message))
    if err != nil {
        fmt.Printf("Client %d: Write error: %v\n", id, err)
        return
    }
    
    // 接收响应
    _, response, err := conn.ReadMessage()
    if err != nil {
        fmt.Printf("Client %d: Read error: %v\n", id, err)
        return
    }
    
    fmt.Printf("Client %d: Received: %s\n", id, response)
}

func main() {
    var wg sync.WaitGroup
    
    for i := 0; i < 10; i++ {
        wg.Add(1)
        go clientWorker(i, &wg)
        time.Sleep(100 * time.Millisecond)
    }
    
    wg.Wait()
}

四、心跳机制

服务端心跳检测

func handler(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        fmt.Printf("Upgrade error: %v\n", err)
        return
    }
    defer conn.Close()
    
    conn.SetReadDeadline(time.Now().Add(10 * time.Second))
    
    conn.SetPongHandler(func(string) error {
        conn.SetReadDeadline(time.Now().Add(10 * time.Second))
        return nil
    })
    
    // 启动心跳定时器
    ticker := time.NewTicker(5 * time.Second)
    defer ticker.Stop()
    
    go func() {
        for range ticker.C {
            err := conn.WriteMessage(websocket.PingMessage, nil)
            if err != nil {
                fmt.Printf("Ping error: %v\n", err)
                return
            }
        }
    }()
    
    // 消息处理循环
    for {
        _, message, err := conn.ReadMessage()
        if err != nil {
            fmt.Printf("Read error: %v\n", err)
            break
        }
        fmt.Printf("Received: %s\n", message)
    }
}

客户端心跳处理

func main() {
    conn, _, err := websocket.DefaultDialer.Dial("ws://localhost:8080/ws", nil)
    if err != nil {
        panic(err)
    }
    defer conn.Close()
    
    // 处理pong
    conn.SetPongHandler(func(appData string) error {
        fmt.Println("Received pong")
        return nil
    })
    
    // 发送ping
    go func() {
        ticker := time.NewTicker(5 * time.Second)
        defer ticker.Stop()
        
        for range ticker.C {
            err := conn.WriteMessage(websocket.PingMessage, nil)
            if err != nil {
                fmt.Printf("Ping error: %v\n", err)
                return
            }
        }
    }()
    
    // 接收消息
    for {
        _, message, err := conn.ReadMessage()
        if err != nil {
            fmt.Printf("Read error: %v\n", err)
            break
        }
        fmt.Printf("Received: %s\n", message)
    }
}

五、消息压缩

package main

import (
    "github.com/gorilla/websocket"
    "github.com/pierrec/lz4/v4"
)

func compressMessage(message []byte) ([]byte, error) {
    compressed := make([]byte, lz4.CompressBlockBound(len(message)))
    n, err := lz4.CompressBlock(message, compressed)
    if err != nil {
        return nil, err
    }
    return compressed[:n], nil
}

func decompressMessage(compressed []byte) ([]byte, error) {
    decompressed := make([]byte, 1024*1024) // 1MB缓冲区
    n, err := lz4.UncompressBlock(compressed, decompressed)
    if err != nil {
        return nil, err
    }
    return decompressed[:n], nil
}

func handler(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        return
    }
    defer conn.Close()
    
    for {
        _, message, err := conn.ReadMessage()
        if err != nil {
            break
        }
        
        // 压缩响应
        compressed, err := compressMessage([]byte("Response: " + string(message)))
        if err != nil {
            continue
        }
        
        err = conn.WriteMessage(websocket.BinaryMessage, compressed)
        if err != nil {
            break
        }
    }
}

六、安全WebSocket (WSS)

package main

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

var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}

func handler(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        fmt.Printf("Upgrade error: %v\n", err)
        return
    }
    defer conn.Close()
    
    // 消息处理
    for {
        _, message, err := conn.ReadMessage()
        if err != nil {
            break
        }
        fmt.Printf("Received: %s\n", message)
    }
}

func main() {
    http.HandleFunc("/ws", handler)
    
    // 使用TLS
    err := http.ListenAndServeTLS(":8443", "cert.pem", "key.pem", nil)
    if err != nil {
        fmt.Printf("Server error: %v\n", err)
    }
}

七、WebSocket与HTTP共存

package main

import (
    "fmt"
    "net/http"
    "text/template"
    
    "github.com/gorilla/websocket"
)

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

func wsHandler(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        fmt.Printf("Upgrade error: %v\n", err)
        return
    }
    defer conn.Close()
    
    for {
        _, message, err := conn.ReadMessage()
        if err != nil {
            break
        }
        conn.WriteMessage(websocket.TextMessage, message)
    }
}

func homeHandler(w http.ResponseWriter, r *http.Request) {
    tmpl := template.Must(template.New("home").Parse(`
        <!DOCTYPE html>
        <html>
        <body>
            <script>
                var ws = new WebSocket("ws://localhost:8080/ws");
                ws.onopen = function() {
                    ws.send("Hello, Server!");
                };
                ws.onmessage = function(event) {
                    console.log("Received:", event.data);
                };
            </script>
        </body>
        </html>
    `))
    
    tmpl.Execute(w, nil)
}

func main() {
    http.HandleFunc("/", homeHandler)
    http.HandleFunc("/ws", wsHandler)
    
    fmt.Println("Server starting on :8080")
    http.ListenAndServe(":8080", nil)
}

八、性能优化

连接池

package main

import (
    "sync"
    "time"
    
    "github.com/gorilla/websocket"
)

type ConnectionPool struct {
    connections []*websocket.Conn
    mu          sync.Mutex
    maxSize     int
}

func NewConnectionPool(maxSize int) *ConnectionPool {
    return &ConnectionPool{
        connections: make([]*websocket.Conn, 0, maxSize),
        maxSize:     maxSize,
    }
}

func (p *ConnectionPool) Get() (*websocket.Conn, error) {
    p.mu.Lock()
    defer p.mu.Unlock()
    
    if len(p.connections) > 0 {
        conn := p.connections[len(p.connections)-1]
        p.connections = p.connections[:len(p.connections)-1]
        return conn, nil
    }
    
    // 创建新连接
    conn, _, err := websocket.DefaultDialer.Dial("ws://localhost:8080/ws", nil)
    if err != nil {
        return nil, err
    }
    
    return conn, nil
}

func (p *ConnectionPool) Put(conn *websocket.Conn) {
    p.mu.Lock()
    defer p.mu.Unlock()
    
    if len(p.connections) < p.maxSize {
        conn.SetReadDeadline(time.Now().Add(30 * time.Second))
        p.connections = append(p.connections, conn)
    } else {
        conn.Close()
    }
}

消息批量发送

type MessageQueue struct {
    messages []interface{}
    mu       sync.Mutex
    cond     *sync.Cond
    maxBatch int
}

func NewMessageQueue(maxBatch int) *MessageQueue {
    mq := &MessageQueue{
        messages: make([]interface{}, 0, maxBatch),
        maxBatch: maxBatch,
    }
    mq.cond = sync.NewCond(&mq.mu)
    return mq
}

func (mq *MessageQueue) Enqueue(msg interface{}) {
    mq.mu.Lock()
    mq.messages = append(mq.messages, msg)
    
    if len(mq.messages) >= mq.maxBatch {
        mq.cond.Signal()
    }
    mq.mu.Unlock()
}

func (mq *MessageQueue) DequeueBatch() []interface{} {
    mq.mu.Lock()
    defer mq.mu.Unlock()
    
    for len(mq.messages) < mq.maxBatch {
        mq.cond.Wait()
    }
    
    batch := mq.messages
    mq.messages = make([]interface{}, 0, mq.maxBatch)
    return batch
}

九、总结

WebSocket提供了实时双向通信能力，在Go语言中可以使用gorilla/websocket库轻松实现。关键要点包括：

  1. 基础使用：升级HTTP连接为WebSocket连接
  2. 广播机制：管理多个客户端连接，实现消息广播
  3. 心跳检测：定期发送ping/pong消息保持连接
  4. 消息压缩：使用LZ4等算法压缩消息
  5. 安全通信：配置TLS实现WSS
  6. 性能优化：连接池、批量发送等优化策略

WebSocket广泛应用于实时聊天、实时通知、在线游戏等场景，是构建实时Web应用的重要技术。

到此这篇关于Go语言WebSocket实时通信的文章就介绍到这了,更多相关Go语言WebSocket实时通信内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！

您可能感兴趣的文章:
阅读全文