transport/bufWriter: fast-fail on error returned from flushKeepBuffer()

internal/transport/http_util.go

@@ -329,20 +329,24 @@ func (w *bufWriter) Write(b []byte) (n int, err error) {
 		b := w.pool.Get().(*[]byte)
 		w.buf = *b
 	}
+	var bytesWritten int
 	for len(b) > 0 {
 		nn := copy(w.buf[w.offset:], b)
 		b = b[nn:]
 		w.offset += nn
-		n += nn
 		if w.offset >= w.batchSize {
-			err = w.flushKeepBuffer()
+			bytesWritten, err = w.flushKeepBuffer()
+			n += bytesWritten
+			if err != nil {
+				return n, err
+			}
 		}
 	}
-	return n, err
+	return n, nil
 }
 
 func (w *bufWriter) Flush() error {
-	err := w.flushKeepBuffer()
+	_, err := w.flushKeepBuffer()
 	// Only release the buffer if we are in a "shared" mode
 	if w.buf != nil && w.pool != nil {
 		b := w.buf
@@ -352,17 +356,18 @@ func (w *bufWriter) Flush() error {
 	return err
 }
 
-func (w *bufWriter) flushKeepBuffer() error {
+func (w *bufWriter) flushKeepBuffer() (int, error) {
 	if w.err != nil {
-		return w.err
+		return 0, w.err
 	}
 	if w.offset == 0 {
-		return nil
+		return 0, nil
 	}
-	_, w.err = w.conn.Write(w.buf[:w.offset])
+	var n int
+	n, w.err = w.conn.Write(w.buf[:w.offset])
 	w.err = toIOError(w.err)
 	w.offset = 0
-	return w.err
+	return n, w.err
 }
 
 type ioError struct {

internal/transport/http_util_test.go

@@ -19,7 +19,10 @@
 package transport
 
 import (
+	"errors"
 	"fmt"
+	"io"
+	"net"
 	"reflect"
 	"testing"
 	"time"
@@ -215,6 +218,39 @@ func (s) TestParseDialTarget(t *testing.T) {
 	}
 }
 
+type badNetworkConn struct {
+	net.Conn
+}
+
+func (c *badNetworkConn) Write([]byte) (int, error) {
+	return 0, io.EOF
+}
+
+// This test ensures Write() on a broken network connection does not lead to
+// an infinite loop. See https://github.com/grpc/grpc-go/issues/7389 for more details.
+func (s) TestWriteBadConnection(t *testing.T) {
+	data := []byte("test_data")
+	// Configure the bufWriter with a batchsize that results in data being flushed
+	// to the underlying conn, midway through Write().
+	writeBufferSize := (len(data) - 1) / 2
+	writer := newBufWriter(&badNetworkConn{}, writeBufferSize, getWriteBufferPool(writeBufferSize))
+
+	errCh := make(chan error, 1)
+	go func() {
+		_, err := writer.Write(data)
+		errCh <- err
+	}()
+
+	select {
+	case <-time.After(time.Second):
+		t.Fatalf("Write() did not return in time")
+	case err := <-errCh:
+		if !errors.Is(err, io.EOF) {
+			t.Fatalf("Write() = %v, want error presence = %v", err, io.EOF)
+		}
+	}
+}
+
 func BenchmarkDecodeGrpcMessage(b *testing.B) {
 	input := "Hello, %E4%B8%96%E7%95%8C"
 	want := "Hello, 世界"