diff --git a/abci/server/socket_server.go b/abci/server/socket_server.go index f98b9d153..9496faf47 100644 --- a/abci/server/socket_server.go +++ b/abci/server/socket_server.go @@ -259,14 +259,26 @@ func (s *SocketServer) handleResponses( responses <-chan *types.Response, ) { bw := bufio.NewWriter(conn) - for res := range responses { - if err := types.WriteMessage(res, bw); err != nil { - closeConn <- fmt.Errorf("error writing message: %w", err) - return - } - if err := bw.Flush(); err != nil { - closeConn <- fmt.Errorf("error flushing write buffer: %w", err) + for { + select { + case <-ctx.Done(): return + case res := <-responses: + if err := types.WriteMessage(res, bw); err != nil { + select { + case <-ctx.Done(): + case closeConn <- fmt.Errorf("error writing message: %w", err): + } + return + } + if err := bw.Flush(); err != nil { + select { + case <-ctx.Done(): + case closeConn <- fmt.Errorf("error flushing write buffer: %w", err): + } + + return + } } } } diff --git a/abci/tests/client_server_test.go b/abci/tests/client_server_test.go index 2581714e8..a97c0c7c4 100644 --- a/abci/tests/client_server_test.go +++ b/abci/tests/client_server_test.go @@ -4,6 +4,7 @@ import ( "context" "testing" + "github.com/fortytw2/leaktest" "github.com/stretchr/testify/assert" abciclientent "github.com/tendermint/tendermint/abci/client" @@ -13,6 +14,8 @@ import ( ) func TestClientServerNoAddrPrefix(t *testing.T) { + t.Cleanup(leaktest.Check(t)) + ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -27,9 +30,11 @@ func TestClientServerNoAddrPrefix(t *testing.T) { assert.NoError(t, err, "expected no error on NewServer") err = server.Start(ctx) assert.NoError(t, err, "expected no error on server.Start") + t.Cleanup(server.Wait) client, err := abciclientent.NewClient(logger, addr, transport, true) assert.NoError(t, err, "expected no error on NewClient") err = client.Start(ctx) assert.NoError(t, err, "expected no error on client.Start") + t.Cleanup(client.Wait) }