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
}
CodePudding user response:
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 change 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.