Golang - 扩展 websocket 客户端以实现与不同服务器的多个连接

在使用 Gorilla Websocket 客户端时，需要扩展客户端代码以建立与多个服务器的连接。最有效的方法是将每个连接的代码移动到一个函数中，并使用 `sync.WaitGroup` 等待所有连接完成。在 `main()` 函数中，使用 `sync.WaitGroup` 跟踪启动的 goroutine，并启动一个新的 goroutine 来连接每个服务器。主循环将监听信号并根据需要调用取消函数并在 `WaitGroup` 上等待。这种方法允许根据需要随时取消连接，并提供对底层上下文的访问，以便在必要时使用超时或取消功能。

我有一个 websocket 客户端。实际上，它比下面所示的基本代码复杂得多。 我现在需要扩展此客户端代码以打开与多个服务器的连接。最终，从服务器收到消息时需要执行的任务是相同的。 处理这个问题的最佳方法是什么？ 正如我上面所说，接收消息时执行的实际代码比示例中显示的要复杂得多。

package main

import (
        "flag"
        "log"
        "net/url"
        "os"
        "os/signal"
        "time"

        "github.com/gorilla/websocket"
)

var addr = flag.String("addr", "localhost:1234", "http service address")

func main() {
        flag.Parse()
        log.SetFlags(0)

        interrupt := make(chan os.Signal, 1)
        signal.Notify(interrupt, os.Interrupt)

        // u := url.URL{Scheme: "ws", Host: *addr, Path: "/echo"}
        u := url.URL{Scheme: "ws", Host: *addr, Path: "/"}
        log.Printf("connecting to %s", u.String())

        c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
        if err != nil {
                log.Fatal("dial:", err)
        }
        defer c.Close()

        done := make(chan struct{})

        go func() {
                defer close(done)
                for {
                        _, message, err := c.ReadMessage()
                        if err != nil {
                                log.Println("read:", err)
                                return
                        }
                        log.Printf("recv: %s", message)
                }
        }()

        ticker := time.NewTicker(time.Second)
        defer ticker.Stop()

        for {
                select {
                case <-done:
                        return
                case t := <-ticker.C:
                        err := c.WriteMessage(websocket.TextMessage, []byte(t.String()))
                        if err != nil {
                                log.Println("write:", err)
                                return
                        }
                case <-interrupt:
                        log.Println("interrupt")

                        // Cleanly close the connection by sending a close message and then
                        // waiting (with timeout) for the server to close the connection.
                        err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
                        if err != nil {
                                log.Println("write close:", err)
                                return
                        }
                        select {
                        case <-done:
                        case <-time.After(time.Second):
                        }
                        return
                }
        }
}

解决方案

修改中断处理以在中断时关闭通道。这允许多个 goroutine 通过等待通道关闭来等待事件。

shutdown := make(chan struct{})
interrupt := make(chan os.signal, 1)
signal.notify(interrupt, os.interrupt)
go func() {
    <-interrupt
    log.println("interrupt")
    close(shutdown)
}()

将每个连接的代码移至函数中。此代码是问题的复制和粘贴，有两处更改：将中断通道替换为关闭通道；当函数完成时，该函数会通知sync.waitgroup。

func connect(u string, shutdown chan struct{}, wg *sync.waitgroup) {
    defer wg.done()

    log.printf("connecting to %s", u)
    c, _, err := websocket.defaultdialer.dial(u, nil)
    if err != nil {
        log.fatal("dial:", err)
    }
    defer c.close()

    done := make(chan struct{})

    go func() {
        defer close(done)
        for {
            _, message, err := c.readmessage()
            if err != nil {
                log.println("read:", err)
                return
            }
            log.printf("recv: %s", message)
        }
    }()

    ticker := time.newticker(time.second)
    defer ticker.stop()

    for {
        select {
        case <-done:
            return
        case t := <-ticker.c:
            err := c.writemessage(websocket.textmessage, []byte(t.string()))
            if err != nil {
                log.println("write:", err)
                return
            }
        case <-shutdown:
            // cleanly close the connection by sending a close message and then
            // waiting (with timeout) for the server to close the connection.
            err := c.writemessage(websocket.closemessage, websocket.formatclosemessage(websocket.closenormalclosure, ""))
            if err != nil {
                log.println("write close:", err)
                return
            }
            select {
            case <-done:
            case <-time.after(time.second):
            }
            return
        }
    }
}

在 main() 中声明 sync.WaitGroup。对于要连接的每个 websocket 端点，增加 waitgroup 并启动​​一个 goroutine 来连接该端点。启动 goroutine 后，在 waitgroup 上等待 goroutine 完成。

var wg sync.WaitGroup
for _, u := range endpoints { // endpoints is []string 
                              // where elements are URLs 
                              // of endpoints to connect to.
    wg.Add(1)
    go connect(u, shutdown, &wg)
}
wg.Wait()

上面的代码经过编辑后可以在 gorilla 的 echo 示例服务器上运行，代码为 posted on the playground。

与每个不同服务器的通信是否完全独立于其他服务器？如果是的话，我会以这样的方式四处走动：

• 在main中创建一个带有取消函数的context
• 在 main 中创建一个 waitgroup 来跟踪启动的 goroutine
• 对于每个服务器，添加到等待组，从传递上下文和等待组引用的主函数中启动一个新的 goroutine
• main 进入 for/select 循环，监听信号，如果信号到达，则调用 cancelfunc 并在 waitgroup 上等待。
• main 还可以监听来自 goroutine 的结果 chan，并且可能打印结果本身（goroutines 不应该直接这样做）。
• 正如我们所说，每个goroutine都有对工作组、上下文的引用，并且可能还有一个返回结果的chan。现在，该方法会根据 goroutine 是否必须只做一件事情，或者是否需要做一系列事情进行划分。对于第一种方法
• 如果只要做一件事，我们遵循 here 所描述的方法（观察到，为了异步，他会依次启动一个新的 goroutine 来执行 dosomething() 步骤，该步骤将在通道上返回结果） 这使得它能够随时接受取消信号。由您决定您想要的非阻塞程度以及响应取消信号的提示程度。此外，将关联的上下文传递给 goroutine 的好处是您可以调用启用的 context大多数库函数的版本。例如，如果您希望拨号的超时时间为 1 分钟，您可以创建一个新的上下文，其中包含传递的超时时间，然后使用该上下文创建 dialcontext。这允许拨号因超时或父级（您在 main 中创建的）上下文的 cancelfunc 被调用而停止。
• 如果需要做更多事情，我通常更喜欢用 goroutine 做一件事，让它调用一个新的 goroutine 来执行下一步（将所有引用传递到管道中）并退出。

这种方法可以很好地进行取消扩展，并且能够在任何步骤停止管道，并且可以轻松地为可能需要很长时间的步骤提供交易线支持上下文。

本篇关于《Golang - 扩展 websocket 客户端以实现与不同服务器的多个连接》的介绍就到此结束啦，但是学无止境，想要了解学习更多关于Golang的相关知识，请关注golang学习网公众号！