前往Golang学习之Websocket通信

Golang是一种非常强大的编程语言，其函数式编程和并发性使得它在Web编程中得到了越来越广泛的应用。在本文中，我们将介绍如何使用Golang实现WebSocket通信，为Web编程提供更好的支持。

WebSocket是一种协议，通过该协议，Web应用程序可以打开一个TCP连接，只需进行一次握手。它支持全双工通信，这意味着服务器和客户端可以同时发送和接收信息，从而使实时数据传输变得更加容易。

接下来，我们将介绍如何使用Golang实现WebSocket通信。我们将从新建WebSocket服务开始，然后讲解如何监听WebSocket请求、发送和接收消息以及关闭WebSocket连接。

新建WebSocket服务

要使用Golang实现WebSocket通信，我们首先需要导入" net/http "和"golang.org/x/net/websocket "包。然后，在我们的代码中，我们需要编写一个自定义的WebSocket处理器函数，这个处理函数需要实现websocket.Handler接口。接下来，我们需要使用http.Handle()函数将该处理器注册为一个HTTP处理器。

package main

import (
    "fmt"
    "net/http"

    "golang.org/x/net/websocket"
)

// WebSocket处理器
func wsHandler(ws *websocket.Conn) {
    fmt.Println("WebSocket connected")
    for {
        var recvMessage string
        // 从WebSocket中读取消息
        if err := websocket.Message.Receive(ws, &recvMessage); err != nil {
            // 连接已关闭
            fmt.Println("WebSocket disconnected")
            return
        }
        // 将消息原样返回给客户端
        if err := websocket.Message.Send(ws, recvMessage); err != nil {
            // 发送消息失败
            fmt.Println("Failed to send message")
        }
    }
}

func main() {
    // 将WebSocket处理器注册为HTTP处理器
    http.Handle("/ws", websocket.Handler(wsHandler))
    // 启动HTTP服务
    if err := http.ListenAndServe(":8080", nil); err != nil {
        fmt.Println(err)
    }
}

在上面的代码中，我们定义了一个名为"wsHandler"的WebSocket处理器函数。当WebSocket连接建立时，它会输出一条"WebSocket connected"的消息。然后，它将等待WebSocket客户端发送消息，然后直接将收到的消息原样返回给客户端。当WebSocket连接关闭时，它会输出一条"WebSocket disconnected"的消息。最后，我们将"/ws"路径绑定到我们的WebSocket处理器函数，并且使用http.ListenAndServe()函数开启一个HTTP服务，监听8080端口。

监听WebSocket请求

现在，让我们编写一个客户端程序来测试我们的WebSocket服务。我们可以使用浏览器中的JavaScript代码或者使用下面的Golang代码。代码中我们定义了一个"NewWebSocketClient"函数来创建WebSocket客户端。

package main

import (
    "fmt"
    "time"

    "golang.org/x/net/websocket"
)

// 创建Websocket客户端
func NewWebSocketClient(url string) (*websocket.Conn, error) {
    var wsConfig *websocket.Config
    origin := "http://localhost/"
    // 创建WebSocket配置
    wsConfig, err := websocket.NewConfig(url, origin)
    if err != nil {
        fmt.Println("Failed to create WebSocket config:", err.Error())
        return nil, err
    }
    wsConfig.Protocol = []string{"chat", "superchat"} // 设置协议
    wsConfig.Location = nil                           // 使用默认路径
    wsConfig.Header = make(http.Header)               // 创建HTTP头
    ws, err := websocket.DialConfig(wsConfig)         // 创建WebSocket客户端
    if err != nil {
        fmt.Println("Failed to create WebSocket client:", err.Error())
        return nil, err
    }
    return ws, nil
}

func main() {
    // 创建WebSocket客户端
    ws, err := NewWebSocketClient("ws://localhost:8080/ws")
    if err != nil {
        fmt.Println("Failed to create WebSocket client:", err.Error())
        return
    }
    defer ws.Close()

    // 向WebSocket发送消息
    message := "Hello World"
    if err := websocket.Message.Send(ws, message); err != nil {
        fmt.Println("Failed to send message:", err.Error())
        return
    }

    // 等待WebSocket返回消息
    var recvMessage string
    if err := websocket.Message.Receive(ws, &recvMessage); err != nil {
        fmt.Println("Failed to receive message:", err.Error())
        return
    }
    fmt.Println("Received message:", recvMessage)

    // 关闭WebSocket连接
    ws.Close()
    time.Sleep(time.Second * 1) // 等待1秒钟
}

