如何在WebSocket连接关闭时终止Go程序运行？

问题描述 #

我开发了一个连接BitMart WebSocket的Go程序，希望在WebSocket保持连接时持续运行，连接关闭时自动终止程序。当前程序能正常运行，但会无限阻塞无法停止；WebSocket对象有私有属性isClose表示连接状态，但无法直接访问，且不想修改SDK源码。

当前代码：

func main() {
    //try to connect to rabbitmq
    rabbitCh, err := rabbitmq.OpenChannel()
    if err != nil {
        log.Println(err)
        os.Exit(1)
    }
    defer rabbitCh.Close()

    app := Config{
        Rabbit: rabbitCh,
    }

    // https://developer-pro.bitmart.com/en/spot/#public-depth-channel
    ws := bitmart.NewWS(bitmart.Config{WsUrl: bitmart.WS_URL})
    println(ws.isClose)
    err = ws.Connection(func(message string) {
        OnMessage(message, app)
    })
    if err != nil {
        log.Println(err)
    }

    channels := []string{
        "spot/trade:BTC_USDT",
    }
    forever := make(chan bool)
    go ws.SubscribeWithoutLogin(channels)

    <-forever
}

已尝试方案：

• 使用forever通道，但程序即使按Ctrl+C也无法终止；
• 尝试访问私有属性ws.isClose，但权限不足，且不想修改SDK。

解决方案 #

可以通过监听系统中断信号+利用WebSocket连接的生命周期回调来实现需求，无需修改SDK源码：

1. 核心思路 #

• 用通道监听程序退出信号（如Ctrl+C），收到信号时主动关闭WebSocket连接；
• 将Connection方法放入goroutine执行，该方法会在WebSocket连接关闭后返回，此时触发程序退出逻辑；
• 通过一个统一的退出通道等待所有终止条件触发。

2. 修改后的代码 #

import (
    "log"
    "os"
    "os/signal"
    "syscall"

    // 导入你的bitmart和rabbitmq包
    "your-path/bitmart"
    "your-path/rabbitmq"
)

func main() {
    // 连接RabbitMQ
    rabbitCh, err := rabbitmq.OpenChannel()
    if err != nil {
        log.Fatal(err)
    }
    defer rabbitCh.Close()

    app := Config{
        Rabbit: rabbitCh,
    }

    // 初始化WebSocket客户端
    ws := bitmart.NewWS(bitmart.Config{WsUrl: bitmart.WS_URL})

    // 创建退出通道
    quit := make(chan struct{})

    // 监听系统中断信号（Ctrl+C、kill命令）
    sigChan := make(chan os.Signal, 1)
    signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
    go func() {
        <-sigChan
        log.Println("收到退出信号，正在关闭WebSocket连接...")
        _ = ws.Close() // 调用SDK的Close方法主动关闭连接
        close(quit)
    }()

    // 在goroutine中执行WebSocket连接逻辑
    go func() {
        err := ws.Connection(func(message string) {
            OnMessage(message, app)
        })
        if err != nil {
            log.Printf("WebSocket连接关闭或出错: %v", err)
        }
        // 连接关闭后，触发退出
        close(quit)
    }()

    // 订阅指定通道
    channels := []string{"spot/trade:BTC_USDT"}
    if err := ws.SubscribeWithoutLogin(channels); err != nil {
        log.Fatalf("订阅失败: %v", err)
    }

    // 等待退出信号
    <-quit
    log.Println("程序已终止")
}

3. 关键说明 #

• 信号监听：通过signal.Notify监听SIGINT（Ctrl+C）和SIGTERM（kill命令）信号，收到信号后调用ws.Close()主动关闭WebSocket连接，确保资源正确释放；
• Connection方法的异步执行：将ws.Connection()放入goroutine，因为该方法会在WebSocket连接正常运行时持续阻塞，连接关闭后才会返回。返回后关闭quit通道，触发程序退出；
• 统一退出通道：quit通道作为所有终止条件的统一入口，无论是连接自然关闭还是收到中断信号，都会通过关闭该通道让主goroutine退出，从而终止程序。

额外说明 #

查看BitMart Go SDK的Connection方法实现可知，当WebSocket连接关闭时，该方法会退出内部循环并返回，因此我们可以利用这一特性，在方法返回后触发程序退出逻辑，无需依赖私有属性isClose。

内容的提问来源于stack exchange，提问作者Isslerman