Go WebSocket服务优雅关闭测试：关闭exit通道不生效问题排查

根本原因 #

你直接关闭exit通道后立刻关闭了done通道，执行时序存在竞态：

1. 关闭exit后，Run方法的select确实会触发执行s.Shutdown()，但s.Shutdown()内部的wg.Wait()等待逻辑是异步执行的（因为Run是跑在单独的goroutine里的）
2. 你在关闭exit后没有等待Shutdown执行完成，就直接关闭了done，测试主线程立刻唤醒执行断言，此时25个请求还没全部处理完成，自然计数不符合预期。

而你直接调用s.Shutdown()的版本是同步等待Shutdown（包括内部的wg.Wait()）执行完成后，才关闭done，此时所有请求都处理完毕，断言自然通过。

修复方案 #

修改测试代码的逻辑，把done的关闭时机放到Run方法返回之后，保证Shutdown逻辑完全执行完成再触发测试断言：

func TestServer_GracefulShutdown(t *testing.T) {
    done := make(chan struct{})
    exit := make(chan struct{})
    ctx := context.Background()

    finishedRequestCount := atomic.NewInt32(0)
    ln, err := net.Listen("tcp", "localhost:")
    require.Nil(t, err)

    handler := HandlerFunc(func(conn *websocket.Conn) {
        for {
            _, _, err := conn.ReadMessage()
            if err != nil {
                return
            }
            time.Sleep(100 * time.Millisecond)
            finishedRequestCount.Inc()
        }
    })
    s, err := makeServer(ctx, handler) // server create
    require.Nil(t, err)
    // AfterFunc只负责关闭exit触发关闭流程
    time.AfterFunc(50*time.Millisecond, func() {
        close(exit)
    })
    go func() {
        fmt.Printf("Running...")
        require.Nil(t, s.Run(ctx, exit, ln))
        // Run返回意味着Shutdown已经执行完成，再关闭done
        close(done)
    }()
    for i := 0; i < connCount; i++ {
        go func() {
            err := clientRun(ln)
            require.Nil(t, err)
        }()
    }

    <-done

    assert.Equal(t, int32(totalCount), finishedRequestCount.Load())
}

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