websocket gracefully shutdown - go

Question

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
}

Answer 1

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.