How to set up Go WebSocket server?

How to build a Go WebSocket server: Install the gorilla/websocket library. Create an HTTP server to handle WebSocket handshake requests. Upgrade HTTP requests to WebSocket connections. Handle WebSocket messages on the server side, including sending and receiving data. The server can be extended to allow clients to subscribe to specific channels and receive messages only from those channels.

Build a Go WebSocket server

WebSocket is a full-duplex communication protocol that allows the client and server to perform two-way data transmission on a single TCP connection. The Go language provides powerful support for building WebSocket servers. This article will introduce how to use Go to create a WebSocket server, with practical examples.

Getting Started

First, you need a Go development environment, which includes the Go compiler and a code editor. You also need to install the gorilla/websocket library, which provides a WebSocket implementation in Go.

go get -u github.com/gorilla/websocket

Create Server

A WebSocket server is essentially a traditional HTTP server, but it handles WebSocket handshake requests. Here's how to create a WebSocket server in Go:

package main

import (
    "fmt"
    "log"
    "net/http"

    "github.com/gorilla/websocket"
)

func main() {
    // 升级 HTTP 请求到 WebSocket 连接
    upgrader := websocket.Upgrader{
        ReadBufferSize:  1024,
        WriteBufferSize: 1024,
    }

    http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
        conn, err := upgrader.Upgrade(w, r, nil)
        if err != nil {
            log.Fatal(err)
        }

        // websocket.Conn 是一个双向的连接
        for {
            // 从客户端接收消息
            messageType, p, err := conn.ReadMessage()
            if err != nil {
                log.Fatal(err)
            }

            // 将消息返回客户端
            if err := conn.WriteMessage(messageType, p); err != nil {
                log.Fatal(err)
            }
        }
    })

    http.ListenAndServe(":8080", nil)
}

Practical Example

Now, let's extend this server to allow clients to subscribe to specific channels and receive messages only from those channels.

package main

import (
    "fmt"
    "log"
    "net/http"
    "sync"

    "github.com/gorilla/websocket"
)

type channel struct {
    sync.Mutex
    clients map[*websocket.Conn]bool
}

var channels = struct {
    sync.RWMutex
    m map[string]*channel
}{m: make(map[string]*channel)}

func NewChannel(name string) *channel {
    channels.Lock()
    defer channels.Unlock()

    if _, ok := channels.m[name]; !ok {
        channels.m[name] = &channel{clients: make(map[*websocket.Conn]bool)}
    }
    return channels.m[name]
}

func main() {
    upgrader := websocket.Upgrader{
        ReadBufferSize:  1024,
        WriteBufferSize: 1024,
    }

    http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
        conn, err := upgrader.Upgrade(w, r, nil)
        if err != nil {
            log.Fatal(err)
        }

        // 订阅和取消订阅频道
        go func() {
            for {
                // 从客户端接收消息
                messageType, p, err := conn.ReadMessage()
                if err != nil {
                    log.Fatal(err)
                }

                // 如果消息是一种订阅请求，则将连接添加到相应频道
                if messageType == websocket.TextMessage && string(p)[:4] == "join" {
                    channelName := string(p)[5:]
                    channel := NewChannel(channelName)

                    // 锁定频道的客户端列表
                    channel.Lock()
                    channel.clients[conn] = true
                    // 解锁频道的客户端列表
                    channel.Unlock()
                }

                // 如果消息是一种取消订阅请求，则将连接从相应频道中移除
                if messageType == websocket.TextMessage && string(p)[:7] == "leave" {
                    channelName := string(p)[8:]
                    channel := NewChannel(channelName)

                    channel.Lock()
                    delete(channel.clients, conn)
                    channel.Unlock()
                }
            }
        }()

        // 广播消息到客户端
        go func() {
            for {
                // 分别广播来自每个频道的消息
                channels.RLock()
                for _, channel := range channels.m {
                    channel.RLock()
                    for client := range channel.clients {
                        if err := client.WriteMessage(websocket.TextMessage, "hello"); err != nil {
                            // 如果写入失败，则从频道中移除连接
                            channel.Lock()
                            delete(channel.clients, client)
                            channel.Unlock()
                        }
                    }
                    channel.RUnlock()
                }
                channels.RUnlock()
            }
        }()
    })

    http.ListenAndServe(":8080", nil)
}

Now you can start the server and use a WebSocket client to connect to the /ws endpoint, subscribe to the channel, and receive messages from the channel.

The above is the detailed content of How to set up Go WebSocket server?. For more information, please follow other related articles on the PHP Chinese website!