在上面的代码中，我们使用了"golang.org/x/net/websocket"包中的函数"websocket.DialConfig()"来创建WebSocket客户端。我们还定义了一个"NewWebSocketClient"函数来封装WebSocket客户端的创建过程，这使得我们可以轻松地从其他地方使用该函数来创建WebSocket客户端。

发送和接收消息

现在我们已经成功建立了WebSocket连接，我们可以开始发送和接收消息了。我们可以使用"go-routine"在后台不断地读取WebSocket消息。

package main

import (
    "fmt"
    "time"

    "golang.org/x/net/websocket"
)

// 创建WebSocket客户端
func NewWebSocketClient(url string) (*websocket.Conn, error) {
    var wsConfig *websocket.Config
    origin := "http://localhost/"
    // 创建WebSocket配置
    wsConfig, err := websocket.NewConfig(url, origin)
    if err != nil {
        fmt.Println("Failed to create WebSocket config:", err.Error())
        return nil, err
    }
    wsConfig.Protocol = []string{"chat", "superchat"} // 设置协议
    wsConfig.Location = nil                           // 使用默认路径
    wsConfig.Header = make(http.Header)               // 创建HTTP头
    ws, err := websocket.DialConfig(wsConfig)         // 创建WebSocket客户端
    if err != nil {
        fmt.Println("Failed to create WebSocket client:", err.Error())
        return nil, err
    }
    return ws, nil
}

// 读取消息
func readWebSocket(ws *websocket.Conn) {
    for {
        var recvMessage string
        // 从WebSocket中读取消息
        if err := websocket.Message.Receive(ws, &recvMessage); err != nil {
            fmt.Println("Failed to receive message:", err.Error())
            ws.Close()
            return
        }
        fmt.Println("Received message:", recvMessage)
    }
}

func main() {
    // 创建WebSocket客户端
    ws, err := NewWebSocketClient("ws://localhost:8080/ws")
    if err != nil {
        fmt.Println("Failed to create WebSocket client:", err.Error())
        return
    }
    defer ws.Close()

    // 使用go-routine读取WebSocket消息
    go readWebSocket(ws)

    // 循环发送消息
    for i := 0; i < 10; i++ {
        message := fmt.Sprintf("Hello World %d", i)
        if err := websocket.Message.Send(ws, message); err != nil {
            fmt.Println("Failed to send message:", err.Error())
            break
        }
        fmt.Println("Sent message:", message)
        time.Sleep(time.Second * 1) // 等待1秒钟
    }

    // 关闭WebSocket连接
    ws.Close()
    time.Sleep(time.Second * 1) // 等待1秒钟
}

上述代码中我们使用了"readWebSocket"函数来接收WebSocket消息。我们通过循环调用"websocket.Message.Receive()"来从WebSocket连接中读取消息。当出现错误时，我们会关闭WebSocket连接并退出循环。

关闭WebSocket连接

最后，我们需要确保在完成WebSocket通信后正确关闭WebSocket连接。为了实现这一点，我们需要在WebSocket客户端和WebSocket服务器端的代码中分别增加一些逻辑。

在客户端中，我们必须在发送完所有消息后关闭WebSocket连接。在执行关闭之前，我们需要确保所有未完成的go-routine已经执行完毕。

package main

import (
    "fmt"
    "sync"
    "time"

    "golang.org/x/net/websocket"
)

// 创建WebSocket客户端
func NewWebSocketClient(url string) (*websocket.Conn, error) {
    var wsConfig *websocket.Config
    origin := "http://localhost/"
    // 创建WebSocket配置
    wsConfig, err := websocket.NewConfig(url, origin)
    if err != nil {
        fmt.Println("Failed to create WebSocket config:", err.Error())
        return nil, err
    }
    wsConfig.Protocol = []string{"chat", "superchat"} // 设置协议
    wsConfig.Location = nil                           // 使用默认路径
    wsConfig.Header = make(http.Header)               // 创建HTTP头
    ws, err := websocket.DialConfig(wsConfig)         // 创建WebSocket客户端
    if err != nil {
        fmt.Println("Failed to create WebSocket client:", err.Error())
        return nil, err
    }
    return ws, nil
}

