websocket gracefully shutdown

I have a websocket server. I wrote a test for him that tests his ability to shutdown gracefully. 5 connections are created and each sends 5 requests. After a while, shutdown starts. All 25 requests must be fulfilled. If I close the exit channel, then the test does not work as it should.

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

And if I just call the s.shutdown function, then everything is ok.

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

My test

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())
}

My run func

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
	}
}

My shutdown

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

What is happening if the following code is executed?

close(exit)
close(done)

Both channels are closed almost at the same time. The first triggers the Shutdown function which waits for a graceful shutdown. But the second triggers the evaluation of

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

It is triggered while the graceful shutdown is still running or hasn't even started yet.

If you execute the Shutdown function directly it will block until finished and only then close(done) will start the assertion. That is why this works:

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

You can move the close(done) to the following location to make the test work while using the exit channel to close:

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

This way done will be closed after the Shutdown function was executed.

As discussed in the comments I strongly suggest to use contexts instead of channels to close. They have the complexity of close channels hidden away.