如何检测Go WebSocket已关闭并退出对应的处理函数？

Go里没有和.NET CancellationToken完全一致的机制，但可以通过**通道(channel)**结合错误检测，或者使用context.Context来实现类似的取消逻辑，解决WebSocket关闭时退出所有无限循环的问题。以下是优化后的代码：

import (
    "context"
    "log"

    "github.com/docker/docker/api/types"
    "github.com/docker/docker/client"
    "github.com/gorilla/websocket"
)

func handleDockerShell(c *websocket.Conn, containerId string) {
    cli, err := client.NewClientWithOpts(client.FromEnv)
    if err != nil {
        log.Println(err)
        return
    }
    defer cli.Close()

    execConfig := types.ExecConfig{
        Tty:         true,
        Detach:      false,
        AttachStdin: true,
        AttachStderr: true,
        AttachStdout: true,
        Cmd:         []string{"bash"},
    }

    idResponse, err := cli.ContainerExecCreate(context.Background(), containerId, execConfig)
    if err != nil {
        log.Println(err)
        return
    }

    hijackedResponse, err := cli.ContainerExecAttach(context.Background(), idResponse.ID, types.ExecStartCheck{Tty: true})
    if err != nil {
        log.Println(err)
        return
    }
    defer hijackedResponse.Conn.Close()

    // 创建退出信号通道，关闭时通知所有goroutine终止
    quit := make(chan struct{})
    // 启动WebSocket读协程，转发输入到Docker
    go func() {
        defer close(quit)
        for {
            _, inputBytes, err := c.ReadMessage()
            if err != nil {
                // WebSocket关闭或读取失败，触发退出
                log.Println("WebSocket读取错误:", err)
                return
            }
            _, err = hijackedResponse.Conn.Write(inputBytes)
            if err != nil {
                log.Println("Docker连接写入错误:", err)
                return
            }
        }
    }()

    // 启动Docker读协程，转发输出到WebSocket
    go func() {
        defer close(quit)
        outputBuf := make([]byte, 1024)
        for {
            select {
            case <-quit:
                // 收到退出信号，立即终止循环
                return
            default:
                n, err := hijackedResponse.Conn.Read(outputBuf)
                if err != nil {
                    log.Println("Docker连接读取错误:", err)
                    return
                }
                err = c.WriteMessage(websocket.BinaryMessage, outputBuf[:n])
                if err != nil {
                    log.Println("WebSocket写入错误:", err)
                    return
                }
            }
        }
    }()

    // 主协程等待退出信号，避免函数提前返回
    <-quit
    // 主动关闭WebSocket连接
    _ = c.Close()
}

核心改进说明： #

• 统一退出信号：用quit chan struct{}作为所有协程的终止触发器，任何一个协程检测到错误（WebSocket关闭、Docker连接断开）时，会关闭该通道，其他协程通过select监听通道实现快速退出。
• 错误不忽略：不再丢弃ReadMessage、Write等操作的错误，一旦出现异常立即触发退出逻辑。
• 资源自动清理：通过defer语句确保Docker客户端、连接，以及WebSocket连接在函数退出时被正确关闭，避免资源泄漏。
• 替换无效主循环：原来的无限Sleep循环被替换为监听退出通道，函数在收到终止信号后立即返回，不再无意义阻塞。

如果需要更贴近CancellationToken的层级传递能力，也可以使用context.WithCancel：

ctx, cancel := context.WithCancel(context.Background())
defer cancel()

// 在协程中通过select监听ctx.Done()
go func() {
    for {
        select {
        case <-ctx.Done():
            return
        default:
            // 执行读写操作
        }
    }
}()

// 触发取消时调用cancel()

这种方式本质和通道机制一致，ctx.Done()就是内置的取消信号通道，适合需要跨层级传递取消信号的场景。

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