// 读取消息
func readWebSocket(ws *websocket.Conn, wg *sync.WaitGroup) {
    defer wg.Done() // 缩小go-routine的生存范围
    for {
        var recvMessage string
        // 从WebSocket中读取消息
        if err := websocket.Message.Receive(ws, &recvMessage); err != nil {
            fmt.Println("Failed to receive message:", err.Error())
            ws.Close()
            return
        }
        fmt.Println("Received message:", recvMessage)
    }
}

func main() {
    // 创建WebSocket客户端
    ws, err := NewWebSocketClient("ws://localhost:8080/ws")
    if err != nil {
        fmt.Println("Failed to create WebSocket client:", err.Error())
        return
    }
    defer ws.Close()

    var wg sync.WaitGroup
    go readWebSocket(ws, &wg)

    // 循环发送消息
    for i := 0; i < 10; i++ {
        message := fmt.Sprintf("Hello World %d", i)
        if err := websocket.Message.Send(ws, message); err != nil {
            fmt.Println("Failed to send message:", err.Error())
            break
        }
        fmt.Println("Sent message:", message)
        time.Sleep(time.Second * 1) // 等待1秒钟
    }

    // 阻塞等待所有go-routine执行完成
    wg.Wait()

    // 关闭WebSocket连接
    fmt.Println("Closing WebSocket connection ...")
    if err := ws.Close(); err != nil {
        fmt.Println("Failed to close WebSocket connection:", err.Error())
    }
}

在上述代码中，我们使用了"sync.WaitGroup"来确保所有go-routine执行完毕后再执行WebSocket连接的关闭操作。在主go-routine中，我们首先设定一个"WaitGroup"对象，然后在启动"readWebSocket"go-routine之前调用"WaitGroup.Add()"方法，把计数器加1。然后，在"readWebSocket"函数执行完成后，我们会在"go-routine"中调用"WaitGroup.Done()"方法，将计数器减1。最后，在主go-routine中调用"WaitGroup.Wait()"方法阻塞等待所有go-routine执行完成，然后再关闭WebSocket连接。

在服务器端代码中，我们可以在WebSocket处理器函数中监听WebSocket连接关闭事件并执行相应的清理操作。

package main

import (
    "fmt"
    "net/http"

    "golang.org/x/net/websocket"
)

// WebSocket处理器
func wsHandler(ws *websocket.Conn) {
    fmt.Println("WebSocket connected")
    for {
        var recvMessage string
        // 从WebSocket中读取消息
        if err := websocket.Message.Receive(ws, &recvMessage); err != nil {
            // 连接已关闭
            fmt.Println("WebSocket disconnected")
            return
        }
        // 将消息原样返回给客户端
        if err := websocket.Message.Send(ws, recvMessage); err != nil {
            // 发送消息失败
            fmt.Println("Failed to send message")
        }
    }
}

func main() {
    // 将WebSocket处理器注册为HTTP处理器
    http.Handle("/ws", websocket.Handler(wsHandler))
    // 启动HTTP服务
    if err := http.ListenAndServe(":8080", nil); err != nil {
        fmt.Println(err)
    }
}

在上述代码中，当WebSocket连接关闭时，WebSocket处理器函数"wsHandler"会输出"WebSocket disconnected"的消息。我们可以在这个处理函数中添加其他清理逻辑，比如记录日志或者释放资源。

总结

在这篇文章中，我们展示了如何使用Golang实现WebSocket通信。我们从新建WebSocket服务开始，讲解了如何监听WebSocket请求、发送和接收消息以及关闭WebSocket连接。我们还介绍了如何使用go-routine来实现非阻塞式的WebSocket通信，并且在代码中添加了一些逻辑，以确保在WebSocket通信完成后正确关闭WebSocket连接。

WebSocket是实现实时数据传输的一种非常强大的协议。在Web开发中，使用Golang和WebSocket通信可以为开发人员提供更大的灵活性和更好的维护性。我们希望这篇文章能够帮助你更深入地了解WebSocket通信和Golang编程。