websocket 优雅地关闭

Question

我有一个 websocket 服务器。我为他写了一个测试，测试他正常关机的能力。创建了 5 个连接，每个连接发送 5 个请求。一段时间后，关机开始。必须满足所有 25 个请求。如果我关闭 exit 频道，那么测试将无法正常工作。

    time.AfterFunc(50*time.Millisecond, func() {
        close(exit)
        close(done)
    })

如果我只调用 s.shutdown 函数，那么一切正常。

    time.AfterFunc(50*time.Millisecond, func() {
        require.Nil(t, s.Shutdown())
        close(done)
    })

我的测试

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)
    time.AfterFunc(50*time.Millisecond, func() {
        close(exit)
        close(done)
    })
    go func() {
        fmt.Printf("Running...")
        require.Nil(t, s.Run(ctx, exit, ln))
    }()
    for i := 0; i < connCount; i++ {
        go func() {
            err := clientRun(ln)
            require.Nil(t, err)
        }()
    }

    <-done

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

我的运行函数

func (s *Server) Run(ctx context.Context, exit <-chan struct{}, ln net.Listener) error {
    errs := make(chan error, 1)

    go func() {
        err := s.httpServer.Run(ctx, exit, ln)
        if err != nil {
            errs <- err
        }
    }()

    select {
    case <-ctx.Done():
        return s.Close()
    case <-exit:
        return s.Shutdown()
    case err := <-errs:
        return err
    }
}

我的关机

func (s *Server) Shutdown() error {
    err := s.httpServer.Shutdown() // we close the possibility to connect to any conn
    s.wg.Wait()
    return err
}

Answer 1

如果执行以下代码会发生什么？

close(exit)
close(done)

两个频道几乎同时关闭。第一个触发等待正常关闭的 Shutdown 函数。但是第二个触发了

的评估

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

在正常关机仍在运行或尚未开始时触发。

如果直接执行 Shutdown 函数，它将阻塞直到完成，然后 close(done) 才会开始断言。这就是为什么它有效：

require.Nil(t, s.Shutdown())
close(done)

您可以将 close(done) 移动到以下位置以使测试工作，同时使用 exit 通道关闭：

go func() {
    fmt.Printf("Running...")
    require.Nil(t, s.Run(ctx, exit, ln))
    close(done)
}()

这样done会在Shutdown函数执行后关闭

正如评论中所讨论的，我强烈建议使用上下文而不是通道来关闭。他们隐藏了封闭渠道的复杂性。

websocket 优雅地关闭

websocket gracefully shutdown

go

websocket