How to start and close Websocket connection in Golang

Websocket is a network protocol that enables real-time two-way communication between a client and a server. It can achieve efficient data transmission and low-latency interaction effects at the same time without using HTTP polling. Golang is an open source, high-performance programming language with excellent concurrency performance and excellent code quality. In Websocket communication, Golang also has excellent implementation methods. This article will introduce how to start and close Websocket connections in Golang.

Implementation of Websocket in Golang

To implement Websocket in Golang, you need to perform the following steps:

1. Import the corresponding package

In Go language, implementing Websocket requires the support of WebSocket and Http packages. The code is as follows:

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

Among them, github.com/gorilla/websocket is an excellent WebSocket package that supports multiple protocols. It contains some WebSocket-related structures and functions, which we can use to quickly implement a WebSocket connection.

2. Define WebSocket connection processing function

In the function that handles WebSocket connection, we need to complete the following tasks:

a. Define WebSocket upgrade

For an HTTP request, if it needs to be upgraded to a WebSocket connection, we need to process it and return the relevant upgrade request. The code is as follows:

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

func websocketHandler(w http.ResponseWriter, r *http.Request) {
    // 升级为WebSocket连接
    ws, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println(err)
        return
    }
    
    // 实现WebSocket通信
    ...
}

In this code, we use the Upgrader structure, which is the function that implements WebSocket upgrade in the gorilla/websocket library.

Upgrader includes the read and write buffer size and CheckOrigin function. CheckOrigin is used to check whether the request comes from a legitimate source. If this method returns false, the upgrade request will be rejected. Since we don't need to check validity in this example, our CheckOrigin always returns true.

After calling the Upgrader.Upgrade method, we will get a WebSocket object ws. On this object, we can call methods such as ReadMessage and WriteMessage to complete subsequent WebSocket communication.

b. Read WebSocket messages

After getting the WebSocket connection, we need to continuously read the messages sent by the client. To read WebSocket messages, you can use the ReadMessage method. The code is as follows:

func websocketHandler(w http.ResponseWriter, r *http.Request) {
    ...
    
    for {
        // 读取消息
        _, message, err := ws.ReadMessage()
        if err != nil {
            log.Println(err)
            break
        }
        
        // 处理消息
        ...
    }
}

In this code, we use a for loop to continuously read messages from WebSocket. If the read fails, the WebSocket connection has been closed and we exit the loop.

c. Send WebSocket message

After processing the message sent by the client, we may need to send some processing results to the client. To send WebSocket messages, you can use the WriteMessage method. The code is as follows:

func websocketHandler(w http.ResponseWriter, r *http.Request) {
    ...
    
    for {
        ...
        
        // 发送消息
        err = ws.WriteMessage(websocket.TextMessage, []byte(reply))
        if err != nil {
            log.Println(err)
            break
        }
    }
}

In this code, we send a message to the client through the WriteMessage method. The first parameter of this method is the type of message, which we can specify as TextMessage or BinaryMessage; the second parameter is a byte array that contains the data to be sent.

3. Start the WebSocket service

After completing the above steps, we can start a WebSocket service. The code is as follows:

func main() {
    http.HandleFunc("/ws", websocketHandler)
    log.Println("Server started at http://127.0.0.1:8080")
    http.ListenAndServe(":8080", nil)
}

In this code, we call the http.HandleFunc function to map the "/ws" path to our WebSocket processing function websocketHandler. Finally, we call http.ListenAndServe to start the service.

Close WebSocket connection

After using WebSocket, we need to manually close the connection. To close the WebSocket connection normally, you need to pay attention to the following points:

a. Server side actively closes

When the server side needs to close the WebSocket connection, it can call the WebSocket.Close method. The code is as follows:

func websocketHandler(w http.ResponseWriter, r *http.Request) {
    defer ws.Close() // 关闭WebSocket连接
    
    for {
        ...
    }
}

In this code, we use the defer statement to automatically close the WebSocket connection when the websocket function ends.

b. Client actively closes

When the client needs to close the WebSocket connection, it can send a Close message to the server. The server will receive the Close message and actively close the WebSocket connection after receiving the message.

We can add a branch to process the Close message in the server-side websocket function. The code is as follows:

func websocketHandler(w http.ResponseWriter, r *http.Request) {
    ...
    
    for {
        // 读取消息
        messageType, message, err := ws.ReadMessage()
        if err != nil {
            log.Println(err)
            break
        }
        
        // 处理消息
        switch messageType {
        case websocket.CloseMessage:
            // 收到关闭消息，关闭连接
            return
        case websocket.TextMessage:
            // 收到文本消息，处理消息
            ...
        case websocket.BinaryMessage:
            ...
        }
    }
    
    // 关闭WebSocket连接
    ws.Close()
}

In this code, we add a processing branch with the message type CloseMessage. After receiving the Close message, we immediately exit the websocket function and close the WebSocket connection.

c. Abnormal closing

When an abnormality occurs in the WebSocket connection, the connection needs to be closed in time. For example, when the network is disconnected, the connection between the client and the server times out, etc., the WebSocket connection needs to be closed.

We can catch exceptions in the websocket function and close the WebSocket connection when the exception is caught. The code is as follows:

func websocketHandler(w http.ResponseWriter, r *http.Request) {
    defer ws.Close() // 当函数退出时关掉WebSocket连接
    
    for {
        _, message, err := ws.ReadMessage()
        if err != nil {
            if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
                log.Printf("error: %v", err)
            }
            break
        }
        
        // 处理消息
        ...
    }
}

In this code, we manually close the WebSocket connection through the defer statement before the websocket function returns. In the ReadMessage operation, we added the ability to capture exception errors and perform log output.

Summary

In Golang, Websocket connection can be achieved through the gorilla/websocket library. We can use the Upgrader structure in the processing function of WebSocket communication to upgrade the HTTP protocol to WebSocket. When connecting to WebSocket, we need to implement some reading and sending operations of WebSocket messages.

In order to ensure the normal closing of the WebSocket connection, we need to pay attention to the following points:

1. The server must manually close the WebSocket connection after completing the operation;
2. When the client receives the Close message, the server should immediately close the WebSocket connection;
3. When the WebSocket connection occurs When an exception occurs, the connection needs to be closed promptly after catching the exception.

The above is the detailed content of How to start and close Websocket connection in Golang. For more information, please follow other related articles on the PHP Chinese website!