abci: socket server shutdown response handler (#7547)

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

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